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.

Segmentación de imágenes

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar el bloc de notas

Este tutorial se enfoca en la tarea de segmentación de imágenes, usando una U-Net modificada.

¿Qué es la segmentación de imágenes?

En una tarea de clasificación de imágenes, la red asigna una etiqueta (o clase) a cada imagen de entrada. Pero supongamos que desea conocer la forma de ese objeto, qué pixel pertenece a qué objeto, etc. En este caso, necesita asignar una clase a cada pixel de la imagen; esta tarea se conoce como segmentación. Un modelo de segmentación devuelve información mucho más detallada sobre la imagen. La segmentación de imágenes tiene muchas aplicaciones en el diagnóstico médico por imagen, los coches autónomos y las imágenes por satélite, por nombrar sólo algunas.

Este tutorial usa el conjunto de datos Oxford-IIIT Pet (Parkhi et al, 2012). El conjunto de datos incluye imágenes de 37 razas de mascotas, con 200 imágenes por raza (~100 cada una en divisiones de entrenamiento y prueba). Cada imagen incluye las etiquetas correspondientes y máscaras a nivel de pixel. Las máscaras son etiquetas de clase para cada pixel. A cada pixel se le asigna una de tres categorías:

• Clase 1: Pixel perteneciente a la mascota.

• Clase 2: Pixel que bordea la mascota.

• Clase 3: Ninguna de las anteriores/un pixel circundante.

!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

Descargar el conjunto de datos Oxford-IIIT Pets

El conjunto de datos está disponible en Conjuntos de datos TensorFlow. Las máscaras de segmentación se incluyen a partir de la versión 3.

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

Además, los valores de color de la imagen se normalizan al rango [0, 1]. Por último, como se ha mencionado anteriormente, los pixeles de la máscara de segmentación se etiquetan como {1, 2, 3}. Para que resulte más cómodo, se resta 1 de la máscara de segmentación, lo que da como resultado las etiquetas: {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

El conjunto de datos ya contiene las divisiones de entrenamiento y prueba necesarias, así que siga usando las mismas divisiones:

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)

La siguiente clase ejecuta un simple aumento volteando aleatoriamente una imagen. Vaya al tutorial Aumento de imágenes para obtener más información.

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

Construya la canalización de entrada, aplicando el aumento tras la división en lotes de las entradas:

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)

Visualice un ejemplo de imagen y su máscara correspondiente desde el conjunto de datos:

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 el modelo

El modelo que se usa aquí es una U-Net modificada. Una U-Net consta de un codificador (downsampler) y un decodificador (upsampler). Si quiere aprender características robustas y reducir el número de parámetros entrenables, use un modelo preentrenado (MobileNetV2) como codificador. Para el decodificador, usará el bloque upsample, que ya está implementado en el ejemplo pix2pix del repositorio de ejemplos de TensorFlow. (Eche un vistazo al tutorial pix2pix: Traducción imagen a imagen con un GAN condicional en un bloc de notas).

Como se ha mencionado, el codificador es un modelo MobileNetV2 preentrenado. Usted usará el modelo de tf.keras.applications. El codificador consiste en salidas específicas de las capas intermedias del modelo. Tenga en cuenta que el codificador no se entrenará durante el proceso de entrenamiento.

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

El decodificador/upsampler es simplemente una serie de bloques de upsample implementados en ejemplos de 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)

Tenga en cuenta que el número de filtros de la última capa se establece en función del número de output_channels. Esto será un canal de salida por clase.

Entrenar el modelo

Ahora sólo queda compilar y entrenar el modelo.

Dado que se trata de un problema de clasificación multiclase, use la función de pérdida tf.keras.losses.SparseCategoricalCrossentropy con el argumento from_logits como True, ya que las etiquetas son enteros escalares en lugar de vectores de puntuaciones para cada pixel de cada clase.

Al ejecutar la inferencia, la etiqueta asignada al pixel es el canal con el valor más alto. Esto es lo que hace la función create_mask.

OUTPUT_CLASSES = 3

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

Trace la arquitectura del modelo resultante:

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

Pruebe el modelo para comprobar lo que predice antes del entrenamiento:

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()

La retrollamada definida a continuación se usa para observar cómo mejora el modelo mientras se entrena:

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()

Hacer predicciones

Ahora, haga algunas predicciones. Para ahorrar tiempo, mantuvimos un número reducido de épocas, pero puede fijarlo más alto para obtener resultados más precisos.

show_predictions(test_batches, 3)

Opcional: Clases desequilibradas y ponderaciones por clase

Los conjuntos de datos de segmentación semántica pueden estar muy desequilibrados, lo que significa que los pixeles de una clase concreta pueden estar más presentes dentro de las imágenes que los de otras clases. Dado que los problemas de segmentación pueden tratarse como problemas de clasificación por pixel, puede abordar el problema del desequilibrio asignando ponderaciones a la función de pérdida para compensarlo. Es una forma sencilla y elegante de lidiar con este problema. Consulte el tutorial Clasificación en datos desequilibrados para saber más.

Para evitar ambigüedades, Model.fit no admite el argumento class_weight para objetivos con más de 3 dimensiones.

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}")

Por tanto, en este caso deberá implementar usted mismo la ponderación. Lo hará usando ponderaciones de muestreo: Además de los pares (data, label), Model.fit también acepta los tripletes (data, label, sample_weight).

Model.fit de Keras propaga el sample_weight a las pérdidas y métricas, que también aceptan un argumento sample_weight. La ponderación de la muestra se multiplica por el valor de la muestra antes del paso de reducción. Por ejemplo:

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()

Así, para hacer ponderaciones de muestra para este tutorial, se necesita una función que tome un par (data, label) y devuelva un triplete (data, label, sample_weight) donde el sample_weight es una imagen de 1 canal que contiene la ponderación de clase para cada pixel.

La implementación más sencilla posible es usar la etiqueta como índice en una 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

Los elementos del conjunto de datos resultante contienen 3 imágenes cada uno:

train_batches.map(add_sample_weights).element_spec

Ahora ya puede entrenar un modelo en este conjunto de datos 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)

Siguientes pasos

Ahora que ya sabe qué es la segmentación de imágenes y cómo funciona, puede probar este tutorial con diferentes salidas de capas intermedias, o incluso con diferentes modelos preentrenados. También puede retarse a sí mismo probando el reto de enmascaramiento de imágenes de Carvana alojado en Kaggle.

También puede consultar la API de detección de objetos de Tensorflow para ver otro modelo que puede volver a entrenar con sus propios datos. Los modelos preentrenados están disponibles en TensorFlow Hub.