Segmentación de Imágenes Una Guía Detallada

La Segmentación de Imágenes Una Guía Detallada

¿Cómo puede una computadora distinguir entre diferentes tipos de objetos en una imagen? Guía paso a paso.

Tabla de contenidos

1. Introducción, Motivación
2. Extracción de datos
3. Visualización de las imágenes
4. Construcción de un modelo U-Net simple
5. Métricas y la función de pérdida
6. Construcción del modelo U-Net completo
7. Resumen
8. Referencias

Enlaces relevantes

• Cuaderno de trabajo en Kaggle (recomendado; regístrese en la competencia de Kaggle, haga una copia del cuaderno y utilice esa GPU gratuita)
• Cuaderno de Colab (es necesario registrarse en la competencia de Kaggle para acceder al conjunto de datos)
• Competencia de Kaggle sobre enmascaramiento de imágenes de Carvana

Introducción, Motivación

Segmentación de imágenes se refiere a la capacidad de las computadoras (o más precisamente, modelos almacenados en computadoras) para tomar una imagen y asignar a cada píxel en la imagen una categoría correspondiente. Por ejemplo, es posible ejecutar la imagen de un gato frente a una cerca blanca mostrada arriba a través de un segmentador de imágenes y obtener la imagen segmentada a continuación:

En este ejemplo, segmenté la imagen a mano. Esta es una operación tediosa y nos encantaría poder automatizarla. En esta guía te guiaré a través del proceso de entrenar un algoritmo para realizar la segmentación de imágenes. Muchas guías en internet y en libros de texto son útiles hasta cierto punto, pero todas fallan en abordar los detalles explícitos de la implementación. Aquí dejaré tan pocos aspectos sin cubrir como sea posible, para ayudarte a ahorrar tiempo al implementar la segmentación de imágenes en tus propios conjuntos de datos.

• Aplicación de métodos de aprendizaje automático para buscar defectos en rieles (Parte 2)
• Aceptando lo Inevitable La Era de las Empresas con IA en Primer Lugar
• Comparando manzanas con naranjas con python

En primer lugar, coloquemos nuestra tarea en el contexto más amplio del aprendizaje automático. La definición de aprendizaje automático es evidente por sí misma: estamos enseñando a las máquinas a aprender cómo resolver problemas que nos encantaría automatizar. Hay muchos problemas que los humanos desearían automatizar; en este artículo nos enfocamos en un subconjunto de problemas en visión por computadora. La visión por computadora busca enseñar a una computadora cómo ver. Es trivial darle a un niño de seis años una imagen de un gato frente a una cerca blanca y pedirle que segmente la imagen en píxeles de ‘gato’ y píxeles de ‘fondo’ (después de explicarle qué significa ‘segmentar’ al niño confundido, por supuesto). Y sin embargo, durante décadas las computadoras han luchado intensamente con este problema.

¿Por qué las computadoras tienen dificultades para hacer lo que un niño de seis años puede hacer? Podemos empatizar con la computadora pensando en cómo uno aprende a leer mediante el braille. Imagina que te entregan un ensayo escrito en braille y supón que no tienes conocimiento de cómo leerlo. ¿Cómo procederías? ¿Qué necesitarías para descifrar el braille al inglés?

Lo que necesitas es un método para transformar esta entrada en una salida que sea legible para ti. En matemáticas, esto lo llamamos una asignación. Decimos que nos gustaría aprender una función f(x) que asigna nuestra entrada x, que es ilegible, a una salida y, que es legible.

Con muchos meses de práctica y un buen tutor, cualquier persona puede aprender la asignación necesaria del braille al inglés. Por analogía, una computadora procesando una imagen es un poco como alguien que se encuentra con el braille por primera vez; parece como un montón de tonterías. La computadora necesita aprender la asignación necesaria f(x) para transformar un montón de números correspondientes a píxeles en algo que pueda usar para segmentar la imagen. Y desafortunadamente, el modelo de computadora no tiene miles de años de evolución, biología y años de experiencia viendo el mundo; es esencialmente ‘nacido’ cuando inicias tu programa. Esto es lo que esperamos enseñar a nuestro modelo en visión por computadora.

¿Por qué querríamos realizar la segmentación de imágenes en primer lugar? Uno de los casos de uso más obvios es Zoom. Muchas personas prefieren usar fondos virtuales durante las videoconferencias para evitar que sus compañeros de trabajo vean a su perro haciendo volteretas en la sala de estar. La segmentación de imágenes es crucial para esta tarea. Otro caso de uso poderoso es la imagen médica. Al realizar tomografías computarizadas de órganos de pacientes, podría ser útil tener un algoritmo que segmente automáticamente los órganos en las imágenes para que los profesionales médicos puedan determinar cosas como lesiones, la presencia de tumores, etc. Aquí hay un gran ejemplo de una competencia de Kaggle centrada en esta tarea.

Existen varios tipos de segmentación de imágenes, desde simples hasta complejas. En este artículo, nos ocuparemos del tipo más simple de segmentación de imágenes: la segmentación binaria. Esto significa que solo habrá dos clases diferentes de objetos, por ejemplo, ‘gato’ y ‘fondo’. Ni más ni menos.

Ten en cuenta que el código que presento aquí ha sido ligeramente reorganizado y editado para mayor claridad. Para ejecutar un código funcional, consulta los enlaces de código en la parte superior del artículo. Utilizaremos el conjunto de datos Carvana Image Masking Challenge de Kaggle. Tendrás que registrarte en este desafío para obtener acceso al conjunto de datos y conectar tu clave de API de Kaggle en el cuaderno de Colab para que funcione (si no deseas usar el cuaderno de Kaggle). Por favor, consulta esta publicación de discusión para obtener detalles sobre cómo hacerlo.

Una cosa más; por mucho que me gustaría profundizar en cada idea de este código, presumiré que tienes algún conocimiento práctico de redes neuronales convolucionales, capas de max pooling, capas conectadas densamente, capas de eliminación y conectores residuales. Desafortunadamente, discutir estos conceptos en detalle requeriría un nuevo artículo y está fuera del alcance de este, donde nos enfocamos en los aspectos prácticos de la implementación.

Extrayendo datos

Los datos relevantes para este artículo se encontrarán en las siguientes carpetas:

• train_hq.zip: Carpeta que contiene imágenes de alta calidad de entrenamiento de automóviles
• test_hq.zip: Carpeta que contiene imágenes de alta calidad de prueba de automóviles
• train_masks.zip: Carpeta que contiene máscaras para el conjunto de entrenamiento

En el contexto de la segmentación de imágenes, una máscara es la imagen segmentada. Estamos tratando de lograr que nuestro modelo aprenda a mapear una imagen de entrada a una máscara de segmentación de salida. Por lo general, se asume que la máscara verdadera (también conocida como verdad básica) es dibujada a mano por un experto humano.

El primer paso será descomprimir las carpetas desde su origen /kaggle/input:

def getZippedFilePaths():  
    zip_file_names = []
    for dirname, _, filenames in os.walk('/kaggle/input'):
        for filename in filenames:
            if filename.split('.')[-1] == 'zip':
                zip_file_names.append((os.path.join(dirname, filename)))
    return zip_file_names

zip_file_names = getZippedFilePaths()
items_to_remove = ['/kaggle/input/carvana-image-masking-challenge/train.zip',
                    '/kaggle/input/carvana-image-masking-challenge/test.zip']
zip_file_names = [item for item in zip_file_names if item not in items_to_remove]

for zip_file_path in zip_file_names:
    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
        zip_ref.extractall()

Este código obtiene las rutas de archivo para todos los archivos .zip en su entrada y los extrae en su directorio /kaggle/output. Tenga en cuenta que de manera intencionada no se extraen las fotos de baja calidad; el repositorio de Kaggle solo puede contener 20 GB de datos y este paso es necesario para evitar exceder este límite.

Visualizando las imágenes

El primer paso en la mayoría de los problemas de visión por computadora es inspeccionar su conjunto de datos. ¿Con qué exactamente estamos tratando? Primero necesitamos organizar nuestras imágenes en conjuntos de datos ordenados para su visualización. (Esta guía utilizará TensorFlow; la conversión a PyTorch no debería ser demasiado difícil).

# Agregar todos los nombres de ruta en una lista ordenada
train_hq_dir = '/kaggle/working/train_hq/'
train_masks_dir = '/kaggle/working/train_masks/'
test_hq_dir = '/kaggle/working/test_hq/'

X_train_id = sorted([os.path.join(train_hq_dir, filename) for filename in os.listdir(train_hq_dir)], key=lambda x: x.split('/')[-1].split('.')[0])
y_train = sorted([os.path.join(train_masks_dir, filename) for filename in os.listdir(train_masks_dir)], key=lambda x: x.split('/')[-1].split('.')[0])
X_test_id = sorted([os.path.join(test_hq_dir, filename) for filename in os.listdir(test_hq_dir)], key=lambda x: x.split('/')[-1].split('.')[0])

X_train_id = X_train_id[:1000]
y_train = y_train[:1000]

X_train, X_val, y_train, y_val = train_test_split(X_train_id, y_train, test_size=0.2, random_state=42)

# Crear objetos Dataset a partir de la lista de rutas de archivo
X_train = tf.data.Dataset.from_tensor_slices(X_train)
y_train = tf.data.Dataset.from_tensor_slices(y_train)
X_val = tf.data.Dataset.from_tensor_slices(X_val)
y_val = tf.data.Dataset.from_tensor_slices(y_val)
X_test = tf.data.Dataset.from_tensor_slices(X_test_id)

img_height = 96
img_width = 128
num_channels = 3
img_size = (img_height, img_width)

# Aplicar el preprocesamiento
X_train = X_train.map(preprocess_image)
y_train = y_train.map(preprocess_target)
X_val = X_val.map(preprocess_image)
y_val = y_val.map(preprocess_target)
X_test = X_test.map(preprocess_image)

# Agregar etiquetas a los objetos de dataframe (one-hot-encoded)
train_dataset = tf.data.Dataset.zip((X_train, y_train))
val_dataset = tf.data.Dataset.zip((X_val, y_val))

# Aplicar el tamaño de lote al conjunto de datos
BATCH_SIZE = 32
batched_train_dataset = train_dataset.batch(BATCH_SIZE)
batched_val_dataset = val_dataset.batch(BATCH_SIZE)
batched_test_dataset = X_test.batch(BATCH_SIZE)

# Agregar la optimización automática para el preprocesamiento
AUTOTUNE = tf.data.experimental.AUTOTUNE
batched_train_dataset = batched_train_dataset.prefetch(buffer_size=AUTOTUNE)
batched_val_dataset = batched_val_dataset.prefetch(buffer_size=AUTOTUNE)
batched_test_dataset = batched_test_dataset.prefetch(buffer_size=AUTOTUNE)

Desglosemos esto:

• Primero creamos una lista ordenada de todas las rutas de archivo a todas las imágenes en el conjunto de entrenamiento, conjunto de prueba y máscaras de verdad. Ten en cuenta que estas aún no son imágenes; hasta este punto, solo estamos buscando las rutas de archivo de las imágenes.
• Luego solo tomamos las primeras 1000 rutas de archivo a imágenes/máscaras en el conjunto de datos de Carvana. Esto se hace para reducir la carga computacional y acelerar el entrenamiento. Si tienes acceso a varias poderosas GPUs (¡qué suerte!), siéntete libre de usar todas las imágenes para obtener un rendimiento aún mejor. También creamos una división de entrenamiento/validación de 80/20. Cuantos más datos (imágenes) incluyas, más debe inclinarse esta división hacia el conjunto de entrenamiento. No es raro ver divisiones de 98/1/1 para conjuntos de entrenamiento/validación/prueba al tratar con conjuntos de datos muy grandes. Cuantos más datos haya en el conjunto de entrenamiento, mejor será tu modelo en general.
• Luego creamos objetos de conjunto de datos de TensorFlow (TF) utilizando el método tf.data.Dataset.from_tensor_slices(). Utilizar un objeto de conjunto de datos es un método común para manejar conjuntos de entrenamiento, validación y prueba, en lugar de mantenerlos como matrices Numpy. En mi experiencia, el preprocesamiento de datos es mucho más rápido y fácil cuando se utilizan objetos de conjunto de datos. Ver este enlace para la documentación.
• A continuación, especificamos la altura, el ancho y el número de canales de nuestras imágenes de entrada. Las imágenes de alta calidad reales son mucho más grandes que 96 píxeles por 128 píxeles; este “reducción de la escala” de nuestras imágenes se realiza para reducir la carga computacional (las imágenes más grandes requieren más tiempo de entrenamiento). Si tienes el poder de procesamiento necesario (GPU), no recomendaría la reducción de escala.
• Luego usamos la función .map() de nuestros objetos de conjunto de datos para preprocesar nuestras imágenes. Esto convierte las rutas de archivo en imágenes y realiza un preprocesamiento adecuado. Más sobre esto en un momento.
• Una vez que hayamos preprocesado nuestras imágenes de entrenamiento sin procesar y nuestras máscaras de referencia, necesitamos una forma de combinar las imágenes con sus máscaras. Para lograr esto, utilizamos la función .zip() de los objetos de conjunto de datos. Esto toma dos listas de datos y une el primer elemento de cada lista y los coloca en una tupla. Hace lo mismo para el segundo elemento, tercero, etc. El resultado final es una sola lista llena de tuplas de la forma (imagen, máscara).
• Luego usamos la función .batch() para crear lotes de 32 imágenes a partir de nuestras mil imágenes. El agrupamiento es una parte importante del flujo de trabajo de aprendizaje automático, ya que nos permite procesar varias imágenes a la vez, en lugar de una a la vez. Esto acelera el entrenamiento.
• Finalmente, utilizamos la función .prefetch(). Este es otro paso que ayuda a acelerar el entrenamiento. Cargar y preprocesar los datos puede ser un cuello de botella en los flujos de trabajo de entrenamiento. Esto puede llevar a tiempos de GPU o CPU inactivos, que nadie quiere. Mientras tu modelo está realizando la propagación hacia adelante y hacia atrás, la función .prefetch() puede preparar el próximo lote. La variable AUTOTUNE en TensorFlow calcula dinámicamente cuántos lotes se prefieren según los recursos de tu sistema; esto generalmente se recomienda.

Veamos más de cerca el paso de preprocesamiento:

def preprocess_image(file_path):    # Cargar y decodificar la imagen    img = tf.io.read_file(file_path)    # Puedes ajustar los canales según tus imágenes (3 para RGB)    img = tf.image.decode_jpeg(img, channels=3) # Devuelto como uint8    # Normalizar los valores de píxeles a [0, 1]    img = tf.image.convert_image_dtype(img, tf.float32)    # Cambiar el tamaño de la imagen a las dimensiones deseadas    img = tf.image.resize(img, [96, 128], method='nearest')    return imgdef preprocess_target(file_path):    # Cargar y decodificar la imagen    mask = tf.io.read_file(file_path)    # Normalizar entre 0 y 1 (solo dos clases)    mask = tf.image.decode_image(mask, expand_animations=False, dtype=tf.float32)    # Obtener solo un valor para el tercer canal    mask = tf.math.reduce_max(mask, axis=-1, keepdims=True)    # Cambiar el tamaño de la imagen a las dimensiones deseadas    mask = tf.image.resize(mask, [96, 128], method='nearest')    return mask

Estas funciones hacen lo siguiente:

• En primer lugar, convertimos las rutas de archivo en un tensor de tipo de datos ‘string’ utilizando tf.io.read_file(). Un tensor es una estructura de datos especial en TensorFlow similar a matrices multidimensionales en otras bibliotecas matemáticas, pero con propiedades y métodos especiales que son útiles para el aprendizaje profundo. Para citar la documentación de TensorFlow: tf.io.read_file() “no realiza ningún análisis, simplemente devuelve el contenido tal como está”. Básicamente, esto significa que devuelve un archivo binario (1 y 0) en el tipo de datos string que contiene la información de la imagen.
• En segundo lugar, necesitamos decodificar los datos binarios. Para hacer esto, necesitamos usar el método apropiado en TensorFlow. Dado que los datos de imagen en bruto están en formato .jpeg, usamos el método tf.image.decode_jpeg(). Dado que las máscaras están en formato GIF, podemos usar tf.io.decode_gif(), o usar el más general tf.image.decode_image(), que puede manejar cualquier tipo de archivo. Cuál elijas realmente no importa. Establecemos expand_animations=False porque estas no son realmente animaciones, son solo imágenes.
• Luego usamos convert_image_dtype() para convertir nuestros datos de imagen en float32. Esto solo se hace para las imágenes, no para la máscara, ya que la máscara ya se decodificó en float32. Hay dos tipos de datos comunes utilizados en el procesamiento de imágenes: float32 y uint8. Float32 representa un número de punto flotante (decimal) que ocupa 32 bits en la memoria de la computadora. Son firmados (lo que significa que el número puede ser negativo) y pueden variar en valor desde 0 hasta 2³² = 4294967296, aunque por convención en el procesamiento de imágenes normalizamos estos valores para que estén entre 0 y 1, donde 1 es el máximo de un color. Uint8 representa un entero sin signo (positivo) que va de 0 a 255 y solo ocupa 8 bits en la memoria. Por ejemplo, podemos representar el color naranja quemado como (Rojo: 204, Verde: 85, Azul: 0) para uint8 o (Rojo: 0,8, Verde: 0,33, Azul: 0) para float32. Float32 suele ser la mejor opción, ya que ofrece más precisión y ya viene normalizado, lo que ayuda a mejorar el entrenamiento. Sin embargo, uint8 ahorra memoria, y esto puede ser mejor dependiendo de las limitaciones de memoria. El uso de float32 en convert_image_dtype normaliza automáticamente los valores.
• En la segmentación binaria, esperamos que nuestras máscaras tengan forma (lote, altura, ancho, canales), con canales = 1. En otras palabras, queremos que una clase (coche) esté representada por el número 1 y la otra clase (fondo) esté representada por el número 0. No hay ninguna razón para que el número de canales sea 3, como en las imágenes RGB. Desafortunadamente, después de la decodificación, viene con tres canales, y el número de clase se repite tres veces. Para solucionar esto, usamos tf.math.reduce_max(mask, axis=-1, keepdims=True) para tomar el máximo de los valores en los tres canales y deshacernos del resto. Entonces, un valor de canal de (1,1,1) se reduce a (1) y un valor de canal de (0,0,0) se reduce a (0).
• Finalmente, redimensionamos las imágenes/máscaras a nuestras dimensiones deseadas (pequeñas). Ten en cuenta que las imágenes que mostré anteriormente del coche con la máscara de referencia se ven borrosas; esta reducción de escala se hizo a propósito para reducir la carga computacional y permitir que el entrenamiento ocurra relativamente rápido. Usar method=’nearest’ como predeterminado es una buena idea; de lo contrario, la función siempre devolverá un float32, lo cual es malo si deseas que esté en uint8.

Una vez que tengamos nuestros conjuntos de datos organizados, ahora podemos ver nuestras imágenes:

# Ver imágenes y etiquetas asociadaspara imágenes, máscaras en batched_val_dataset.take(1):    número_coche = 0    for ranura_imagen in range(16):        ax = plt.subplot(4, 4, ranura_imagen + 1)        if ranura_imagen % 2 == 0:            plt.imshow((imágenes[número_coche]))             nombre_clase = 'Imagen'        else:            plt.imshow(máscaras[número_coche], cmap = 'gray')            plt.colorbar()            nombre_clase = 'Máscara'            número_coche += 1                    plt.title(nombre_clase)        plt.axis("off")

Aquí estamos utilizando el método .take() para ver el primer lote de datos en nuestro conjunto de datos batched_val_dataset. Dado que estamos haciendo una segmentación binaria, queremos que nuestra máscara solo contenga dos valores: 0 y 1. La representación de las barras de color en la máscara confirma que tenemos la configuración correcta. Note que hemos agregado el argumento cmap = ‘gray’ a la función imshow() de la máscara para indicarle a plt que queremos que estas imágenes se presenten en escala de grises.

Construyendo un modelo U-Net simple

En una carta fechada el 5 de febrero de 1675 a su rival Robert Hooke, Isaac Newton dijo:

“Si he logrado ver más lejos, ha sido porque he subido a hombros de gigantes”.

En esta misma línea, nos apoyaremos en los investigadores anteriores de aprendizaje automático que han descubierto qué arquitecturas funcionan mejor para la tarea de segmentación de imágenes. No es una mala idea experimentar con arquitecturas propias; sin embargo, los investigadores que nos han precedido han recorrido muchos caminos sin salida para descubrir los modelos que funcionan. Estas arquitecturas no son necesariamente las mejores, ya que la investigación aún está en curso y es posible que se encuentre una mejor arquitectura.

Una de las arquitecturas más conocidas se llama U-Net, así llamada porque las partes de downsampling y upsampling de la red se pueden visualizar como una U (como se ve a continuación). En un artículo titulado “U-Net: Convolutional Networks for Biomedical Image Segmentation” escrito por Ronneberger, Fisher y Brox [1], los autores describen cómo crear una red convolucional completa (FCN) que funciona de manera efectiva para la segmentación de imágenes. Completamente convolucional significa que no hay capas densamente conectadas; todas las capas son convolucionales.

Hay algunas cosas a tener en cuenta:

• La red consta de una serie de bloques repetitivos de dos capas convolucionales, con padding = ‘same’ y stride = 1 para que las salidas de las convoluciones no se reduzcan dentro del bloque.
• Cada bloque está seguido por una capa de max pooling, que reduce a la mitad el ancho y la altura del mapa de características.
• El siguiente bloque duplica el número de filtros. Y el patrón continúa. Este patrón de reducir el espacio de características mientras se aumenta el número de filtros debería ser familiar si has estudiado CNNs. Esto completa lo que los autores llaman el “camino de contracción”.
• La capa “bottleneck” se encuentra en la parte inferior de la ‘U’. Esta capa captura características altamente abstractas (líneas, curvas, ventanas, puertas, etc.), pero con una resolución espacial significativamente reducida.
• A continuación comienza lo que ellos llaman el “camino de expansión”. En resumen, esto revierte las contracciones, con cada bloque consistiendo nuevamente en dos capas convolucionales. Cada bloque va seguido de una capa de upsampling, que en TensorFlow llamamos capa Conv2DTranspose. Esto toma un mapa de características más pequeño y duplica la altura y el ancho.
• El siguiente bloque luego reduce a la mitad el número de filtros. Repite el proceso hasta obtener la misma altura y anchura que las imágenes con las que comenzaste. Finalmente, termina con una capa convolucional 1×1 para reducir el número de canales a 1. Queremos terminar con un solo canal porque esto es segmentación binaria, por lo que deseamos un solo filtro donde los valores de los píxeles correspondan a nuestras dos clases. Usamos una activación sigmoidal para ajustar los valores de los píxeles entre 0 y 1.
• También hay conexiones de salto en la arquitectura U-Net, lo que permite que la red retenga información espacial detallada incluso después del downsampling y luego upsampling. Normalmente se pierde mucha información en este proceso. Al pasar la información de un bloque de contracción y en el bloque de expansión correspondiente, podemos preservar esta información espacial. Hay una simetría agradable en la arquitectura.

Comenzaremos haciendo una versión simple del U-Net. Esto será un FCN, pero sin conexiones residuales y sin capas de max pooling.

data_augmentation = tf.keras.Sequential([        tfl.RandomFlip(mode="horizontal", seed=42),        tfl.RandomRotation(factor=0.01, seed=42),        tfl.RandomContrast(factor=0.2, seed=42)])def get_model(img_size):    inputs = Input(shape=img_size + (3,))    x = data_augmentation(inputs)        # Camino de contracción    x = tfl.Conv2D(64, 3, strides=2, activation="relu", padding="same", kernel_initializer='he_normal')(x)     x = tfl.Conv2D(64, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)     x = tfl.Conv2D(128, 3, strides=2, activation="relu", padding="same", kernel_initializer='he_normal')(x)     x = tfl.Conv2D(128, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)     x = tfl.Conv2D(256, 3, strides=2, padding="same", activation="relu", kernel_initializer='he_normal')(x)     x = tfl.Conv2D(256, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)        # Camino de expansión    x = tfl.Conv2DTranspose(256, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)    x = tfl.Conv2DTranspose(256, 3, activation="relu", padding="same", kernel_initializer='he_normal', strides=2)(x)    x = tfl.Conv2DTranspose(128, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)    x = tfl.Conv2DTranspose(128, 3, activation="relu", padding="same", kernel_initializer='he_normal', strides=2)(x)    x = tfl.Conv2DTranspose(64, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)    x = tfl.Conv2DTranspose(64, 3, activation="relu", padding="same", kernel_initializer='he_normal', strides=2)(x)    outputs = tfl.Conv2D(1, 3, activation="sigmoid", padding="same")(x)    model = keras.Model(inputs, outputs)         return modelcustom_model = get_model(img_size=img_size)

Aquí tenemos la misma estructura básica que el U-Net, con un camino de contracción y un camino de expansión. Una cosa interesante a tener en cuenta es que en lugar de usar una capa de max pooling para reducir a la mitad el espacio de características, utilizamos una capa convolucional con strides = 2. Según Chollet [2], esto reduce el espacio de características a la mitad mientras preserva más información espacial que las capas de max pooling. Afirma que siempre que la información de ubicación sea importante (como en la segmentación de imágenes), es una buena idea evitar capas de max pooling destructivas y seguir utilizando convoluciones estratificadas en su lugar (esto es curioso, porque la famosa arquitectura U-Net sí usa max pooling). También observa que estamos realizando algunas técnicas de aumento de datos para ayudar a nuestro modelo a generalizarse a ejemplos no vistos.

Algunos detalles importantes: establecer el kernel_intializer en ‘he_normal’ para las activaciones ReLU hace una diferencia sorprendentemente grande en términos de estabilidad del entrenamiento. Inicialmente subestimé el poder de la inicialización del kernel. En lugar de inicializar los pesos al azar, la he_normalización inicializa los pesos con una media de 0 y una desviación estándar igual a la raíz cuadrada de (2 / número de unidades de entrada a la capa). En el caso de las CNN, el número de unidades de entrada se refiere al número de canales en los mapas de características de la capa anterior. Se ha encontrado que esto conduce a una convergencia más rápida, mitiga los gradientes desvanecidos y mejora el aprendizaje. Consulte la referencia [3] para obtener más detalles.

Métricas y Función de Pérdida

Hay varias métricas y funciones de pérdida comunes que se pueden usar para la segmentación binaria. Aquí, utilizaremos el coeficiente de dice como métrica y la correspondiente pérdida de dice para el entrenamiento, ya que esto es lo que requiere la competición.

Veamos primero las matemáticas detrás del coeficiente de dice:

El coeficiente de Dice se define como la intersección entre dos conjuntos (X e Y), dividido por la suma de cada conjunto, multiplicado por 2. El coeficiente de Dice estará entre 0 (si los conjuntos no tienen intersección) y 1 (si los conjuntos se superponen perfectamente). Ahora vemos por qué esto es una gran métrica para la segmentación de imágenes.

La ecuación anterior es una definición general del coeficiente de Dice; cuando se aplica a cantidades vectoriales (en lugar de conjuntos), usamos la definición más específica:

Aquí, estamos iterando sobre cada elemento (píxel) en cada máscara. x representa el i-ésimo píxel en la máscara predicha y y representa el píxel correspondiente en la máscara de verdad. En la parte superior estamos realizando el producto elemento a elemento, y en la parte inferior estamos sumando todos los elementos en cada máscara de manera independiente. N representa el número total de píxeles (que debería ser el mismo tanto para las máscaras predichas como para las máscaras objetivo). Recuerda que en nuestras máscaras, los números serán todos 0 o 1, por lo que un píxel con un valor de 1 en la máscara de verdad y un píxel correspondiente en la máscara predicha con un valor de 0 no contribuirá al puntaje de Dice, como se espera (1 x 0 = 0).

La pérdida de Dice se definirá simplemente como 1 – Puntaje de Dice. Dado que el puntaje de Dice está entre 0 y 1, la pérdida de Dice también estará entre 0 y 1. De hecho, la suma del puntaje de Dice y la pérdida de Dice debe ser igual a 1. Están inversamente relacionados.

Echemos un vistazo a cómo se implementa esto en código:

from tensorflow.keras import backend as Kdef dice_coef(y_true, y_pred, smooth=10e-6):    y_true_f = K.flatten(y_true)    y_pred_f = K.flatten(y_pred)    intersection = K.sum(y_true_f * y_pred_f)    dice = (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)    return dicedef dice_loss(y_true, y_pred):    return 1 - dice_coef(y_true, y_pred)

Aquí estamos aplanando dos máscaras 4-D (lote, altura, ancho, canales=1) en vectores 1-D, y calculando los puntajes de Dice para todas las imágenes en el lote. Ten en cuenta que agregamos un valor de suavizado tanto al numerador como al denominador para evitar tener un problema de 0/0 si las dos máscaras no se superponen.

Finalmente, comenzamos el entrenamiento. Estamos utilizando la detención temprana para evitar el sobreajuste y guardando el mejor modelo.

custom_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001,                                                        epsilon=1e-06),                                                         loss=[dice_loss],                                                         metrics=[dice_coef])callbacks_list = [    keras.callbacks.EarlyStopping(        monitor="val_loss",        patience=2,    ),    keras.callbacks.ModelCheckpoint(        filepath="best-custom-model",        monitor="val_loss",        save_best_only=True,    )]history = custom_model.fit(batched_train_dataset, epochs=20,                    callbacks=callbacks_list,                    validation_data=batched_val_dataset)

Podemos determinar los resultados de nuestro entrenamiento con el siguiente código:

def display(display_list):    plt.figure(figsize=(15, 15))    title = ['Imagen de entrada', 'Máscara verdadera', 'Máscara predicha']    for i in range(len(display_list)):        plt.subplot(1, len(display_list), i+1)        plt.title(title[i])        plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))        plt.axis('off')    plt.show()    def create_mask(pred_mask):    mask = pred_mask[..., -1] >= 0.5    pred_mask[..., -1] = tf.where(mask, 1, 0)    # Devolver solo la primera máscara del lote    return pred_mask[0]def show_predictions(model, dataset=None, num=1):    """    Muestra la primera imagen de cada una de las num lotes    """    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, ...]))])custom_model = keras.models.load_model("/kaggle/working/best-custom-model", custom_objects={'dice_coef': dice_coef, 'dice_loss': dice_loss})show_predictions(model = custom_model, dataset = batched_train_dataset, num = 6)

Después de 10 épocas, alcanzamos una puntuación de dados de validación superior de 0.8788. No es terrible, pero no es genial. En una GPU P100 esto me tomó alrededor de 20 minutos. Aquí hay una máscara de muestra para nuestra revisión:

Destacando algunos puntos interesantes:

• Observe que create_mask es la función que asigna valores de píxeles a 0 o 1. Un valor de píxel de < 0.5 se cortará a 0 y asignaremos ese píxel a la categoría “fondo”. Un valor ≥ 0.5 se aumentará a 1 y asignaremos ese píxel a la categoría “coche”.
• ¿Por qué las máscaras salieron amarillas y moradas, en lugar de negras y blancas? Usamos: tf.keras.preprocessing.image.array_to_img() para convertir la salida de la máscara de un tensor a una imagen PIL. Luego pasamos la imagen a plt.imshow(). Desde la documentación vemos que el mapa de colores predeterminado para imágenes de un solo canal es “viridis” (las imágenes RGB de 3 canales se muestran como están). El mapa de colores viridis transforma los valores bajos en un morado profundo y los valores altos en amarillo. Este mapa de colores aparentemente puede ayudar a las personas con daltonismo a obtener una vista precisa del color en una imagen. Podríamos haber solucionado esto agregando cmap=”escala de grises” como argumento, pero esto habría arruinado nuestra imagen de entrada. Ver más aquí en este enlace.

Construyendo la U-Net completa

Ahora pasamos a usar la arquitectura completa de U-Net, con conexiones residuales, capas de max pooling y la inclusión de capas de dropout para la regularización. Observe la ruta de contracción, la capa bottleneck y la ruta de expansión. Las capas de dropout se pueden agregar en la ruta de contracción, al final de cada bloque.

def conv_block(inputs=None, n_filters=64, dropout_prob=0, max_pooling=True):    conv = Conv2D(n_filters,                    3,                     activation='relu',                  padding='same',                  kernel_initializer='he_normal')(inputs)    conv = Conv2D(n_filters,                    3,                     activation='relu',                  padding='same',                  kernel_initializer='he_normal')(conv)    if dropout_prob > 0:        conv = Dropout(dropout_prob)(conv)    if max_pooling:        next_layer = MaxPooling2D(pool_size=(2, 2))(conv)    else:        next_layer = conv    skip_connection = conv    return next_layer, skip_connectiondef upsampling_block(expansive_input, contractive_input, n_filters=64):    up = Conv2DTranspose(        n_filters,            3,            strides=(2, 2),        padding='same',        kernel_initializer='he_normal')(expansive_input)    # Combinar la salida anterior y el contractive_input    merge = concatenate([up, contractive_input], axis=3)    conv = Conv2D(n_filters,                     3,                       activation='relu',                  padding='same',                  kernel_initializer='he_normal')(merge)    conv = Conv2D(n_filters,                     3,                       activation='relu',                  padding='same',                  kernel_initializer='he_normal')(conv)    return convdef unet_model(input_size=(96, 128, 3), n_filters=64, n_classes=1):    inputs = Input(input_size)        inputs = data_augmentation(inputs)    # Ruta de contracción (codificación)    cblock1 = conv_block(inputs, n_filters)    cblock2 = conv_block(cblock1[0], n_filters*2)    cblock3 = conv_block(cblock2[0], n_filters*4)    cblock4 = conv_block(cblock3[0], n_filters*8, dropout_prob=0.3)    # Capa bottleneck    cblock5 = conv_block(cblock4[0], n_filters*16, dropout_prob=0.3, max_pooling=False)        # Ruta de expansión (decodificación)    ublock6 = upsampling_block(cblock5[0], cblock4[1],  n_filters*8)    ublock7 = upsampling_block(ublock6, cblock3[1],  n_filters*4)    ublock8 = upsampling_block(ublock7, cblock2[1],  n_filters*2)    ublock9 = upsampling_block(ublock8, cblock1[1],  n_filters)    conv9 = Conv2D(n_filters,                   3,                   activation='relu',                   padding='same',                   kernel_initializer='he_normal')(ublock9)    conv10 = Conv2D(n_classes, 1, padding='same', activation="sigmoid")(conv9)    model = tf.keras.Model(inputs=inputs, outputs=conv10)    return model

Luego compilamos la U-Net. Estoy usando 64 filtros para el primer bloque convolucional. Este es un hiperparámetro que querrías ajustar para obtener resultados óptimos.

unet = unet_model(input_size=(img_height, img_width, num_channels), n_filters=64, n_classes=1)unet.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001, epsilon=1e-06),             loss=[dice_loss],              metrics=[dice_coef])callbacks_list = [    keras.callbacks.EarlyStopping(        monitor="val_loss",        patience=2,    ),    keras.callbacks.ModelCheckpoint(        filepath="best-u_net-model",        monitor="val_loss",        save_best_only=True,    )]history = unet.fit(batched_train_dataset, epochs=20,                    callbacks=callbacks_list,                    validation_data=batched_val_dataset)

Después de 16 épocas, obtengo un puntaje de validación de dice de 0.9416, mucho mejor que con la U-Net simple. Esto no debería ser muy sorprendente; al observar la cantidad de parámetros, hay un aumento de un orden de magnitud desde la U-Net simple hasta la U-Net completa. En una GPU P100, esto me llevó unos 32 minutos. Luego echamos un vistazo a las predicciones:

unet = keras.models.load_model("/kaggle/working/best-u_net-model", custom_objects={'dice_coef': dice_coef, 'dice_loss': dice_loss})show_predictions(model = unet, dataset = batched_train_dataset, num = 6)

Estas predicciones son mucho mejores. Al observar varias predicciones, se puede notar que la antena que sobresale de los autos es un desafío para la red. Dado que las imágenes están muy pixeladas, no puedo culpar a la red por no detectar esto.

Para mejorar el rendimiento, uno debería ajustar hiperparámetros que incluyen:

• Número de bloques de submuestreo y sobre muestreo
• Número de filtros
• Resolución de la imagen
• Tamaño del conjunto de entrenamiento
• Función de pérdida (quizás combinando la pérdida de dice con la pérdida de entropía cruzada binaria)
• Ajustar los parámetros del optimizador. La estabilidad del entrenamiento parece ser un problema para ambos modelos. Según la documentación del optimizador Adam: “El valor predeterminado de 1e-7 para epsilon puede no ser un buen valor predeterminado en general”. Aumentar epsilon en un orden de magnitud o más puede ayudar con la estabilidad del entrenamiento.

Ya podemos ver el camino hacia una puntuación excelente en el desafío de Carvana. ¡Qué lástima que haya terminado!

Resumen

Este artículo fue un análisis en profundidad sobre el tema de la segmentación de imágenes, específicamente la segmentación binaria. Si te llevas algo, recuerda lo siguiente:

• El objetivo de la segmentación de imágenes es encontrar una correspondencia entre los valores de píxeles de entrada en una imagen y los números de salida que tu modelo pueda usar para asignar clases a cada píxel.
• Uno de los primeros pasos es organizar tus imágenes en objetos Dataset de TensorFlow y echar un vistazo a tus imágenes y máscaras correspondientes.
• No es necesario reinventar la rueda cuando se trata de la arquitectura del modelo: sabemos por experiencia que una U-Net funciona bien.
• El puntaje de dice es una métrica común que se utiliza para monitorear el éxito de las predicciones de tu modelo. También podemos obtener nuestra función de pérdida a partir de esto.

Trabajo futuro podría consistir en convertir las capas de max pooling en la arquitectura canónica de U-Net en capas de convolución con saltos.

¡Buena suerte en tus problemas de segmentación de imágenes!

Referencias

[1] O. Ronneberger, P. Fischer, y T. Brox, U-Net: Convolutional Networks for Biomedical Image Segmentation (2015), MICCAI 2015 International Conference

[2] F. Chollet, Deep Learning with Python (2021), Manning Publications Co.

[3] K. He, X. Zhang, S. Ren, J. Sun, Adentrándose en los Rectificadores: Superando el Desempeño a Nivel Humano en la Clasificación de ImageNet (2015), Conferencia Internacional sobre Visión por Computadora (ICCV)

We will continue to update Zepes; if you have any questions or suggestions, please contact us!