Avances en la Segmentación de Imágenes de Microscopía con Transformers

Segmentar imágenes en microscopía es un paso clave para analizar estas imágenes. Se han desarrollado muchas técnicas a lo largo de los años, desde métodos más antiguos hasta modelos avanzados de deep learning. Un modelo popular en el campo de la Segmentación de imágenes biomédicas es U-Net. Recientemente, nuevos modelos llamados transformers han mostrado promesas para mejorar la forma en que segmentamos imágenes de microscopía. En este artículo, vamos a ver diferentes modelos basados en transformers, incluyendo UNETR, el Segment Anything Model, y Swin-UPerNet, y vamos a comparar su rendimiento con el modelo U-Net establecido.

#Importancia de la Segmentación

Identificar y separar diferentes objetos en imágenes de microscopía es esencial para un análisis completo. Segmentar con precisión varias estructuras dentro de estas imágenes, como los núcleos celulares, ayuda a los investigadores a reunir información importante sobre ellos. Sin embargo, lograr una segmentación efectiva puede ser complicado debido a la naturaleza complicada de los datos de microscopía.

Los algoritmos de deep learning son herramientas que pueden ayudar con las tareas de segmentación porque pueden aprender y reconocer características clave en las imágenes. Durante muchos años, el modelo U-Net, que está basado en Redes Neuronales Convolucionales (CNNs), ha sido ampliamente utilizado y ha dado buenos resultados en la segmentación de imágenes de microscopía. Sin embargo, han surgido muchos modelos nuevos, destacando los transformers como particularmente prometedores. Los transformers capturan características complejas en las imágenes y consideran mejor los contextos locales que las CNNs tradicionales, lo que puede ayudar a refinar los procesos de segmentación.

#Modelos Transformers en Segmentación de Imágenes

En esta discusión, vamos a evaluar modelos de segmentación populares que utilizan transformers de visión como codificadores de imágenes. Estos incluyen dos tipos principales: Vision Transformer (ViT) en el modelo UNETR, y el Swin Transformer en el modelo Swin-UPerNet. También veremos el Segment Anything Model (SAM), que utiliza entradas definidas por el usuario para mejorar los resultados de segmentación. A lo largo de esta comparación, utilizaremos el modelo U-Net como nuestro punto de referencia.

El Swin Transformer, cuando se empareja con el decodificador UPerNet, ha mostrado resultados prometedores en la segmentación de imágenes. Sin embargo, su método de procesamiento de parches de imagen a veces puede llevar a la pérdida de detalles importantes. Además, su dependencia de la interpolación bilineal puede afectar la precisión de la segmentación final. Para contrarrestar estos problemas, proponemos modificaciones a la arquitectura para mejorar el rendimiento y la captura de detalles.

Al evaluar estos modelos basados en transformers, nuestro objetivo es mostrar sus posibles ventajas y desafíos en comparación con el modelo U-Net tradicional.

#Trabajo Relacionado

Mientras que U-Net sigue siendo uno de los modelos más utilizados para la segmentación en aplicaciones biomédicas, en los últimos años hemos visto que han surgido numerosos modelos basados en transformers. Estos se pueden agrupar generalmente en dos categorías: modelos transformer-CNN y modelos híbridos. En los modelos transformer-CNN, los transformers actúan como el codificador de imagen principal mientras que las capas CNN en el decodificador generan máscaras de segmentación. Ejemplos incluyen UNETR, Swin UNETR y Swin-UPerNet. Los modelos híbridos incorporan tanto capas de transformers como de CNN pero mantienen las capas de CNN en el decodificador.

A pesar de la naturaleza flexible de los modelos híbridos, los modelos transformer-CNN suelen tener un mejor rendimiento. Esto se debe en gran parte al beneficio de los transformers preentrenados que han sido entrenados en grandes conjuntos de datos. Por lo tanto, los modelos híbridos a menudo son menos preferidos en la investigación.

Recientemente, los modelos fundamentales han ganado atención. Estos son entrenados en conjuntos de datos masivos y pueden generalizar sin necesidad de entrenamiento adicional para tareas específicas. Un ejemplo de esto es el Segment Anything Model (SAM), que utiliza entradas definidas por el usuario como cajas delimitadoras o puntos para la segmentación.

El Swin Transformer se ha vuelto popular para muchas tareas de imagen, llevando al desarrollo de nuevos modelos basados en su arquitectura, como Swin-UPerNet. Aunque la investigación sobre el Swin-UPerNet original es limitada, nuestro objetivo es encontrar oportunidades para mejorar su rendimiento mientras mantenemos su estructura fundamental.

#Datos y Metodología

Vamos a realizar una evaluación comparativa de U-Net contra modelos destacados basados en transformers, incluyendo UNETR, Swin-UPerNet y SAM. También vamos a incorporar modificaciones personalizadas a Swin-UPerNet para mejorar su rendimiento. Primero, vamos a describir los conjuntos de datos seleccionados para nuestra evaluación.

Para evaluar el rendimiento de los modelos, elegimos cuatro conjuntos de datos, que ofrecen desafíos únicos en diferentes modalidades de imagen. El conjunto de datos de Microscopia Electrónica consta de imágenes enfocadas en microscopía electrónica, mientras que el conjunto de datos de Siete Líneas Celulares incluye imágenes de campo brillante que apuntan a núcleos celulares. El conjunto de datos LIVECell presenta imágenes de contraste de fase concentrándose en células individuales, y el conjunto de datos MoNuSeg incluye imágenes histopatológicas de diapositivas completas. Esta variedad permite una evaluación completa de los modelos en múltiples escenarios.

#Modelos de Segmentación

U-Net sirve como nuestra línea base debido a su efectividad comprobada en tareas de segmentación. Su arquitectura cuenta con una estructura simétrica de codificador-decodificador con conexiones de salto para mejorar el mapeo de características.

Para los modelos transformers, seleccionamos UNETR, que utiliza ViT para la codificación, y Swin-UPerNet, que emplea el Swin Transformer con mecanismos de atención únicos. También incluimos el Segment Anything Model, que se basa en indicaciones del usuario para la segmentación.

Usamos el marco de modelos de segmentación de PyTorch para construir el modelo U-Net, utilizando una ResNet34 preentrenada como columna vertebral. Para UNETR, lo adaptamos para la segmentación de imágenes 2D, manteniendo el diseño arquitectónico central.

#Modificaciones a Swin-UPerNet

Al evaluar Swin-UPerNet, identificamos problemas que surgen del uso de un tamaño de parche de 4x, lo que reduce el tamaño de entrada y lleva a desalineaciones en el decodificador. El modelo original depende de la interpolación bilineal para alinear la máscara de segmentación final con la imagen de entrada, pero este método puede introducir artefactos y reducir el rendimiento general.

Para resolver estos desafíos, proponemos mejoras arquitectónicas. Estas incluyen:

1. Reducir el tamaño del parche para una captura más fina de detalles.
2. Agregar conexiones de salto desde la imagen de entrada al decodificador para mejorar el contexto local.
3. Introducir etapas adicionales en la columna vertebral de la red para mejorar la flexibilidad.

Estas modificaciones tienen como objetivo aumentar el rendimiento del modelo al segmentar imágenes de microscopía.

#Entrenamiento y Evaluación

Creamos un pipeline de entrenamiento para entrenar y evaluar diferentes modelos de manera eficiente. Las imágenes de entrada fueron preprocesadas usando técnicas estándar, incluyendo normalización y aplicación de aumentos para mejorar el reconocimiento de estructuras.

Cada modelo fue entrenado durante 150 épocas. Usamos un tamaño de lote de 16 y muestreamos imágenes del conjunto de datos para ejemplos de entrenamiento diversos. Las métricas de evaluación para valorar el rendimiento incluyeron las puntuaciones de F1 e IoU, evaluadas en conjuntos de prueba separados.

#Resultados

Comparamos el rendimiento de U-Net contra modelos transformers seleccionados, incluyendo UNETR, Swin-UPerNet y SAM. U-Net demostró consistentemente un gran rendimiento a través de los conjuntos de datos. UNETR tuvo un rendimiento comparable pero se quedó corto en el conjunto de datos de Microscopia Electrónica. El Swin-UPerNet original quedó atrás tanto de U-Net como de UNETR, subrayando la efectividad de los enfoques tradicionales de CNN.

Al analizar el rendimiento de las modificaciones de Swin-UPerNet, nuestros hallazgos mostraron mejoras significativas, especialmente con la variante Swin-S-TB-Skip. Esta modificación superó al Swin-UPerNet original y a U-Net en varios conjuntos de datos, demostrando los beneficios de nuestros cambios propuestos.

#Conclusión

Este estudio resalta los avances realizados en la segmentación de imágenes de microscopía a través de modelos basados en transformers. Si bien U-Net sigue siendo altamente efectivo, los modelos más nuevos muestran un gran potencial, especialmente con modificaciones cuidadosas. Nuestros hallazgos también revelan áreas para el desarrollo futuro, enfatizando la necesidad de refinar estas arquitecturas para su uso práctico en diversas herramientas de análisis de microscopía.

Al explorar y mejorar estos modelos, podemos aumentar aún más la precisión y eficiencia de las tareas de segmentación en microscopía, haciéndolas más valiosas para la investigación científica y aplicaciones prácticas.