U-Net é uma arquitetura de rede neural convolucional desenvolvida para segmentação de imagens, particularmente em contextos biomédicos. Foi introduzida em 2015 por Olaf Ronneberger, Philipp Fischer e Thomas Brox na Universidade de Freiburg, com base em redes totalmente convolucionais anteriores. A arquitetura se distingue por seu design simétrico de codificador-decodificador com conexões de salto, o que permite produzir mapas de segmentação precisos a partir de relativamente poucas imagens de treinamento. A segmentação de uma imagem de 512 × 512 leva menos de um segundo em uma GPU moderna (2015) usando a arquitetura U-Net.
A arquitetura U-Net deriva da chamada "rede totalmente convolucional" (FCN), que foi desenvolvida por Evan Shelhamer, Jonathan Long e Trevor Darrell em 2014. A ideia principal é complementar uma rede contrátil usual com camadas sucessivas, onde operações de pooling são substituídas por operadores de upsampling. Portanto, essas camadas aumentam a resolução da saída. Uma camada convolucional sucessiva pode então aprender a montar uma saída precisa com base nessa informação.
Arquitetura da Rede
A rede consiste em um caminho contrátil e um caminho expansivo, o que lhe confere a arquitetura em forma de U. O caminho contrátil é uma rede convolucional típica que consiste na aplicação repetida de convoluções, cada uma seguida por uma unidade linear retificada (ReLU) e uma operação de max pooling. Durante a contração, a informação espacial é reduzida enquanto a informação de características é aumentada. O caminho expansivo combina a informação de características e espacial através de uma sequência de up-convoluções e concatenações com características de alta resolução do caminho contrátil.
Uma modificação importante no U-Net é que há um grande número de canais de características na parte de upsampling, o que permite que a rede propague informação de contexto para camadas de maior resolução. Como consequência, o caminho expansivo é mais ou menos simétrico à parte contrátil, resultando em uma arquitetura em forma de U. A rede usa apenas a parte válida de cada convolução, sem nenhuma camada totalmente conectada. Portanto, a rede pode ser treinada e executada com uma mistura de resoluções de entrada.
Para prever os pixels na região de borda da imagem, o contexto ausente além do limite da imagem é extrapolado espelhando a imagem de entrada nas bordas da imagem. O artigo original do U-Net sugeriu ainda uma estratégia de mosaico, onde imagens grandes são cortadas em mosaicos sobrepostos que são processados independentemente. Isso permite processar imagens de alta resolução que, de outra forma, excederiam a memória GPU disponível. Recentemente, também houve interesse em modelos U-Net baseados em campo receptivo para segmentação de imagens médicas.
Treinamento e Função de Perda
O U-Net é treinado de ponta a ponta usando gradiente descendente estocástico (SGD) com uma função de perda pixel a pixel, tipicamente entropia cruzada ou uma variante ponderada. A função de energia é calculada por um soft-max pixel a pixel sobre o mapa de características final combinado com a perda de entropia cruzada. Os autores introduziram um mapa de pesos que dá maior importância a pixels próximos às fronteiras entre objetos tocantes, o que ajuda a rede a aprender a separar instâncias individuais em imagens biomédicas. Este mapa de pesos é pré-computado para cada imagem de treinamento e é usado para penalizar erros nas regiões de borda mais severamente.
A rede usa aumento de dados para aumentar o tamanho efetivo do conjunto de treinamento, incluindo deformações elásticas, rotações e variações de intensidade. Isso é particularmente importante em imagens biomédicas, onde dados anotados são escassos. A implementação original usou um momentum alto (0,99) e uma taxa de aprendizado baixa (0,00001) para garantir convergência estável.
Aplicações em Imagens Biomédicas
Existem muitas aplicações do U-Net na segmentação de imagens biomédicas, como a segmentação de diferentes órgãos e sistemas corporais em tomografias computadorizadas (TC) e ressonâncias magnéticas (RM). Casos específicos incluem segmentação de imagens cerebrais (BRATS) e segmentação de imagens do fígado (siliver07), bem como predição de sítios de ligação de proteínas. Implementações do U-Net também encontraram uso nas ciências físicas, por exemplo, na análise de micrografias de materiais.
Variações do U-Net também foram aplicadas para reconstrução de imagens médicas. Aqui estão algumas variantes e aplicações do U-Net, como segue:
- Regressão pixel a pixel usando U-Net e sua aplicação em pansharpening.
- U-Net 3D: Aprendizado de Segmentação Volumétrica Densa a partir de Anotação Esparsa.
- TernausNet: U-Net com Codificador VGG11 Pré-Treinado no ImageNet para Segmentação de Imagens.
- Tradução de imagem para imagem para estimar corantes fluorescentes.
- Na predição de sítios de ligação de estrutura de proteínas.
Uso em Modelos de Difusão
A arquitetura U-Net também foi empregada em modelos de difusão para remoção iterativa de ruído em imagens. Essa tecnologia fundamenta muitos modelos modernos de geração de imagens, como DALL-E, Midjourney e Stable Diffusion. Nesses modelos, um U-Net é treinado para prever o ruído adicionado a uma imagem em cada etapa do processo de difusão, permitindo que o modelo refine gradualmente uma entrada de ruído aleatório em uma imagem coerente. As conexões de salto no U-Net são particularmente benéficas aqui porque preservam detalhes de alta frequência durante o processo de remoção de ruído.
Exploração para Modelos de Linguagem
O U-Net também está sendo explorado para modelos de linguagem. A tokenização não é uma etapa separada, permitindo que o modelo entenda mais facilmente a ortografia e, simultaneamente, vetorize / tokenize conceitos de nível superior. Esta é uma área emergente de pesquisa que aproveita a capacidade do U-Net de capturar características multiescala, potencialmente oferecendo uma alternativa à arquitetura transformador dominante no processamento de linguagem natural.
História e Legado
O U-Net foi criado por Olaf Ronneberger, Philipp Fischer e Thomas Brox em 2015 e relatado no artigo "U-Net: Convolutional Networks for Biomedical Image Segmentation". É uma melhoria e desenvolvimento da FCN: Evan Shelhamer, Jonathan Long, Trevor Darrell (2014). "Fully convolutional networks for semantic segmentation". O artigo se tornou um dos mais citados no campo do aprendizado profundo, e a arquitetura foi adaptada para inúmeras tarefas além da imagem biomédica.
O sucesso do U-Net inspirou muitas variantes, incluindo aquelas com diferentes codificadores (como VGG ou ResNet), mecanismos de atenção e processamento multiescala. Sua influência se estende a outras áreas do aprendizado de máquina, incluindo IA generativa e visão computacional. A arquitetura também é uma linha de base comum em desafios de imagem médica e é amplamente implementada em estruturas como TensorFlow e PyTorch.
Implementações
Várias implementações de código aberto do U-Net estão disponíveis. O Tensorflow Unet de J Akeret (2017) é uma implementação popular. O código-fonte original do U-Net está disponível no grupo de Reconhecimento de Padrões e Processamento de Imagens do Departamento de Ciência da Computação da Universidade de Freiburg, Alemanha. Essas implementações facilitaram a adoção generalizada do U-Net tanto na pesquisa quanto na indústria.