U-Net

U-Net es una red neuronal convolucional (CNN) para la segmentación semántica, tarea que consiste en asignar una clase a cada píxel de una imagen. Olaf Ronneberger, Philipp Fischer y Thomas Brox, de la Universidad de Friburgo, la presentaron en 2015 en U-Net: Convolutional Networks for Biomedical Image Segmentation. Un codificador reduce la imagen para captar su contenido; un decodificador amplía el resultado hasta recuperar la resolución completa, y las conexiones de salto transfieren los detalles del codificador al decodificador en cada escala, lo que permite obtener límites de objetos nítidos. U-Net se entrenó bien con unas pocas decenas de imágenes de microscopía anotadas, ganó por amplio margen en dos categorías del reto de seguimiento celular ISBI 2015 y sigue siendo un modelo de referencia en el análisis de imágenes médicas. Más adelante, su diseño se convirtió en la red de eliminación de ruido de muchos modelos de difusión.

Cómo funciona U-Net#

La red tiene dos mitades simétricas, representadas como los dos brazos de una U:

• Ruta de contracción (codificador): Pares repetidos de convoluciones 3×3 con activaciones ReLU, cada par seguido de una agrupación máxima 2×2 que reduce a la mitad la resolución. El número de canales se duplica en cada nivel, de 64 a 1024, por lo que la red intercambia detalle espacial por un campo receptivo más amplio y características más ricas.
• Ruta de expansión (decodificador): En cada nivel, una convolución ascendente 2×2 duplica la resolución y reduce a la mitad los canales; el resultado se concatena con el mapa de características del codificador del mismo nivel y dos convoluciones 3×3 los combinan. Una convolución final 1×1 convierte las características de cada píxel en puntuaciones de clase.

La red original tiene 23 capas convolucionales y cuatro pasos de agrupación. Utilizaba convoluciones sin relleno, por lo que una imagen de entrada de 572×572 producía un mapa de segmentación de 388×388, y las imágenes grandes se segmentaban en teselas solapadas. Actualmente, la mayoría de las implementaciones añaden relleno a las convoluciones para que la salida tenga el mismo tamaño que la entrada.

Conexiones de salto#

La agrupación descarta la ubicación exacta de cada característica. En el cuello de botella, la red sabe que una región contiene una membrana celular, pero no por qué píxeles pasa. Las conexiones de salto copian el mapa de características de alta resolución de cada nivel del codificador a través de la U hasta el decodificador, que combina el «qué» de la ruta ampliada con el «dónde» del codificador. U-Net concatena ambos a lo largo de la dimensión de canales, mientras que las conexiones residuales de ResNet suman características dentro de un bloque. Este código de PyTorch crea un nivel de codificador, un cuello de botella y un nivel de decodificador:

import torch
from torch import nn

# One encoder level, a bottleneck, and one decoder level of a U-Net
encoder = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.ReLU())
bottleneck = nn.Sequential(nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding=1), nn.ReLU())
upsample = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
decoder = nn.Sequential(nn.Conv2d(32, 16, 3, padding=1), nn.ReLU())
head = nn.Conv2d(16, 2, kernel_size=1)  # 2 classes, e.g. cell and background

image = torch.randn(1, 3, 128, 128)
skip = encoder(image)  # full-resolution features kept for the decoder
context = bottleneck(skip)  # half resolution, twice the channels, wider context

# Skip connection: concatenate upsampled context with the saved high-resolution features
merged = torch.cat([upsample(context), skip], dim=1)
logits = head(decoder(merged))
print(f"context {tuple(context.shape)} -> merged {tuple(merged.shape)} -> logits {tuple(logits.shape)}")

Muestra context (1, 32, 64, 64) -> merged (1, 32, 128, 128) -> logits (1, 2, 128, 128): los 16 canales de contexto ampliados y los 16 canales de salto se apilan en 32, y la salida contiene dos puntuaciones de clase para cada uno de los píxeles de entrada, de 128×128.

Entrenamiento con pocos datos#

El conjunto de entrenamiento original para estructuras neuronales en microscopía electrónica contenía 30 imágenes de 512×512 píxeles. U-Net lo compensó con una intensa aumentación de datos, especialmente con deformaciones elásticas aleatorias que imitan la flexión del tejido, y con una pérdida de entropía cruzada ponderada que da importancia a los estrechos espacios de fondo entre células en contacto, para que las células vecinas aparezcan como regiones separadas. Actualmente, los modelos de segmentación médica suelen entrenarse con la pérdida Dice, que mide directamente el solapamiento de máscaras y funciona bien cuando una estructura pequeña ocupa solo un pequeño porcentaje de los píxeles.

Variantes y modelos derivados#

• 3D U-Net y V-Net: Sustituyen las operaciones 2D por operaciones 3D para segmentar volúmenes, como las tomografías computarizadas y las resonancias magnéticas.
• Attention U-Net y UNet++: Añaden compuertas de atención o rutas de salto anidadas y densamente conectadas para refinar qué características del codificador llegan al decodificador.
• nnU-Net: Un marco de trabajo que configura automáticamente el preprocesamiento, la arquitectura y el entrenamiento para un nuevo conjunto de datos. Un estudio de MICCAI 2024 concluyó que las U-Net basadas en CNN bien ajustadas de este marco siguen superando a las arquitecturas Transformer y Mamba más recientes en la segmentación médica 3D cuando se validan rigurosamente.
• U-Net de difusión: DDPM y Stable Diffusion 1.x utilizan una U-Net para predecir el ruido de una imagen o de una representación latente, con capas de atención añadidas para que la indicación guíe la generación. La U-Net de Stable Diffusion 1.x tiene unos 860 millones de parámetros, y ControlNet se acopla al codificador de esta U-Net. Los generadores más recientes sustituyen U-Net por Transformer de difusión.

U-Net frente a modelos relacionados#

• Redes totalmente convolucionales (FCN): La FCN de 2014 convirtió clasificadores en predictores por píxel y fusionó predicciones gruesas y finas sumando mapas de puntuaciones de clase. U-Net creó un decodificador completo que concatena características en cada nivel.
• Autoencoders: También tienen la estructura de un codificador y un decodificador, pero se entrenan para reconstruir la entrada a través de un cuello de botella estrecho. Las conexiones de salto permitirían a un autoencoder evitar ese cuello de botella, mientras que U-Net las necesita para producir etiquetas alineadas por píxel.
• Modelos de segmentación de instancias: U-Net genera una clase por píxel, por lo que dos objetos en contacto de una misma clase comparten una etiqueta, salvo que el entrenamiento o el posprocesamiento los separen.

Los modelos de segmentación semántica de Ultralytics YOLO26, identificados por el sufijo -sem, utilizan una versión reducida del decodificador U-Net. Su cuello amplía dos veces la salida a escala 1/32 de la red troncal, concatenando por el camino características de la red troncal a escalas 1/16 y 1/8. Después clasifica los píxeles a escala 1/8 y amplía las puntuaciones de clase 8 veces hasta alcanzar el tamaño de entrada.

Aplicaciones en el mundo real#

• Análisis de imágenes por satélite: Los mapas de cobertura del suelo etiquetan cada píxel como agua, bosque, cultivo o edificios para hacer un seguimiento de la deforestación y el crecimiento urbano.
• Inspección industrial: Las máscaras de píxeles de arañazos, grietas y corrosión miden el tamaño y la forma exactos de un defecto, algo que una caja delimitadora no puede hacer.
• Diagnóstico por imagen: Las U-Net delimitan órganos, tumores y lesiones en tomografías computarizadas y resonancias magnéticas para el diagnóstico y la planificación de radioterapia, y separan células individuales en imágenes de microscopía. En el reto de seguimiento celular ISBI 2015, U-Net, entrenada con 20 imágenes parcialmente anotadas de células HeLa, alcanzó un IoU medio del 77,5 %, frente al 46 % del segundo mejor método.