Segmentación de imágenes con CNNs

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

La segmentación de imágenes es la tarea de asignar una etiqueta a cada píxel de una imagen. Mientras que la clasificación responde «¿qué hay en esta imagen?» y la detección responde «¿dónde está cada objeto?» con un rectángulo, la segmentación responde «¿a qué clase pertenece cada píxel?».

El resultado de una segmentación es un mapa de segmentación (o segmentation mask): una imagen del mismo tamaño que la original donde cada píxel tiene un valor entero que representa su clase. Esto permite delimitar los objetos con precisión a nivel de píxel.

📊 Clasificación vs. Detección vs. Segmentación

Para entender bien qué aporta la segmentación, comparémosla con las otras dos tareas fundamentales de visión por computador:

🧩 Tipos de segmentación

Existen tres tipos principales de segmentación, cada uno con un nivel diferente de detalle:

Segmentación semántica

Asigna una clase a cada píxel, pero no distingue entre instancias del mismo objeto. Si hay tres coches en la imagen, todos sus píxeles reciben la etiqueta «coche», pero no se sabe que son tres coches distintos.

Donde C es el número de clases. Cada píxel de la salida toma un valor entero que indica su clase.

Segmentación por instancias

Detecta y segmenta cada instancia individual de un objeto. Los tres coches del ejemplo anterior recibirían cada uno una máscara distinta: coche_1, coche_2, coche_3. Sin embargo, las regiones de fondo no se clasifican.

Segmentación panóptica

Combina ambas: todos los píxeles reciben una clase (como la semántica) y, además, las instancias de objetos «contables» (things: personas, coches) se distinguen individualmente. Las regiones como cielo, carretera o hierba (stuff) solo reciben la clase semántica sin distinguir instancias.

🌍 Aplicaciones de la segmentación

La segmentación es una de las tareas más demandadas en visión por computador. Algunas aplicaciones clave:

🔬 Técnicas clásicas de segmentación

Antes de las redes neuronales, la segmentación se basaba en técnicas de procesamiento de imagen que explotan diferencias de color, intensidad o textura. Aunque hoy las redes profundas dominan, estos métodos siguen siendo útiles como preprocesamiento, para datasets pequeños, o cuando la interpretabilidad es clave.

Umbralización (thresholding)

La técnica más simple: se elige un umbral T y cada píxel se clasifica según su intensidad:

Funciona bien cuando hay buen contraste entre el objeto y el fondo. Sin embargo, elegir T manualmente es poco robusto.

Método de Otsu

El método de Otsu (1979) automatiza la elección del umbral. Busca el valor de T que maximiza la varianza entre clases (fondo vs. primer plano):

Donde \omega_0, \omega_1 son las proporciones de píxeles en cada clase y \mu_0, \mu_1 sus medias de intensidad. Se calcula para todos los posibles T \in [0, 255] y se elige el que maximiza \sigma_B^2.

Detección de bordes

Los bordes son transiciones abruptas de intensidad que a menudo delimitan objetos. Los operadores clásicos calculan el gradiente de la imagen:

• Sobel: kernels de 3×3 que aproximan las derivadas parciales.
• Canny: pipeline más sofisticado (suavizado Gaussiano → gradiente → supresión de no-máximos → doble umbral + histéresis).
• Laplaciano de Gausianas (LoG): detecta cruces por cero de la segunda derivada.

Watershed (línea divisoria de aguas)

Inspirado en la topografía: se interpreta la imagen como un relieve donde la intensidad es la altitud. Se «inundan» los valles y las líneas divisorias donde se encuentran aguas de distintas cuencas definen los contornos de los objetos.

Superpíxeles

Los superpíxeles agrupan píxeles vecinos con color y textura similar en regiones compactas. No son una segmentación semántica, pero reducen la complejidad de la imagen de millones de píxeles a cientos o miles de regiones homogéneas.

El algoritmo más popular es SLIC (Simple Linear Iterative Clustering), que aplica k-means en el espacio (L, a, b, x, y) combinando color (CIELAB) y posición espacial.

🧠 ¿Por qué Deep Learning para segmentación?

Las técnicas clásicas funcionan bien en escenarios controlados, pero tienen limitaciones fundamentales en imágenes naturales complejas:

🔄 FCN: de clasificación a segmentación

En 2015, Long et al. propusieron las Fully Convolutional Networks (FCN), la primera arquitectura deep learning exitosa para segmentación semántica. La idea revolucionaria fue simple pero potente: reemplazar las capas fully-connected de una red de clasificación por capas convolucionales, permitiendo que la red acepte imágenes de cualquier tamaño y produzca un mapa de probabilidades denso.

La idea clave

Una CNN de clasificación como VGG-16 tiene dos partes:

1. Backbone convolucional: extrae features (ej. 7×7×512 en VGG-16)
2. Cabeza FC (fully-connected): aplana las features y produce un vector de clases

El problema de las capas FC es que destruyen la información espacial: sabemos qué hay en la imagen, pero no dónde. La solución de FCN fue convertir las capas FC en convoluciones 1×1:

Upsampling: recuperar la resolución

Después de pasar por el backbone, el mapa de features tiene una resolución mucho menor que la imagen original (ej. 7×7 para una entrada de 224×224, un factor ×32). Para obtener una predicción a la resolución original, FCN usa convoluciones transpuestas (transposed convolutions) para hacer upsampling:

Donde s es el stride, p el padding y k el tamaño del kernel. Con s=2, k=4, p=1 se consigue duplicar la resolución espacial.

Variantes: FCN-32s, FCN-16s, FCN-8s

El paper original propuso tres variantes que difieren en la granularidad del upsampling:

Skip connections en FCN

Las skip connections son la innovación que hace a FCN-8s muy superior a FCN-32s. La idea es combinar predicciones de distintas profundidades:

• Las capas profundas tienen features semánticamente ricas pero espacialmente burdas (saben qué hay pero no dónde exactamente).
• Las capas superficiales tienen features con alta resolución espacial pero menos semántica (saben dónde pero no qué).
• Al sumar ambas, se obtiene lo mejor de cada una.

Limitaciones de FCN

A pesar de su impacto, FCN tiene varias limitaciones que motivaron arquitecturas posteriores:

• Upsampling brusco: incluso FCN-8s produce bordes poco definidos porque el upsampling ×8 es grande.
• Sin contexto global: no captura relaciones a larga distancia en la imagen.
• No aprovecha features intermedias: solo usa 2-3 niveles de skip connections.
• Entrenamiento complejo: requiere inicializar desde un modelo pre-entrenado (VGG-16).

🏗️ Arquitectura Encoder-Decoder

La arquitectura encoder-decoder es el paradigma dominante en segmentación semántica. La idea es intuitiva:

1. Encoder (codificador): comprime la imagen progresivamente, extrayendo features cada vez más abstractas y reduciendo la resolución espacial.
2. Decoder (decodificador): expande las features comprimidas de vuelta a la resolución original, generando la máscara de segmentación.

A diferencia de FCN, donde el upsampling es un paso final brusco (×8, ×16 o ×32), el decoder reconstruye gradualmente la resolución, permitiendo refinamientos progresivos.

🔗 Conexión con Autoencoders

La arquitectura encoder-decoder para segmentación está inspirada en los autoencoders, un tipo de red neuronal que aprende a comprimir y reconstruir datos:

La diferencia clave entre un autoencoder y una red de segmentación es:

🏆 U-Net: la arquitectura estrella

U-Net (Ronneberger et al., 2015) es probablemente la arquitectura más influyente en segmentación. Diseñada originalmente para segmentación de imágenes biomédicas, su elegancia y efectividad la han convertido en el estándar de facto para muchas tareas de segmentación.

¿Por qué se llama «U-Net»?

El nombre viene de la forma de U que tiene la arquitectura cuando se dibuja: el encoder desciende por la izquierda, el bottleneck está en el fondo, y el decoder asciende por la derecha, con conexiones horizontales (skip connections) que cruzan de izquierda a derecha.

Camino contractivo (Encoder)

El encoder de U-Net sigue un patrón clásico de CNN, con 4 bloques de downsampling. Cada bloque contiene:

1. Dos convoluciones 3×3 + ReLU (sin padding en el paper original)
2. Max pooling 2×2 con stride 2 (reduce resolución a la mitad)

El número de canales se duplica en cada bloque: 64 → 128 → 256 → 512 → 1024 (bottleneck).

Bottleneck

En la base de la U, el bottleneck procesa las features más comprimidas (la resolución más baja y el mayor número de canales: 1024). Consta de dos convoluciones 3×3 + ReLU, sin max pooling.

Camino expansivo (Decoder)

El decoder reconstruye la resolución progresivamente. Cada bloque:

1. Upsampling 2×2 (convolución transpuesta) que duplica la resolución y reduce canales a la mitad
2. Concatenación con el mapa de features correspondiente del encoder (skip connection)
3. Dos convoluciones 3×3 + ReLU

Skip connections: la clave del éxito

Las skip connections son lo que hace a U-Net especial. En lugar de sumar features como en FCN, U-Net las concatena:

Capa final

La última capa de U-Net es una convolución 1×1 que mapea los 64 canales del último bloque del decoder a C canales (uno por clase). Se aplica softmax (o sigmoid para segmentación binaria) para obtener probabilidades:

¿Qué hizo especial a U-Net?

• Funciona con pocos datos: diseñada para imagen biomédica donde las muestras anotadas son escasas. Data augmentation agresiva compensa.
• Skip connections por concatenación: preservan toda la información espacial del encoder.
• Decoder simétrico: el decoder es un espejo del encoder, creando una arquitectura balanceada y elegante.
• Loss ponderada: los píxeles de borde entre objetos cercanos reciben mayor peso, mejorando la separación de instancias tocándose.

Donde d_1 y d_2 son las distancias al borde de las dos células más cercanas. Los píxeles entre células cercanas reciben mayor peso, forzando a la red a aprender a separarlas.

import torch
import torch.nn as nn

class DoubleConv(nn.Module):
    """Bloque básico de U-Net: (Conv3x3 → BN → ReLU) × 2"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)


class UNet(nn.Module):
    def __init__(self, in_channels=3, num_classes=2):
        super().__init__()

        # Encoder (camino contractivo)
        self.enc1 = DoubleConv(in_channels, 64)
        self.enc2 = DoubleConv(64, 128)
        self.enc3 = DoubleConv(128, 256)
        self.enc4 = DoubleConv(256, 512)
        self.pool = nn.MaxPool2d(2)

        # Bottleneck
        self.bottleneck = DoubleConv(512, 1024)

        # Decoder (camino expansivo)
        self.up4 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
        self.dec4 = DoubleConv(1024, 512)   # 512+512 por concat

        self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.dec3 = DoubleConv(512, 256)    # 256+256 por concat

        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.dec2 = DoubleConv(256, 128)    # 128+128 por concat

        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.dec1 = DoubleConv(128, 64)     # 64+64 por concat

        # Clasificador final (1×1 conv)
        self.final = nn.Conv2d(64, num_classes, 1)

    def forward(self, x):
        # Encoder
        e1 = self.enc1(x)               # → 64 canales
        e2 = self.enc2(self.pool(e1))    # → 128 canales
        e3 = self.enc3(self.pool(e2))    # → 256 canales
        e4 = self.enc4(self.pool(e3))    # → 512 canales

        # Bottleneck
        b = self.bottleneck(self.pool(e4))  # → 1024 canales

        # Decoder con skip connections
        d4 = self.dec4(torch.cat([self.up4(b), e4], dim=1))
        d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1))
        d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
        d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))

        return self.final(d1)  # → num_classes canales


# Ejemplo de uso
model = UNet(in_channels=3, num_classes=21)  # 21 clases PASCAL VOC
x = torch.randn(1, 3, 256, 256)
out = model(x)
print(f"Input: {x.shape}")     # [1, 3, 256, 256]
print(f"Output: {out.shape}")  # [1, 21, 256, 256]

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

def double_conv(x, filters):
    """Bloque (Conv3x3 → BN → ReLU) × 2"""
    x = layers.Conv2D(filters, 3, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters, 3, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    return x

def build_unet(input_shape=(256, 256, 3), num_classes=2):
    inputs = layers.Input(shape=input_shape)

    # Encoder
    e1 = double_conv(inputs, 64)
    p1 = layers.MaxPool2D(2)(e1)

    e2 = double_conv(p1, 128)
    p2 = layers.MaxPool2D(2)(e2)

    e3 = double_conv(p2, 256)
    p3 = layers.MaxPool2D(2)(e3)

    e4 = double_conv(p3, 512)
    p4 = layers.MaxPool2D(2)(e4)

    # Bottleneck
    b = double_conv(p4, 1024)

    # Decoder con skip connections (concatenación)
    u4 = layers.Conv2DTranspose(512, 2, strides=2, padding='same')(b)
    u4 = layers.Concatenate()([u4, e4])
    d4 = double_conv(u4, 512)

    u3 = layers.Conv2DTranspose(256, 2, strides=2, padding='same')(d4)
    u3 = layers.Concatenate()([u3, e3])
    d3 = double_conv(u3, 256)

    u2 = layers.Conv2DTranspose(128, 2, strides=2, padding='same')(d3)
    u2 = layers.Concatenate()([u2, e2])
    d2 = double_conv(u2, 128)

    u1 = layers.Conv2DTranspose(64, 2, strides=2, padding='same')(d2)
    u1 = layers.Concatenate()([u1, e1])
    d1 = double_conv(u1, 64)

    # Capa final
    outputs = layers.Conv2D(num_classes, 1, activation='softmax')(d1)

    return Model(inputs, outputs, name='UNet')

model = build_unet(num_classes=21)
model.summary()  # ~31M parámetros

🔀 Variantes de U-Net

El éxito de U-Net inspiró una familia de arquitecturas que mantienen la estructura encoder-decoder con skip connections pero introducen mejoras específicas. Veamos las más relevantes.

Attention U-Net

Attention U-Net (Oktay et al., 2018) introduce attention gates en las skip connections para que el modelo aprenda a enfocarse en las regiones relevantes y suprimir features irrelevantes.

En U-Net estándar, las skip connections pasan todas las features del encoder al decoder. Pero no toda esa información es útil — en imagen médica, por ejemplo, gran parte del fondo es irrelevante. Los attention gates aprenden a «filtrar» qué features pasar:

Donde x_i son las features del encoder, g_i la señal del decoder (gating signal), y \alpha_i \in [0,1] es el coeficiente de atención que pondera cada feature espacial. La salida es \hat{x}_i = \alpha_i \cdot x_i.

U-Net++ (Nested U-Net)

U-Net++ (Zhou et al., 2018) rediseña las skip connections añadiendo bloques convolucionales intermedios entre el encoder y el decoder. En lugar de conexiones directas, crea una red densa de sub-redes anidadas.

La intuición es que el gap semántico entre encoder y decoder puede ser grande (features en escala ×1 vs. features en escala ×16). Los nodos intermedios de U-Net++ reducen progresivamente ese gap:

Donde [\cdot] denota concatenación y \mathcal{H} es un bloque convolucional. Cada nodo recibe features de todos los nodos anteriores en su nivel, más el upsampling del nivel inferior.

V-Net: segmentación 3D

V-Net (Milletari et al., 2016) adapta la idea de U-Net a datos volumétricos 3D (como CT scans o MRI). Las principales diferencias:

• Convoluciones 3D en lugar de 2D (kernels 5×5×5)
• Skip connections residuales (suma) en lugar de concatenación
• Dice loss como función de pérdida (veremos esto en detalle más adelante)
• Convoluciones strided en lugar de max pooling para downsampling

📐 SegNet

SegNet (Badrinarayanan et al., 2017) es otra arquitectura encoder-decoder, pero con un enfoque distinto para el upsampling: usa los índices del max pooling del encoder para guiar el unpooling en el decoder.

La ventaja de usar los pooling indices es que es muy eficiente en memoria: no se necesitan almacenar los feature maps completos del encoder (como en U-Net con concatenación), sino solo los índices de las posiciones de los máximos.

🔍 Convoluciones dilatadas (Atrous)

Antes de ver DeepLab, necesitamos entender la convolución dilatada (o atrous convolution), que es su pieza fundamental.

Una convolución estándar 3×3 tiene un campo receptivo de 3×3. Para capturar contexto más amplio, normalmente apilamos muchas capas o usamos pooling. La convolución dilatada ofrece una alternativa: inserta «huecos» (dilation) entre los pesos del kernel.

Donde r es la tasa de dilatación (dilation rate). Con r=1 tenemos una convolución estándar. Con r=2, el kernel 3×3 tiene un campo receptivo efectivo de 5×5. Con r=4, el campo receptivo es de 9×9.

🏛️ DeepLab: segmentación de alta resolución

La familia DeepLab (Chen et al., Google) es una de las líneas de investigación más influyentes en segmentación semántica. Veamos su evolución:

DeepLab v1 y v2

DeepLab v1 (2015) introdujo el uso de convoluciones dilatadas en un backbone pre-entrenado (VGG-16), manteniendo una resolución de salida de ×8 en lugar de ×32. También incorporó un CRF (Conditional Random Field) como post-procesamiento para refinar bordes.

DeepLab v2 (2017) mejoró con dos innovaciones:

• Backbone más potente: cambió VGG por ResNet-101
• ASPP (Atrous Spatial Pyramid Pooling): aplica convoluciones dilatadas con múltiples tasas en paralelo para capturar contexto a distintas escalas

ASPP: capturando contexto multi-escala

ASPP es la innovación clave de DeepLab. Aplica varias convoluciones dilatadas en paralelo con diferentes tasas de dilatación (r = 6, 12, 18), capturando contexto a múltiples escalas simultáneamente:

DeepLab v3+: el estado del arte

DeepLab v3+ (2018) combinó ASPP con un decoder ligero, creando una arquitectura encoder-decoder donde:

• Encoder: backbone (ResNet/Xception) con convoluciones dilatadas + módulo ASPP
• Decoder: un módulo simple que combina features de baja resolución (del ASPP) con features de alta resolución (del backbone temprano) mediante concatenación y convoluciones

El resultado es un modelo que captura tanto contexto global (ASPP) como detalles finos (skip connection del encoder), logrando resultados excelentes en PASCAL VOC y Cityscapes.

🔮 PSPNet: Pyramid Scene Parsing

PSPNet (Zhao et al., 2017) aborda el problema del contexto global con un módulo de Pyramid Pooling. La motivación: cuando un modelo ve solo un trozo de un objeto grande (ej. un barco visto de cerca), puede confundirlo con algo pequeño (ej. un coche) si le falta contexto de la escena completa.

El Pyramid Pooling Module (PPM) aplica average pooling a 4 escalas distintas (1×1, 2×2, 3×3, 6×6), generando representaciones a diferentes niveles de granularidad que se concatenan con las features originales:

🎭 Mask R-CNN: segmentación por instancias

Todas las arquitecturas anteriores realizan segmentación semántica (no distinguen instancias). Mask R-CNN (He et al., 2017) extiende Faster R-CNN para realizar segmentación por instancias: detecta cada objeto y genera una máscara de segmentación para cada uno.

Arquitectura

Mask R-CNN añade una rama de máscara paralela a las ramas existentes de clasificación y regresión de bounding box de Faster R-CNN:

RoIAlign: alineamiento preciso

Una innovación clave de Mask R-CNN es RoIAlign, que reemplaza el RoI Pooling de Faster R-CNN. El RoI Pooling original cuantiza las coordenadas a enteros, causando desalineamientos de 1-2 píxeles — insignificante para bounding boxes pero catastrófico para máscaras a nivel de píxel.

RoIAlign usa interpolación bilineal para muestrear las features en coordenadas continuas, eliminando la cuantización.

🌐 Segmentación panóptica

La segmentación panóptica (Kirillov et al., 2019) unifica segmentación semántica e instancias en un solo marco. Para cada píxel, se predice:

1. Clase semántica (para todo: stuff y things)
2. ID de instancia (solo para things, es decir, objetos contables)

La métrica principal es PQ (Panoptic Quality), que combina reconocimiento y calidad de segmentación:

📉 Funciones de pérdida para segmentación

La elección de la función de pérdida es crítica en segmentación, especialmente cuando hay desbalance de clases (mucho más fondo que objeto). Veamos las más utilizadas.

Cross-Entropy píxel a píxel

La extensión natural de la cross-entropy de clasificación. Se calcula independientemente para cada píxel y se promedia:

Donde y_{c,h,w} es el ground truth one-hot y \hat{p}_{c,h,w} es la probabilidad predicha para la clase c en la posición (h,w).

Dice Loss

La Dice loss está basada en el coeficiente de Sørensen-Dice, una medida de solapamiento entre dos conjuntos. Mide directamente cuánto se «parecen» la predicción y el ground truth:

\epsilon (ej. 10^{-6}) evita divisiones por cero. La Dice loss es intrínsecamente robusta al desbalance porque normaliza por el tamaño del objeto.

Focal Loss

La Focal loss (Lin et al., 2017) modifica la cross-entropy para reducir la contribución de los píxeles fáciles (clasificados con alta confianza) y enfocarse en los difíciles:

Donde \gamma \geq 0 es el parámetro de enfoque. Con \gamma = 0 es cross-entropy estándar. Con \gamma = 2 (valor típico), un píxel clasificado con \hat{p}_t = 0.9 recibe un peso 100× menor que uno con \hat{p}_t = 0.1.

Tversky Loss

La Tversky loss generaliza la Dice loss con pesos independientes para falsos positivos (FP) y falsos negativos (FN):

Con \alpha = \beta = 0.5 se reduce a Dice loss. Aumentar \beta > \alpha penaliza más los falsos negativos, útil cuando «no detectar» un tumor es peor que una falsa alarma.

📏 Métricas de evaluación

La accuracy (proporción de píxeles correctos) no es útil en segmentación debido al desbalance de clases. Las métricas estándar son:

IoU (Intersection over Union)

También llamada Jaccard Index. Mide el solapamiento entre la predicción y el ground truth para cada clase:

mIoU (Mean IoU)

La métrica principal en segmentación semántica. Es la media de IoU sobre todas las clases:

Coeficiente Dice / F1

Relación con IoU: \text{Dice} = \frac{2 \cdot \text{IoU}}{1 + \text{IoU}}. El Dice siempre es ≥ IoU para el mismo par (predicción, ground truth).

🔧 Data augmentation para segmentación

El data augmentation es especialmente importante en segmentación porque los datasets anotados a nivel de píxel son pequeños. Pero hay una diferencia clave: las transformaciones geométricas deben aplicarse también a la máscara.

💻 Ejemplo práctico: entrenamiento de U-Net

Veamos un ejemplo completo de cómo entrenar un U-Net para segmentación binaria (ej. tumor vs. fondo) usando combinación de Cross-Entropy + Dice loss.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import numpy as np

# ── Loss combinada ──────────────────────────────────────────────
class DiceBCELoss(nn.Module):
    """Combina Binary Cross-Entropy y Dice Loss."""
    def __init__(self, smooth=1e-6):
        super().__init__()
        self.smooth = smooth
        self.bce = nn.BCEWithLogitsLoss()

    def forward(self, logits, targets):
        # BCE loss
        bce_loss = self.bce(logits, targets)

        # Dice loss (sobre probabilidades)
        probs = torch.sigmoid(logits)
        intersection = (probs * targets).sum(dim=(2, 3))
        union = probs.sum(dim=(2, 3)) + targets.sum(dim=(2, 3))
        dice = (2. * intersection + self.smooth) / (union + self.smooth)
        dice_loss = 1 - dice.mean()

        return bce_loss + dice_loss

# ── Dataset de segmentación ─────────────────────────────────────
class SegmentationDataset(Dataset):
    def __init__(self, images, masks, transform=None):
        self.images = images        # lista de arrays HxWx3
        self.masks = masks          # lista de arrays HxW (0 o 1)
        self.transform = transform

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image = self.images[idx]     # HxWx3, float32
        mask = self.masks[idx]       # HxW, float32

        if self.transform:
            # Aplicar MISMA transformación a imagen y máscara
            seed = np.random.randint(2147483647)
            torch.manual_seed(seed)
            image = self.transform(image)
            torch.manual_seed(seed)
            mask = self.transform(mask.unsqueeze(0))  # 1xHxW

        return image, mask

# ── Entrenamiento ───────────────────────────────────────────────
def train_unet(model, train_loader, val_loader, epochs=50, lr=1e-4):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = DiceBCELoss()
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, patience=5, factor=0.5
    )

    best_dice = 0.0

    for epoch in range(epochs):
        # ── Train ──
        model.train()
        train_loss = 0.0
        for images, masks in train_loader:
            images, masks = images.to(device), masks.to(device)

            optimizer.zero_grad()
            outputs = model(images)           # [B, 1, H, W]
            loss = criterion(outputs, masks)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

        # ── Validation ──
        model.eval()
        val_dice = 0.0
        with torch.no_grad():
            for images, masks in val_loader:
                images, masks = images.to(device), masks.to(device)
                outputs = torch.sigmoid(model(images))
                preds = (outputs > 0.5).float()

                # Calcular Dice
                intersection = (preds * masks).sum(dim=(2, 3))
                union = preds.sum(dim=(2, 3)) + masks.sum(dim=(2, 3))
                dice = (2. * intersection + 1e-6) / (union + 1e-6)
                val_dice += dice.mean().item()

        avg_train_loss = train_loss / len(train_loader)
        avg_val_dice = val_dice / len(val_loader)
        scheduler.step(avg_train_loss)

        print(f"Epoch {epoch+1}/{epochs} | "
              f"Loss: {avg_train_loss:.4f} | "
              f"Val Dice: {avg_val_dice:.4f}")

        if avg_val_dice > best_dice:
            best_dice = avg_val_dice
            torch.save(model.state_dict(), 'best_unet.pth')
            print(f"  → Nuevo mejor modelo (Dice: {best_dice:.4f})")

    print(f"\nMejor Dice en validación: {best_dice:.4f}")
    return model

import tensorflow as tf
from tensorflow.keras import backend as K

# ── Dice loss ───────────────────────────────────────────────────
def dice_loss(y_true, y_pred, smooth=1e-6):
    y_pred = tf.sigmoid(y_pred)
    intersection = tf.reduce_sum(y_true * y_pred, axis=[1, 2, 3])
    union = tf.reduce_sum(y_true, axis=[1, 2, 3]) + \
            tf.reduce_sum(y_pred, axis=[1, 2, 3])
    dice = (2. * intersection + smooth) / (union + smooth)
    return 1 - tf.reduce_mean(dice)

def bce_dice_loss(y_true, y_pred):
    bce = tf.keras.losses.binary_crossentropy(
        y_true, y_pred, from_logits=True
    )
    return tf.reduce_mean(bce) + dice_loss(y_true, y_pred)

# ── Dice metric ─────────────────────────────────────────────────
def dice_coefficient(y_true, y_pred, smooth=1e-6):
    y_pred = tf.cast(tf.sigmoid(y_pred) > 0.5, tf.float32)
    intersection = tf.reduce_sum(y_true * y_pred, axis=[1, 2, 3])
    union = tf.reduce_sum(y_true + y_pred, axis=[1, 2, 3])
    return tf.reduce_mean((2. * intersection + smooth) / (union + smooth))

# ── Compilar y entrenar ─────────────────────────────────────────
model = build_unet(input_shape=(256, 256, 3), num_classes=1)

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),
    loss=bce_dice_loss,
    metrics=[dice_coefficient]
)

callbacks = [
    tf.keras.callbacks.ReduceLROnPlateau(patience=5, factor=0.5),
    tf.keras.callbacks.EarlyStopping(patience=15, restore_best_weights=True),
    tf.keras.callbacks.ModelCheckpoint('best_unet.keras', save_best_only=True)
]

history = model.fit(
    train_dataset,
    validation_data=val_dataset,
    epochs=50,
    callbacks=callbacks
)

🏁 Comparativa de arquitecturas

Para cerrar, una comparativa completa de todas las arquitecturas de segmentación que hemos visto: