Self-supervised Pretraining of Cell Segmentation Models

Título: DINOCell: Pré-treinamento Auto-supervisionado para Segmentação de Células em Microscopia

1. O Problema

A segmentação de instâncias em imagens de microscopia é fundamental para a quantificação de morfologia celular, permitindo monitoramento de saúde celular, detecção de senescência e análise de funções biológicas. No entanto, o avanço nessa área enfrenta dois obstáculos principais:

• Escassez de Dados Rotulados: A criação de máscaras de segmentação manuais é demorada, cara e sujeita a inconsistências entre anotadores.
• Desafio de Domínio (Domain Shift): Modelos modernos de "foundation models" (como o Segment Anything Model - SAM) são pré-treinados em grandes conjuntos de dados de imagens naturais. As representações aprendidas nesses dados (baseadas em texturas e "objetos" distintos) não se alinham bem com imagens de microscopia, onde as bordas são menos distintas, as formas são variadas e as células podem estar densamente agrupadas. Isso resulta em desempenho degradado quando esses modelos são aplicados a dados biológicos sem adaptação adequada.

2. Metodologia: DINOCell

Os autores propõem o DINOCell, um framework que utiliza aprendizado auto-supervisionado para adaptar representações visuais ao domínio da microscopia antes do ajuste fino supervisionado. A abordagem segue um processo de três estágios:

• Estágio 1: Pré-treinamento Auto-supervisionado (DINOv2):
  O modelo inicia com pesos pré-treinados do DINOv2 (um Vision Transformer - ViT-B), que aprende representações visuais robustas através de auto-distilação (framework professor-aluno) em grandes conjuntos de dados gerais.

• Estágio 2: Adaptação de Domínio (Domain Adaptation):
  Para reduzir a lacuna entre imagens naturais e microscopia, o encoder do DINOv2 é submetido a treinamento contínuo (auto-supervisionado) em um conjunto de dados não rotulado específico de microscopia.

  • Dados: 130.000 imagens não rotuladas de células, coletadas de múltiplas fontes (Broad Institute, BBBC, IDR), cobrindo mais de 50 tipos celulares e quatro modalidades (contraste de fase, campo claro, DIC e fluorescência).
  • Objetivo: Alinhar as representações aprendidas com os padrões estruturais das células, preservando os priores visuais gerais.

• Estágio 3: Ajuste Fino Supervisionado (Fine-tuning):
  O modelo completo é ajustado para a tarefa de segmentação usando dados rotulados.

  • Arquitetura: Um encoder ViT-B (adaptado) seguido por um decodificador leve com convoluções transpostas.
  • Objetivo de Treinamento: Inspirado no Cellpose, o modelo não prevê máscaras diretamente, mas sim campos de fluxo vetoriais (flows) e probabilidade celular. Os fluxos apontam para o centro de cada célula, permitindo a reconstrução das instâncias via rastreamento de gradiente.
  • Técnica de "Ignore Masking": Para lidar com dados rotulados imperfeitos (com células não anotadas), pixels sem rótulo são ponderados com um peso reduzido (0.05) no cálculo da perda, evitando que o modelo aprenda falsos negativos ou se sobreajuste a erros de anotação.

3. Contribuições Principais

1. Novo Framework de Pré-treinamento: Demonstração de que a adaptação de domínio auto-supervisionada (usando DINOv2) é mais eficaz para microscopia do que o uso direto de priores de tarefas aprendidos em imagens naturais (como no SAM).
2. DINOCell: Uma arquitetura que integra um encoder de visão adaptado ao domínio com um decodificador de fluxo, superando a necessidade de anotações massivas para pré-treinamento inicial.
3. Métrica MMA (Maximum Matching Accuracy): Introdução de uma nova métrica de avaliação que formula um problema de correspondência bipartida máxima para calcular a área de interseção total entre máscaras reais e previstas, normalizada pela união total.
4. Análise de Representação: Visualização via PCA (Análise de Componentes Principais) mostrando que o pré-treinamento e a adaptação de domínio organizam progressivamente os embeddings do encoder para corresponder à morfologia celular, reduzindo ruído e melhorando a separação de objetos.

4. Resultados

O modelo foi avaliado no benchmark LIVECell e em conjuntos de dados "out-of-distribution" (zero-shot).

• Desempenho no LIVECell (In-Distribution):

  • O DINOCell alcançou uma pontuação SEG de 0.784.
  • Isso representa uma melhoria de 10,42% em relação aos modelos baseados em SAM (SAMCell e Cellpose-SAM).
  • O modelo superou modelos especialistas (treinados apenas em um dataset) e generalistas, demonstrando capacidade de ajuste fino superior.

• Desempenho Zero-Shot (Generalização):

  • O DINOCell demonstrou robustez excepcional em três conjuntos de dados não vistos durante o treinamento (PBL-HEK, PBL-N2A e Glioma C6).
  • No conjunto Glioma C6 (altamente variável em morfologia e condições de iluminação), o DINOCell superou o melhor baseline (SAMCell-cyto) em 43,84% na métrica MMA. Enquanto outros modelos falharam em segmentar a maioria das células, o DINOCell manteve alta precisão.
  • Melhorias de 29,51% a 285,49% foram observadas em métricas de zero-shot em comparação com modelos sem pré-treinamento.

• Impacto do Pré-treinamento vs. Adaptação de Domínio:

  • O pré-treinamento com DINOv2 foi o fator mais crítico para o desempenho.
  • A adaptação de domínio adicional trouxe melhorias marginais no conjunto de teste, mas foi crucial para a generalização em cenários zero-shot (especialmente no Glioma C6).

5. Significado e Conclusão

O trabalho valida a hipótese de que, devido à grande diferença entre imagens do mundo real e microscopia, aprender representações visuais adequadas ao domínio é um problema mais limitante do que aprender o mapeamento de segmentação em si.

• Eficiência de Dados: O DINOCell demonstra que é possível obter desempenho de ponta em segmentação celular utilizando grandes quantidades de dados não rotulados para pré-treinamento, reduzindo a dependência de anotações manuais caras.
• Robustez: A abordagem supera significativamente os modelos baseados em foundation models de imagens naturais (como SAM) em cenários biológicos complexos, onde bordas são difusas e aglomerações celulares são comuns.
• Futuro: O estudo sugere que a adaptação de domínio auto-supervisionada deve ser um padrão para aplicações de visão computacional em ciências da vida, onde a distribuição de dados difere substancialmente dos conjuntos de dados gerais.

Em resumo, o DINOCell estabelece um novo estado da arte para segmentação de células, provando que a combinação de pré-treinamento auto-supervisionado em larga escala com adaptação de domínio específica é a chave para modelos robustos e generalizáveis em microscopia.