segmentation.ipynb

Kernel: Python 3

Copyright 2019 The TensorFlow Authors.

Licensed under the Apache License, Version 2.0 (the "License");

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Segmentação de imagens

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

O foco deste tutorial é a tarefa de segmentação de imagens usando uma U-Net modificada.

O que é segmentação de imagens?

Em uma tarefa de classificação de imagens, a rede atribui um rótulo (ou classe) a cada imagem de entrada. Entretanto, vamos supor que você queira saber o formato desse objeto, qual pixel pertence a qual objeto, etc. Nesse caso, você precisa atribuir uma classe a cada pixel da imagem. Essa tarefa é conhecida como segmentação. Um modelo de segmentação retorna informações mais detalhadas sobre a imagem. A segmentação de imagens tem diversas aplicações nos exames de imagens médicos, em carros autônomos e em imagens de satélites, apenas para citar algumas.

Este tutorial usa o dataset Oxford-IIIT Pets (Parkhi et al, 2012). O dataset é composto por imagens de 37 raças de animais domésticos, com 200 imagens por raça (cerca de 100 para o treinamento e 100 para o teste). Cada imagem inclui os rótulos correspondentes e as máscaras de pixels. As máscaras são rótulos (classes) para cada pixel, e é atribuída uma categoria a cada pixel:

• Classe 1: pixels pertencentes ao animal doméstico.

• Classe 2: pixels no contorno do animal doméstico.

• Classe 3: nenhuma das opções acima/pixels nos arredores.

!pip install git+https://github.com/tensorflow/examples.git

import tensorflow as tf

import tensorflow_datasets as tfds

from tensorflow_examples.models.pix2pix import pix2pix

from IPython.display import clear_output
import matplotlib.pyplot as plt

Baixar o dataset Oxford-IIIT Pets

O dataset está disponível no TensorFlow Datasets. As máscaras de segmentação estão incluídas na versão 3 e posteriores.

dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)

Além disso, os valores de cor das imagens são normalizados no intervalo [0, 1]. Por fim, conforme mencionado acima, os pixels na máscara de segmentação são rotulados como {1, 2, 3}. Por questões de conveniência, subtraia 1 da máscara de segmentação, e os rótulos resultantes serão: {0, 1, 2}.

def normalize(input_image, input_mask):
  input_image = tf.cast(input_image, tf.float32) / 255.0
  input_mask -= 1
  return input_image, input_mask

def load_image(datapoint):
  input_image = tf.image.resize(datapoint['image'], (128, 128))
  input_mask = tf.image.resize(
    datapoint['segmentation_mask'],
    (128, 128),
    method = tf.image.ResizeMethod.NEAREST_NEIGHBOR,
  )

  input_image, input_mask = normalize(input_image, input_mask)

  return input_image, input_mask

O dataset já contém as divisões de treinamento e teste necessárias, então continue usando as mesmas divisões:

TRAIN_LENGTH = info.splits['train'].num_examples
BATCH_SIZE = 64
BUFFER_SIZE = 1000
STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE

train_images = dataset['train'].map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
test_images = dataset['test'].map(load_image, num_parallel_calls=tf.data.AUTOTUNE)

A classe abaixo faz uma ampliação simples invertendo uma imagem aleatoriamente. Saiba mais no tutorial Ampliação de imagens.

class Augment(tf.keras.layers.Layer):
  def __init__(self, seed=42):
    super().__init__()
    # both use the same seed, so they'll make the same random changes.
    self.augment_inputs = tf.keras.layers.RandomFlip(mode="horizontal", seed=seed)
    self.augment_labels = tf.keras.layers.RandomFlip(mode="horizontal", seed=seed)
  
  def call(self, inputs, labels):
    inputs = self.augment_inputs(inputs)
    labels = self.augment_labels(labels)
    return inputs, labels

Crie o pipeline de entrada, aplicando a ampliação após dividir as entradas em lotes:

train_batches = (
    train_images
    .cache()
    .shuffle(BUFFER_SIZE)
    .batch(BATCH_SIZE)
    .repeat()
    .map(Augment())
    .prefetch(buffer_size=tf.data.AUTOTUNE))

test_batches = test_images.batch(BATCH_SIZE)

Veja uma imagem de exemplo e sua máscara correspondente do dataset:

def display(display_list):
  plt.figure(figsize=(15, 15))

  title = ['Input Image', 'True Mask', 'Predicted Mask']

  for i in range(len(display_list)):
    plt.subplot(1, len(display_list), i+1)
    plt.title(title[i])
    plt.imshow(tf.keras.utils.array_to_img(display_list[i]))
    plt.axis('off')
  plt.show()

for images, masks in train_batches.take(2):
  sample_image, sample_mask = images[0], masks[0]
  display([sample_image, sample_mask])

Definir o modelo

O modelo em uso aqui é uma U-Net modificada. Uma U-Net consiste de um encoder (downsampler) e decoder (upsampler). Para aprender recursos robustos e reduzir o número de parâmetros treináveis, use um modelo pré-treinado, MobileNetV2, como encoder. Para o decoder, você usará o bloco de upsample, que já está implementado no exemplo Pix2Pix do repositório de exemplos do TensorFlow (confira o tutorial Pix2Pix: Conversão imagem-para-imagem com uma GAN condicional em um notebook).

Conforme mencionado, o encoder é um modelo MobileNetV2 pré-treinado. Você usará o modelo de tf.keras.applications. O encoder consiste de saídas específicas das camadas intermediárias do modelo. É importante salientar que o encoder não será treinado durante o processo de treinamento.

base_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False)

# Use the activations of these layers
layer_names = [
    'block_1_expand_relu',   # 64x64
    'block_3_expand_relu',   # 32x32
    'block_6_expand_relu',   # 16x16
    'block_13_expand_relu',  # 8x8
    'block_16_project',      # 4x4
]
base_model_outputs = [base_model.get_layer(name).output for name in layer_names]

# Create the feature extraction model
down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)

down_stack.trainable = False

O decoder/upsampler é simplesmente uma série de blocos de upsample implementados nos exemplos do TensorFlow:

up_stack = [
    pix2pix.upsample(512, 3),  # 4x4 -> 8x8
    pix2pix.upsample(256, 3),  # 8x8 -> 16x16
    pix2pix.upsample(128, 3),  # 16x16 -> 32x32
    pix2pix.upsample(64, 3),   # 32x32 -> 64x64
]

def unet_model(output_channels:int):
  inputs = tf.keras.layers.Input(shape=[128, 128, 3])

  # Downsampling through the model
  skips = down_stack(inputs)
  x = skips[-1]
  skips = reversed(skips[:-1])

  # Upsampling and establishing the skip connections
  for up, skip in zip(up_stack, skips):
    x = up(x)
    concat = tf.keras.layers.Concatenate()
    x = concat([x, skip])

  # This is the last layer of the model
  last = tf.keras.layers.Conv2DTranspose(
      filters=output_channels, kernel_size=3, strides=2,
      padding='same')  #64x64 -> 128x128

  x = last(x)

  return tf.keras.Model(inputs=inputs, outputs=x)

O número de filtros na última camada é definido como o número de output_channels (canais de saída). Será um canal de saída por classe.

Treinar o modelo

Agora, só falta compilar e treinar o modelo.

Como este é um problema de classificação multiclasse, use a função de perda do tf.keras.losses.SparseCategoricalCrossentropy com o argumento from_logits definido como True, já que os rótulos são inteiros escalares em vez de vetores de pontuações para cada pixel de cada classe.

Ao executar a instância, o rótulo atribuído a cada pixel é o canal com o valor mais alto. É isso que a função create_mask está fazendo.

OUTPUT_CLASSES = 3

model = unet_model(output_channels=OUTPUT_CLASSES)
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

Plote a arquitetura do modelo resultante:

tf.keras.utils.plot_model(model, show_shapes=True)

Teste o modelo para verificar o que ele prevê antes do treinamento:

def create_mask(pred_mask):
  pred_mask = tf.math.argmax(pred_mask, axis=-1)
  pred_mask = pred_mask[..., tf.newaxis]
  return pred_mask[0]

def show_predictions(dataset=None, num=1):
  if dataset:
    for image, mask in dataset.take(num):
      pred_mask = model.predict(image)
      display([image[0], mask[0], create_mask(pred_mask)])
  else:
    display([sample_image, sample_mask,
             create_mask(model.predict(sample_image[tf.newaxis, ...]))])

show_predictions()

O callback definido abaixo é usado para observar como o modelo melhora à medida que é treinado:

class DisplayCallback(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs=None):
    clear_output(wait=True)
    show_predictions()
    print ('\nSample Prediction after epoch {}\n'.format(epoch+1))

EPOCHS = 20
VAL_SUBSPLITS = 5
VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS

model_history = model.fit(train_batches, epochs=EPOCHS,
                          steps_per_epoch=STEPS_PER_EPOCH,
                          validation_steps=VALIDATION_STEPS,
                          validation_data=test_batches,
                          callbacks=[DisplayCallback()])

loss = model_history.history['loss']
val_loss = model_history.history['val_loss']

plt.figure()
plt.plot(model_history.epoch, loss, 'r', label='Training loss')
plt.plot(model_history.epoch, val_loss, 'bo', label='Validation loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss Value')
plt.ylim([0, 1])
plt.legend()
plt.show()

Fazer previsões

Agora, faça previsões. Para poupar tempo, o número de épocas foi mantido baixo, mas você pode usar um valor mais alto para conseguir resultados mais precisos.

show_predictions(test_batches, 3)

Opcional: classes desequilibradas e pesos de classes

Os datasets de segmentação semântica podem ficar muito desequilibrados, ou seja, pixels de uma classe específica podem estar mais presentes dentro das imagens do que de outras classes. Como os problemas de segmentação podem ser tratados como problemas de classificação por pixel, você pode lidar com o problema de desequilíbrio ponderando a função de perda para levar isso em conta. É uma forma simples e elegante de lidar com o problema. Confira mais detalhes no tutorial Classificação ao usar dados desequilibrados.

Para evitar ambiguidade, Model.fit não tem suporte ao argumento class_weight para alvos com mais de três dimensões.

try:
  model_history = model.fit(train_batches, epochs=EPOCHS,
                            steps_per_epoch=STEPS_PER_EPOCH,
                            class_weight = {0:2.0, 1:2.0, 2:1.0})
  assert False
except Exception as e:
  print(f"Expected {type(e).__name__}: {e}")

Portanto, nesse caso, você é que precisa implementar os pesos. Para fazer isso, utilize pesos de amostra: além de pares (data, label), Model.fit também aceita tuplas (data, label, sample_weight).

Model.fit do Keras propaga sample_weight para as perdas e métricas, que também aceitam um argumento sample_weight. O peso da amostra é multiplicado pelo valor da amostra antes do passo de redução. Por exemplo:

label = [0,0]
prediction = [[-3., 0], [-3, 0]] 
sample_weight = [1, 10] 

loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True,
                                               reduction=tf.keras.losses.Reduction.NONE)
loss(label, prediction, sample_weight).numpy()

Portanto, para fazer pesos de amostra para este tutorial, você precisa de uma função que receba um par (data, label) e retorne uma tupla (data, label, sample_weight), em que sample_weight é uma imagem de um canal contendo o peso de classes para cada pixel.

A implementação mais simples possível é usar o rótulo como índice em uma lista class_weight:

def add_sample_weights(image, label):
  # The weights for each class, with the constraint that:
  #     sum(class_weights) == 1.0
  class_weights = tf.constant([2.0, 2.0, 1.0])
  class_weights = class_weights/tf.reduce_sum(class_weights)

  # Create an image of `sample_weights` by using the label at each pixel as an 
  # index into the `class weights` .
  sample_weights = tf.gather(class_weights, indices=tf.cast(label, tf.int32))

  return image, label, sample_weights

Cada elemento do dataset resultante contém três imagens:

train_batches.map(add_sample_weights).element_spec

Agora, você pode treinar um modelo com esse dataset ponderado:

weighted_model = unet_model(OUTPUT_CLASSES)
weighted_model.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy'])

weighted_model.fit(
    train_batches.map(add_sample_weights),
    epochs=1,
    steps_per_epoch=10)

Próximos passos

Agora que você sabe o que é segmentação de imagens e como ela funciona, pode tentar usar este tutorial com diferentes saídas de camadas intermediárias ou até mesmo para diferentes modelos pré-treinados. Além disso, se quiser se desafiar, você pode experimentar o desafio de mascaramento de imagens Carvana, disponível no Kaggle.

Talvez também seja interessante conferir a API de Detecção de Objetos do Tensorflow para ver outro modelo que você pode treinar novamente usando seus próprios dados. Estão disponíveis modelos pré-treinados no TensorFlow Hub.