Tutorial Completo: Segmentación de Imágenes con U-Net desde Cero

¿Qué es U-Net y por qué es tan poderosa para la segmentación de imágenes?

La segmentación de imágenes es una de las tareas más fascinantes y desafiantes dentro del campo de la visión por computadora. A diferencia de la clasificación de imágenes, donde simplemente etiquetamos una imagen completa, la segmentación nos permite identificar y delimitar con precisión cada objeto dentro de una escena. En este contexto, U-Net se ha convertido en una de las arquitecturas más populares y efectivas del mundo del Deep Learning.

Desarrollada originalmente en 2015 por Olaf Ronneberger, Philipp Fischer y Thomas Brox en la Universidad de Friburgo, U-Net fue diseñada específicamente para la segmentación de imágenes biomédicas. Sin embargo, su versatilidad la ha llevado a ser adoptada en industrias tan diversas como la agricultura de precisión, la detección de defectos industriales, el análisis satelital y, por supuesto, el sector de salud digital que está creciendo aceleradamente en Colombia.

Arquitectura de U-Net: Entendiendo su estructura única

El nombre "U-Net" proviene de su forma característica en U cuando se visualiza su arquitectura. Esta red neuronal convolucional está compuesta por dos caminos principales:

1. El Camino Contratante (Encoder)

También conocido como el "camino descendente", este componente sigue la arquitectura típica de una red convolucional. Consiste en la aplicación repetida de:

• Dos convoluciones 3x3 (sin padding), cada una seguida de una función de activación ReLU
• Una operación de max pooling 2x2 con stride 2 para el downsampling
• Duplicación del número de canales de características en cada paso

2. El Camino Expansivo (Decoder)

Este es el componente que hace única a U-Net. El camino expansivo realiza:

• Upsampling del mapa de características
• Una convolución 2x2 ("up-convolution") que reduce a la mitad el número de canales
• Concatenación con el mapa de características correspondiente del camino contratante (las famosas "skip connections")
• Dos convoluciones 3x3, cada una seguida de ReLU

Las skip connections son el elemento diferenciador de U-Net. Al concatenar información de alta resolución del encoder con las características del decoder, la red puede producir segmentaciones mucho más precisas, especialmente en los bordes de los objetos.

Implementación práctica de U-Net con TensorFlow y Keras

A continuación, te mostraremos cómo implementar U-Net desde cero usando Python con TensorFlow 2.x y Keras. Este código ha sido probado y optimizado para funcionar en entornos con recursos limitados, algo importante considerando que muchos desarrolladores colombianos trabajan con equipos de gama media.

Instalación de dependencias

pip install tensorflow numpy matplotlib opencv-python scikit-learn

Construcción del modelo U-Net

import tensorflow as tf
from tensorflow.keras import layers, Model

def conv_block(inputs, num_filters):
    x = layers.Conv2D(num_filters, 3, padding="same")(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)
    x = layers.Conv2D(num_filters, 3, padding="same")(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)
    return x

def encoder_block(inputs, num_filters):
    x = conv_block(inputs, num_filters)
    p = layers.MaxPool2D((2, 2))(x)
    return x, p

def decoder_block(inputs, skip_features, num_filters):
    x = layers.Conv2DTranspose(num_filters, (2, 2), strides=2, padding="same")(inputs)
    x = layers.Concatenate()([x, skip_features])
    x = conv_block(x, num_filters)
    return x

def build_unet(input_shape):
    inputs = layers.Input(input_shape)
    s1, p1 = encoder_block(inputs, 64)
    s2, p2 = encoder_block(p1, 128)
    s3, p3 = encoder_block(p2, 256)
    s4, p4 = encoder_block(p3, 512)
    b1 = conv_block(p4, 1024)
    d1 = decoder_block(b1, s4, 512)
    d2 = decoder_block(d1, s3, 256)
    d3 = decoder_block(d2, s2, 128)
    d4 = decoder_block(d3, s1, 64)
    outputs = layers.Conv2D(1, 1, padding="same", activation="sigmoid")(d4)
    model = Model(inputs, outputs, name="U-Net")
    return model

Compilación y entrenamiento del modelo

model = build_unet((256, 256, 3))
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),
    loss="binary_crossentropy",
    metrics=["accuracy", tf.keras.metrics.MeanIoU(num_classes=2)]
)
model.summary()

Preparación del dataset para segmentación

Uno de los aspectos más críticos al trabajar con U-Net es la preparación correcta de los datos. Para este tutorial, utilizaremos el enfoque de augmentación de datos, fundamental cuando trabajamos con datasets pequeños (situación común en proyectos de IA en Colombia donde los datos etiquetados son escasos).

import numpy as np
import cv2
from sklearn.model_selection import train_test_split

def load_data(image_paths, mask_paths, img_size=(256, 256)):
    images = []
    masks = []
    for img_path, mask_path in zip(image_paths, mask_paths):
        img = cv2.imread(img_path)
        img = cv2.resize(img, img_size)
        img = img / 255.0
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        mask = cv2.resize(mask, img_size)
        mask = mask / 255.0
        mask = np.expand_dims(mask, axis=-1)
        images.append(img)
        masks.append(mask)
    return np.array(images), np.array(masks)

def data_augmentation(image, mask):
    if np.random.rand() > 0.5:
        image = np.fliplr(image)
        mask = np.fliplr(mask)
    if np.random.rand() > 0.5:
        image = np.flipud(image)
        mask = np.flipud(mask)
    return image, mask

Métricas de evaluación: Más allá de la precisión simple

Para evaluar correctamente un modelo de segmentación, necesitamos métricas específicas. Las más utilizadas en la comunidad de Deep Learning son:

Intersection over Union (IoU)

También conocido como el coeficiente de Jaccard, mide la superposición entre la predicción y la máscara real. Un IoU de 0.8 o superior generalmente se considera un buen resultado para aplicaciones prácticas.

Coeficiente de Dice

Similar al IoU pero con diferente formulación matemática. Es especialmente útil cuando las clases están desbalanceadas, lo cual es muy común en imágenes médicas.

def dice_coefficient(y_true, y_pred, smooth=1):
    y_true_flat = tf.keras.backend.flatten(y_true)
    y_pred_flat = tf.keras.backend.flatten(y_pred)
    intersection = tf.keras.backend.sum(y_true_flat * y_pred_flat)
    return (2. * intersection + smooth) / (
        tf.keras.backend.sum(y_true_flat) + 
        tf.keras.backend.sum(y_pred_flat) + smooth
    )

def dice_loss(y_true, y_pred):
    return 1 - dice_coefficient(y_true, y_pred)

Aplicaciones de U-Net en el contexto colombiano

Colombia presenta un ecosistema único de oportunidades para la aplicación de técnicas de segmentación de imágenes. Algunos casos de uso especialmente relevantes incluyen:

Agricultura de Precisión

Con más de 7 millones de hectáreas cultivadas según el DANE, Colombia tiene un enorme potencial para aplicar U-Net en el análisis de imágenes satelitales y de drones para detectar enfermedades en cultivos de café, banano y flores, sectores clave de la economía nacional.

Telemedicina y Diagnóstico Médico

El MinSalud ha impulsado la digitalización del sector salud, y la segmentación de imágenes médicas (radiografías, tomografías) mediante U-Net puede apoyar el diagnóstico en regiones con acceso limitado a especialistas.

Monitoreo Ambiental

La deforestación en la Amazonía colombiana es un problema crítico. U-Net puede procesar imágenes satelitales para detectar cambios en la cobertura boscosa con alta precisión, apoyando iniciativas del IDEAM y organizaciones ambientales.

Optimización del modelo para recursos limitados

Una consideración importante para desarrolladores en Colombia es la optimización del modelo para entornos con recursos computacionales limitados. Algunas estrategias efectivas incluyen:

• Quantización del modelo: Reduce el tamaño del modelo hasta 4x con mínima pérdida de precisión
• Pruning: Elimina conexiones redundantes en la red neuronal
• Transfer Learning: Usa pesos pre-entrenados en ImageNet para acelerar el entrenamiento
• Mixed Precision Training: Utiliza float16 en lugar de float32 para reducir el uso de memoria

from tensorflow.keras.mixed_precision import set_global_policy
set_global_policy("mixed_float16")

policy = tf.keras.mixed_precision.Policy("mixed_float16")
print("Compute dtype: %s" % policy.compute_dtype)
print("Variable dtype: %s" % policy.variable_dtype)

Desafíos y consideraciones futuras

A pesar de su potencia, U-Net presenta algunos desafíos que los profesionales de IA en Colombia deben considerar:

• Necesidad de datos etiquetados: La anotación manual de imágenes es costosa y requiere expertise. Plataformas como Label Studio o Roboflow pueden facilitar este proceso.
• Costo computacional: El entrenamiento completo requiere GPUs potentes. Servicios como Google Colab Pro o AWS ofrecen alternativas accesibles.
• Generalización: Los modelos entrenados en un dominio pueden no funcionar bien en otro sin fine-tuning.

Conclusión

U-Net representa una de las arquitecturas más elegantes y efectivas en el campo de la visión por computadora. Su diseño único con skip connections la hace especialmente poderosa para tareas donde la precisión espacial es crítica. Para los desarrolladores e investigadores de IA en Colombia, dominar esta arquitectura abre puertas a aplicaciones de alto impacto en sectores como la salud, la agricultura y el medio ambiente. Te invitamos a implementar el código de este tutorial, experimentar con tus propios datasets y compartir tus resultados en los comentarios. ¿Ya estás usando U-Net en algún proyecto? ¡Cuéntanos tu experiencia!