U-Net: Convolutional Networks for Biomedical Image Segmentation

Olaf Ronneberger; Philipp Fischer; Thomas Brox

Artículos · arXiv 1505.04597

U-Net: Redes convolucionales para la segmentación de imágenes biomédicas

U-Net: Convolutional Networks for Biomedical Image Segmentation

Olaf Ronneberger, Philipp Fischer, Thomas Brox

Publicado 2015-05-18cs.CV arXiv PDF

Leer con traducción e IA

Resumen en tres líneas

🎯 Entrenar redes de segmentación biomédica píxel a píxel de alta precisión de extremo a extremo utilizando conjuntos de datos anotados muy pequeños.
🔧 Una arquitectura convolucional simétrica en forma de U que combina una trayectoria de contracción para el contexto, una trayectoria expansiva con skip connections para una localización precisa, convoluciones válidas sin relleno, aumento por deformación elástica y un mapa de pesos para los bordes de separación.
📈 Superó a los métodos de ventana deslizante en el desafío de segmentación EM de ISBI (warping error de 0.000353) y ganó el ISBI Cell Tracking Challenge 2015 con un IOU de 0.9203 en PhC-U373 y un IOU de 0.7756 en DIC-HeLa.

Resumen

Existe un amplio consenso en que el entrenamiento exitoso de redes profundas requiere de muchos miles de muestras de entrenamiento anotadas. En este artículo, presentamos una red y una estrategia de entrenamiento que se apoya fuertemente en el uso del aumento de datos para aprovechar de forma más eficiente las muestras anotadas disponibles. La arquitectura consta de una trayectoria de contracción para capturar el contexto y una trayectoria de expansión simétrica que permite una localización precisa. Demostramos que una red de este tipo puede entrenarse de extremo a extremo a partir de muy pocas imágenes y supera al mejor método previo (una red convolucional con ventana deslizante) en el desafío ISBI para la segmentación de estructuras neuronales en pilas de microscopía electrónica. Utilizando la misma red entrenada sobre imágenes de microscopía óptica de luz transmitida (contraste de fases y DIC), ganamos el desafío ISBI cell tracking challenge 2015 en estas categorías por un amplio margen. Además, la red es rápida. La segmentación de una imagen de 512x512 tarda menos de un segundo en una GPU reciente. La implementación completa (basada en Caffe) y las redes entrenadas están disponibles en http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net .

Resumen original (inglés)

There is large consent that successful training of deep networks requires many thousand annotated training samples. In this paper, we present a network and training strategy that relies on the strong use of data augmentation to use the available annotated samples more efficiently. The architecture consists of a contracting path to capture context and a symmetric expanding path that enables precise localization. We show that such a network can be trained end-to-end from very few images and outperforms the prior best method (a sliding-window convolutional network) on the ISBI challenge for segmentation of neuronal structures in electron microscopic stacks. Using the same network trained on transmitted light microscopy images (phase contrast and DIC) we won the ISBI cell tracking challenge 2015 in these categories by a large margin. Moreover, the network is fast. Segmentation of a 512x512 image takes less than a second on a recent GPU. The full implementation (based on Caffe) and the trained networks are available at http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net .

Guía de lectura

Problema y motivación

La segmentación de imágenes biomédicas requiere una clasificación densa a nivel de píxel (localización) en lugar de una categorización a nivel de imagen. Dos desafíos fundamentales dominan este campo:

Los datos de entrenamiento anotados son escasos, consistiendo a menudo en solo unas pocas docenas de imágenes en lugar de miles o millones.
Los métodos previos de aprendizaje profundo, como las redes convolucionales basadas en ventanas deslizantes sobre parches (por ejemplo, Cireşan et al.), sufrían de una severa redundancia, una inferencia lenta y un dilema inherente entre el contexto del parche y la precisión de la localización.

Los autores buscaron una red totalmente convolucional de extremo a extremo que pudiera capturar el contexto global y los detalles locales finos de manera simultánea, procesar imágenes arbitrariamente grandes sin fisuras y entrenarse con éxito a partir de datos extremadamente limitados.

Ideas clave

Arquitectura simétrica en forma de U: Extiende las redes totalmente convolucionales al emparejar el codificador de contracción (submuestreo) con un decodificador expansivo (sobremuestreo) igualmente expresivo y equipado con numerosos canales de características para transferir el contexto hacia arriba.
Conexiones de salto (Skip Connections): Los mapas de características de alta resolución de la trayectoria de contracción se copian y concatenan con los mapas de características sobremuestreados en la trayectoria expansiva (tras un recorte), proporcionando detalles de localización espacial precisos que se pierden durante el pooling.
Estrategia Overlap-Tile con convoluciones válidas: La red utiliza convoluciones sin relleno (unpadded). Las regiones limítrofes se predicen fluidamente mediante mosaicos y extrapolando el contexto faltante por reflexión especular, evitando los límites de memoria de la GPU en imágenes grandes.
Aumento por deformaciones elásticas: Deformaciones elásticas suaves y realistas generadas mediante cuadrículas aleatorias gruesas enseñan a la red invarianza ante deformaciones, reduciendo drásticamente la cantidad requerida de muestras de entrenamiento.
Ponderación de pérdida en bordes de separación: Un mapa de pesos a nivel de píxel precalculado penaliza los errores en los bordes de fondo entre células en contacto de la misma clase, lo que permite una separación nítida entre instancias.

Método

Estructura de la red

La red cuenta con 23 capas convolucionales organizadas en dos trayectorias:

Trayectoria de contracción: Bloques repetidos de dos convoluciones $3 \times 3$ sin relleno (cada una seguida de una ReLU) y un max pooling de $2 \times 2$ con paso 2. Los canales de características se duplican en cada paso de pooling.
Trayectoria expansiva: Bloques repetidos de una operación de sobremuestreo seguida de una convolución de $2 \times 2$ ("up-convolution") que reduce a la mitad los canales de características, concatenación con el mapa de características recortado de la trayectoria de contracción con resolución coincidente, y dos convoluciones $3 \times 3$ seguidas de ReLUs.
Capa final: Una convolución $1 \times 1$ mapea las representaciones de características de 64 canales directamente a las puntuaciones de las $K$ clases deseadas.

Función de pérdida y ponderación de bordes

La red se entrena con entropía cruzada softmax a nivel de píxel: $p_{k} (x) = \frac{exp ( a _{k} ( x ))}{\sum _{k^{'} = 1}^{K} exp ( a _{k^{'}} ( x ))}$ $E = x \in Ω \sum w (x) lo g (p_{ℓ (x)} (x))$ donde $a_{k} (x)$ denota la activación para el canal de características $k$ en el píxel $x$ , $ℓ (x)$ es la etiqueta de clase de referencia (ground truth), y $w (x)$ es un mapa de pesos espacial para la pérdida.

Para balancear las frecuencias de clase y separar objetos en contacto, $w (x)$ se define como: $w (x) = w_{c} (x) + w_{0} \cdot exp (- \frac{( d _{1} ( x ) + d _{2} ( x ) ) ^{2}}{2 σ ^{2}})$ donde $w_{c} (x)$ equilibra las frecuencias de clase, $d_{1} (x)$ es la distancia al borde de la célula más cercana, $d_{2} (x)$ es la distancia al borde de la segunda célula más cercana, y los autores establecen $w_{0} = 10$ y $σ \approx 5$ píxeles.

Optimización

Se entrena usando SGD en Caffe con un tamaño de lote (batch size) de 1 (un único mosaico grande de imagen) y un momento alto (0.99).
Los pesos se inicializan a partir de una distribución gaussiana con desviación estándar $2/ N$ , donde $N$ es el número de nodos de entrada por neurona.

Resultados

Desafío de segmentación EM de ISBI (Estructuras neuronales):
- Obtuvo el mejor puesto con un warping error de $0.000353$ , un Rand error de $0.0382$ y un error de píxel de $0.0611$ (promediado sobre 7 versiones rotadas, sin postprocesamiento adicional).
- Superó al estado del arte previo de Cireşan et al. (warping error de $0.000420$ , Rand error de $0.0504$ ).
ISBI Cell Tracking Challenge 2015:
- PhC-U373 (Contraste de fases): Alcanzó una Intersección sobre Unión (IOU) promedio de $0.9203$ , en comparación con el $0.83$ del segundo mejor algoritmo.
- DIC-HeLa (Contraste de interferencia diferencial): Logró una IOU promedio de $0.7756$ , frente al $0.46$ del segundo mejor algoritmo.
Eficiencia: La segmentación de una imagen de $512 \times 512$ toma menos de un segundo en una GPU moderna, y el tiempo total de entrenamiento dura aproximadamente 10 horas en una GPU NVIDIA Titan (6 GB).

Limitaciones y preguntas abiertas

Convoluciones sin relleno y recorte: Debido a que las convoluciones son válidas (sin relleno), los mosaicos de salida son estrictamente más pequeños que los de entrada, lo que exige elegir cuidadosamente dimensiones de mosaico de entrada divisibles por $2$ a lo largo de todos los pasos de pooling y aplicar reflexión en los bordes de la imagen.
Separación de instancias mediante pesos heurísticos: La separación de instancias contiguas depende de mapas de peso morfológicos hechos a mano en lugar de representaciones embebidas de instancias o ramas de detección entrenadas de extremo a extremo.
Diseño específico para 2D: El artículo se enfoca en cortes 2D, dejando la adaptación a adquisiciones biomédicas volumétricas 3D para investigaciones posteriores.

Términos clave

contracting path: La mitad codificadora de U-Net que aplica repetidamente convoluciones 3x3 sin relleno y max pooling para capturar el contexto y extraer representaciones de características de alto nivel.
expansive path: La mitad decodificadora simétrica de U-Net que sobremuestrea los mapas de características, reduce las dimensiones de canal y reconstruye segmentaciones espaciales precisas.
up-convolution: Una convolución de 2x2 aplicada después del sobremuestreo que reduce a la mitad el número de canales de características antes de la concatenación.
overlap-tile strategy: Un enfoque de inferencia y entrenamiento que divide imágenes arbitrariamente grandes en mosaicos superpuestos, reflejando las regiones fronterizas para proporcionar el contexto necesario a las convoluciones sin relleno.
elastic deformations: Una técnica de aumento de datos que genera vectores de desplazamiento suaves en una cuadrícula gruesa mediante interpolación bicúbica para simular las variaciones de los tejidos biológicos.
weight map: Una función de ponderación de pérdida con variación espacial que combina el balanceo de frecuencias de clase y una penalización exponencial basada en las distancias a las dos células más cercanas para separar objetos en contacto.
warping error: Una métrica de evaluación de segmentación utilizada en el desafío ISBI EM que evalúa la corrección topológica, como la integridad de las membranas neuronales sin cortes.
fully convolutional network: Una red neuronal compuesta enteramente por operaciones convolucionales y de pooling sin capas totalmente conectadas, lo que permite predicciones densas sobre tamaños de entrada variables.

Preguntas

¿Por qué U-Net recorta los mapas de características antes de la concatenación en las skip connections?

Dado que U-Net utiliza convoluciones válidas (sin relleno), cada convolución reduce el tamaño del mapa de características en los píxeles del borde. Recortar el mapa de características de la trayectoria de contracción garantiza que sus dimensiones espaciales coincidan con las del mapa de características sobremuestreado para poder concatenarse.

¿Cómo gestiona U-Net los bordes de imágenes grandes durante la inferencia?

U-Net aplica una estrategia de mosaicos superpuestos (overlap-tile): para predecir la región interior válida de un mosaico, se suministran los píxeles contextuales circundantes como entrada, y cualquier dato faltante fuera de los límites de la imagen se extrapola mediante reflexión especular.

¿Cómo separa U-Net las células que están en contacto entre sí?

Utiliza un mapa de pesos precalculado en la función de pérdida que añade una penalización exponencial considerable a los píxeles de fondo situados entre los límites de la célula más cercana y la segunda más cercana, obligando a la red a conservar bordes finos de separación.

¿Por qué se establece el tamaño de lote (batch size) en 1 durante el entrenamiento?

Los autores prefirieron mosaicos de imagen grandes en lugar de tamaños de lote grandes para maximizar el contexto dentro de la memoria disponible en la GPU (NVIDIA Titan de 6 GB), compensando el tamaño de lote de un único mosaico mediante un momento SGD elevado de 0.99.

La explicación de esta página la escribió una IA a partir del artículo y puede contener errores. Consulte el artículo para los detalles.

Leer con traducción e IA