Avanços na Segmentação de Imagens em Microscopia com Transformers

Segmentar imagens em microscopia é uma parte crucial na análise dessas imagens. Muitas técnicas foram desenvolvidas ao longo dos anos, desde métodos mais antigos até modelos avançados de deep learning. Um modelo popular na área de Segmentação de imagens biomédicas é o U-Net. Recentemente, novos modelos chamados transformers mostraram promessas para melhorar a forma como segmentamos imagens de microscopia. Neste artigo, vamos dar uma olhada em diferentes modelos baseados em transformers, incluindo o UNETR, o Segment Anything Model e o Swin-UPerNet, e comparar o desempenho deles com o modelo U-Net já estabelecido.

#Importância da Segmentação

Identificar e separar diferentes objetos em imagens de microscopia é essencial para uma análise completa. Segmentar com precisão várias estruturas dentro dessas imagens, como núcleos celulares, ajuda os pesquisadores a coletar informações importantes sobre elas. No entanto, conseguir uma segmentação eficaz pode ser difícil por causa da complexidade dos dados de microscopia.

Os algoritmos de deep learning são ferramentas que podem ajudar nas tarefas de segmentação porque conseguem aprender e reconhecer características-chave nas imagens. Durante muitos anos, o modelo U-Net, que é baseado em Redes Neurais Convolucionais (CNNs), tem sido amplamente usado e tem mostrado bons resultados na segmentação de imagens de microscopia. Porém, muitos novos modelos surgiram, e os transformers se destacam como particularmente promissores. Os transformers capturam características complexas em imagens e consideram melhor os contextos locais do que as CNNs tradicionais, o que pode ajudar a refinar os processos de segmentação.

#Modelos Transformer na Segmentação de Imagens

Neste debate, vamos avaliar modelos populares de segmentação que usam transformers de visão como codificadores de imagem. Estes incluem dois tipos principais: Vision Transformer (ViT) no modelo UNETR, e o Swin Transformer no modelo Swin-UPerNet. Também vamos olhar para o Segment Anything Model (SAM), que usa entradas definidas pelo usuário para melhorar os resultados de segmentação. Ao longo dessa comparação, usaremos o modelo U-Net como nosso ponto de referência.

O Swin Transformer, quando combinado com o decodificador UPerNet, tem mostrado resultados promissores na segmentação de imagens. No entanto, o método dele de processar patches de imagem pode, às vezes, levar à perda de detalhes importantes. Além disso, a dependência dele da interpolação bilinear pode afetar a precisão da segmentação final. Para contornar esses problemas, propomos modificações na arquitetura para melhorar o desempenho e a captura de detalhes.

Ao avaliar esses modelos baseados em transformers, nosso objetivo é mostrar as vantagens e desafios potenciais deles em comparação com o modelo tradicional U-Net.

#Trabalhos Relacionados

Embora o U-Net continue sendo um dos modelos mais utilizados para segmentação em aplicações biomédicas, nos últimos anos vários modelos baseados em transformers surgiram. Esses modelos podem ser agrupados em duas categorias: modelos transformer-CNN e modelos híbridos. Nos modelos transformer-CNN, os transformers atuam como o codificador principal da imagem enquanto as camadas CNN no decodificador geram máscaras de segmentação. Exemplos incluem UNETR, Swin UNETR, e Swin-UPerNet. Modelos híbridos incorporam tanto camadas de transformer quanto de CNN, mas mantêm as camadas de CNN no decodificador.

Apesar da natureza flexível dos modelos híbridos, os modelos transformer-CNN costumam apresentar um desempenho melhor. Isso se deve, em grande parte, ao benefício de transformers pré-treinados que foram treinados em grandes conjuntos de dados. Como resultado, os modelos híbridos costumam ser menos preferidos na pesquisa.

Recentemente, modelos fundamentais ganharam atenção. Esses modelos são treinados em conjuntos de dados massivos e conseguem generalizar sem precisar de treinamento adicional para tarefas específicas. Um exemplo é o Segment Anything Model (SAM), que utiliza entradas definidas pelo usuário, como caixas delimitadoras ou pontos para segmentação.

O Swin Transformer se tornou popular para várias tarefas de imagem, levando ao desenvolvimento de novos modelos baseados em sua arquitetura, como o Swin-UPerNet. Embora a pesquisa sobre o Swin-UPerNet original seja limitada, nosso objetivo é encontrar oportunidades para melhorar seu desempenho enquanto mantemos sua estrutura fundamental.

#Dados e Metodologia

Vamos realizar uma avaliação comparativa do U-Net contra modelos notáveis baseados em transformers, incluindo UNETR, Swin-UPerNet e SAM. Também vamos incorporar modificações personalizadas no Swin-UPerNet para aprimorar seu desempenho. Primeiro, vamos descrever os conjuntos de dados selecionados para nossa avaliação.

Para avaliar o desempenho dos modelos, escolhemos quatro conjuntos de dados, que oferecem desafios únicos em diferentes modalidades de imagem. O conjunto de dados de Microscopia Eletrônica consiste em imagens focadas em microscopia eletrônica, enquanto o conjunto de dados de Sete Linhagens Celulares inclui imagens de campo brilhante que focam em núcleos celulares. O conjunto de dados LIVECell apresenta imagens de contraste de fase que se concentram em células individuais, e o conjunto de dados MoNuSeg inclui imagens de histopatologia em lâmina inteira. Essa variedade permite uma avaliação abrangente dos modelos em múltiplos cenários.

#Modelos de Segmentação

O U-Net serve como nossa linha de base devido à sua eficácia comprovada em tarefas de segmentação. Sua arquitetura apresenta uma estrutura simétrica de codificador-decodificador com conexões de salto para melhorar o mapeamento de características.

Para os modelos de transformer, selecionamos o UNETR, que usa ViT para codificação, e o Swin-UPerNet, que emprega o Swin Transformer com mecanismos de atenção exclusivos. Também incluímos o Segment Anything Model, que depende de prompts do usuário para segmentação.

Usamos o framework Segmentation Models PyTorch para construir o modelo U-Net, utilizando um ResNet34 pré-treinado como espinha dorsal. Para o UNETR, adaptamos ele para segmentação de imagem 2D, mantendo o design arquitetônico principal.

#Modificações no Swin-UPerNet

Ao avaliar o Swin-UPerNet, identificamos problemas decorrentes do uso de um tamanho de patch de 4x, que reduz o tamanho de entrada e leva a desalinhamentos no decodificador. O modelo original depende da interpolação bilinear para alinhar a máscara de segmentação final com a imagem de entrada, mas esse método pode introduzir artefatos e reduzir o desempenho geral.

Para resolver esses desafios, propomos melhorias arquitetônicas. Essas incluem:

1. Reduzir o tamanho do patch para capturar detalhes mais finos.
2. Adicionar conexões de salto da imagem de entrada para o decodificador para melhorar o contexto local.
3. Introduzir estágios adicionais na espinha dorsal da rede para melhorar a flexibilidade.

Essas modificações visam aumentar o desempenho do modelo ao segmentar imagens de microscopia.

#Treinamento e Avaliação

Criamos um pipeline de treinamento para treinar e avaliar diferentes modelos de forma eficiente. As imagens de entrada foram pré-processadas usando técnicas padrão, incluindo normalização e aplicação de aumentações para melhor reconhecimento de estruturas.

Cada modelo passou por treinamento por 150 épocas. Usamos um tamanho de lote de 16 e amostramos imagens do conjunto de dados para exemplos de treinamento diversos. As métricas de avaliação para analisar o desempenho incluíram scores F1 e IoU, avaliados em conjuntos de teste separados.

#Resultados

Comparamos o desempenho do U-Net contra modelos transformer selecionados, incluindo UNETR, Swin-UPerNet e SAM. O U-Net demonstrou um desempenho forte consistente em vários conjuntos de dados. O UNETR apresentou um desempenho comparável, mas ficou atrás no conjunto de dados de Microscopia Eletrônica. O Swin-UPerNet original ficou atrás tanto do U-Net quanto do UNETR, destacando a eficácia das abordagens tradicionais de CNN.

Ao analisar o desempenho das modificações no Swin-UPerNet, nossos achados mostraram melhorias significativas, especialmente com a variante Swin-S-TB-Skip. Essa modificação superou o Swin-UPerNet original e o U-Net em vários conjuntos de dados, mostrando os benefícios das mudanças propostas.

#Conclusão

Este estudo destaca os avanços feitos na segmentação de imagens de microscopia através de modelos baseados em transformers. Enquanto o U-Net continua sendo altamente eficaz, modelos mais novos mostram grande potencial, especialmente com modificações pensadas. Nossos achados também revelam áreas para desenvolvimento futuro, enfatizando a necessidade de refinar essas arquiteturas para uso prático em diversas ferramentas de análise de microscopia.

Ao explorar e aprimorar esses modelos, podemos melhorar ainda mais a precisão e a eficiência das tarefas de segmentação em microscopia, tornando-as mais valiosas para a pesquisa científica e aplicações práticas.