U-Net: Redes Convolucionais para Segmentação de Imagens Biomédicas

Guia de leitura

Problema e motivação

A segmentação de imagens biomédicas requer classificação densa em nível de pixel (localização) em vez de categorização no nível da imagem. Dois desafios fundamentais dominam esse domínio:

1. Os dados de treinamento anotados são escassos, consistindo frequentemente em apenas algumas dezenas de imagens em vez de milhares ou milhões.
2. Métodos anteriores de deep learning, como redes convolucionais de janela deslizante (sliding-window) baseadas em patches (por exemplo, Cireşan et al.), sofriam de redundância severa, inferência lenta e um compromisso inerente entre o contexto do patch e a precisão da localização.

Os autores buscaram uma fully convolutional network de ponta a ponta que pudesse capturar o contexto global e detalhes locais finos simultaneamente, processar imagens arbitrariamente grandes de forma contínua e treinar com sucesso a partir de dados extremamente limitados.

Ideias-chave

• Arquitetura Simétrica em Formato de U: Estende redes totalmente convolucionais ao combinar o encoder de contração (downsampling) com um decoder de expansão (upsampling) igualmente expressivo, equipado com inúmeros canais de features para propagar o contexto para as camadas superiores.
• Skip Connections: Mapas de features de alta resolução do caminho de contração são copiados e concatenados com mapas de features sobreamostrados no caminho de expansão (após corte/cropping), fornecendo detalhes espaciais precisos de localização perdidos durante o pooling.
• Estratégia Overlap-Tile com Convoluções Válidas: A rede utiliza convoluções sem preenchimento (unpadded). Regiões de borda são previstas perfeitamente dividindo a imagem em blocos (tiling) e extrapolando o contexto ausente por espelhamento, contornando os limites de memória da GPU em imagens grandes.
• Augmentation por Deformação Elástica: Deformações elásticas realistas e suaves, geradas por meio de grades aleatórias grosseiras, ensinam a rede invariância a deformações, reduzindo drasticamente o número necessário de amostras de treinamento.
• Ponderação da Perda em Bordas de Separação: Um mapa de pesos pixel a pixel pré-computado penaliza erros em bordas de fundo entre células da mesma classe que se tocam, permitindo a separação distinta de instâncias.

Método

Estrutura da Rede

A rede possui 23 camadas convolucionais organizadas em dois caminhos:

• Contracting Path (Caminho de Contração): Blocos repetidos de duas convoluções $3\times 3$ sem preenchimento (cada uma seguida por ReLU) e um max pooling $2\times 2$ com stride 2. O número de canais de features dobra a cada etapa de pooling.
• Expansive Path (Caminho de Expansão): Blocos repetidos de uma operação de upsampling seguida por uma convolução $2\times 2$ ("up-convolution") que reduz os canais de features pela metade, concatenação com o mapa de features recortado correspondente do caminho de contração e duas convoluções $3\times 3$ seguidas por ReLUs.
• Camada Final: Uma convolução $1\times 1$ mapeia as representações de features de 64 canais diretamente para os escores das $K$ classes desejadas.

Função de Perda e Ponderação de Bordas

A rede é treinada com cross-entropy softmax pixel a pixel: $$p_k(\mathbf{x})=\frac{\exp (a_k(\mathbf{x}))}{{\sum }_{k^{\prime }=1}^K\exp (a_{k^{\prime }}(\mathbf{x}))}$$ $$E=\sum _{\mathbf{x}\in \Omega }w(\mathbf{x})\log (p_{\ell (\mathbf{x})}(\mathbf{x}))$$ onde $a_k(\mathbf{x})$ denota a ativação para o canal de features $k$ no pixel $\mathbf{x}$, $\ell (\mathbf{x})$ é o rótulo de classe do ground truth e $w(\mathbf{x})$ é um mapa de pesos espacial da perda.

Para equilibrar as frequências de classes e separar objetos em contato, $w(\mathbf{x})$ é definido como: $$w(\mathbf{x})=w_c(\mathbf{x})+w_0\cdot \exp \left(-\frac{(d_1(\mathbf{x})+d_2(\mathbf{x}){)}^2}{2{\sigma }^2}\right)$$ onde $w_c(\mathbf{x})$ equilibra as frequências de classes, $d_1(\mathbf{x})$ é a distância até a borda da célula mais próxima, $d_2(\mathbf{x})$ é a distância até a borda da segunda célula mais próxima, e os autores definem $w_0=10$ e $\sigma \approx 5$ pixels.

Otimização

• Treinada utilizando SGD no Caffe com tamanho de lote (batch size) igual a 1 (um único tile grande de imagem) e momentum alto (0.99).
• Pesos inicializados a partir de uma distribuição gaussiana com desvio padrão $\sqrt{2/N}$, onde $N$ é o número de nós de entrada por neurônio.

Resultados

• ISBI EM Segmentation Challenge (Estruturas Neuronais):
  • Obteve a melhor classificação com um warping error de $0.000353$, Rand error de $0.0382$ e erro de pixel de $0.0611$ (média sobre 7 versões rotacionadas, sem pós-processamento adicional).
  • Superou o estado da arte anterior de Cireşan et al. (warping error de $0.000420$, Rand error de $0.0504$).
• ISBI Cell Tracking Challenge 2015:
  • PhC-U373 (Contraste de Fase): Alcançou uma Intersection over Union (IOU) média de $0.9203$, em comparação a $0.83$ do segundo colocado.
  • DIC-HeLa (Contraste de Interferência Diferencial): Alcançou uma IOU média de $0.7756$, em comparação a $0.46$ do segundo colocado.
• Eficiência: A segmentação de uma imagem de $512\times 512$ requer menos de um segundo em uma GPU moderna, com o tempo total de treinamento levando aproximadamente 10 horas em uma GPU NVIDIA Titan (6 GB).

Limitações e questões em aberto

• Convoluções sem Preenchimento e Recorte (Cropping): Como as convoluções são válidas (sem padding), os tiles de saída são estritamente menores que os de entrada, exigindo uma seleção cuidadosa das dimensões dos tiles de entrada divisíveis por $2$ em todas as etapas de pooling, além de espelhamento nas bordas da imagem.
• Separação de Instâncias via Pesos Heurísticos: A separação de instâncias em contato depende de mapas de pesos morfológicos manuais em vez de embeddings de instância de ponta a ponta ou ramos dedicados de detecção.
• Design Específico para 2D: O artigo foca em fatias 2D, deixando a adaptação para aquisições biomédicas volumétricas em 3D para pesquisas subsequentes.