Olá pessoal! Sou engenheiro de pesquisa da equipe de visão computacional do Grupo Mail.ru. Neste artigo, vou contar uma história de como criamos um projeto de restauração de fotos baseado em IA para fotos antigas de militares. O que é "restauração de fotos"? Consiste em três etapas:

• encontramos todos os defeitos da imagem: fraturas, arranhões, buracos;
  
• pintamos os defeitos descobertos, com base nos valores de pixel ao seu redor;
  
• nós colorimos a imagem.
  

Além disso, descreverei todas as etapas da restauração de fotos e mostrarei como obtivemos nossos dados, quais redes treinamos, o que realizamos e que erros cometemos.

Procurando defeitos

Queremos encontrar todos os pixels relacionados a defeitos em uma foto enviada. Primeiro, precisamos descobrir que tipo de imagem as pessoas enviarão. Conversamos com os fundadores do projeto "Immortal Regiment", uma organização não comercial que armazena as fotos herdadas da Segunda Guerra Mundial, que compartilharam seus dados conosco. Ao analisá-lo, percebemos que as pessoas enviam principalmente retratos individuais ou de grupo com um número moderado a grande de defeitos.

Então tivemos que coletar um conjunto de treinamento. O conjunto de treinamento para uma tarefa de segmentação é uma imagem e uma máscara em que todos os defeitos são marcados. A maneira mais fácil de fazer isso é permitir que os avaliadores criem as máscaras de segmentação. Obviamente, as pessoas sabem muito bem como encontrar defeitos, mas isso levaria muito tempo.



Pode levar uma hora ou o dia inteiro para marcar os pixels com defeito em uma foto. Portanto, não é fácil coletar um conjunto de treinamento de mais de 100 imagens em poucas semanas. É por isso que tentamos aumentar nossos dados e criar nossos próprios defeitos: tiramos uma boa foto, adicionamos defeitos usando passeios aleatórios na imagem e terminamos com uma máscara mostrando as partes da imagem com os defeitos. Sem aprimoramentos, temos 68 fotos rotuladas manualmente no conjunto de treinamento e 11 fotos no conjunto de validação.

A abordagem de segmentação mais popular: use o Unet com codificador pré-treinado e minimize a soma de BCE ( entropia cruzada binária ) e DICE ( Sørensen - coeficiente de dados ).

Quais problemas surgem quando usamos essa abordagem de segmentação para nossa tarefa?

• Mesmo que pareça que há muitos defeitos na foto, que é muito antiga e surrada, a área com defeitos ainda é muito menor que a não danificada. Para resolver esse problema, podemos aumentar o peso positivo da classe no AEC; um peso ideal seria a proporção de pixels limpos para os defeituosos.
  
• O segundo problema é que, se usarmos um Unet pronto para uso com codificador pré-treinado (Albunet-18, por exemplo), perderemos muitos dados posicionais. A primeira camada do Albunet-18 consiste em uma convolução com um núcleo 5 e um passo que é igual a dois. Permite que a rede trabalhe rapidamente. Negociamos o tempo de operação líquido para obter uma localização melhor dos defeitos: removemos o pool máximo após a primeira camada, diminuímos o passo para 1 e diminuímos o kernel de convolução para 3.
  
• Se trabalharmos com imagens pequenas compactando-as, por exemplo, para 256 x 256 ou 512 x 512 pixels, pequenos defeitos desaparecerão devido à interpolação. Portanto, precisamos trabalhar com imagens maiores. No momento, estamos segmentando defeitos em fotos com tamanho de 1024 x 1024 em produção. É por isso que tivemos que treinar a rede em grandes culturas de imagens. No entanto, isso causa problemas com um tamanho de lote pequeno em uma única GPU.
  
• Durante o treinamento, podemos caber cerca de 20 imagens em uma GPU. Por isso, acabamos com valores imprecisos de média e desvio padrão nas camadas BatchNorm. Podemos resolver esse problema usando o In-place BatchNorm , que, por um lado, economiza espaço na memória e, por outro, possui uma versão Synchronized BatchNorm, que sincroniza as estatísticas de todas as GPUs. Agora calculamos os valores de média e desvio padrão não para 20 imagens em uma única GPU, mas para 80 imagens de 4 GPUs. Isso melhora a convergência líquida.
  

Finalmente, ao aumentar o peso do BCE, alterar a arquitetura e usar o BatchNorm no local, melhoramos a segmentação. No entanto, não custaria muito para fazer algo ainda melhor adicionando o aumento do tempo de teste. Podemos executar a rede uma vez em uma imagem de entrada, espelhá-la e executar novamente a rede para encontrar todos os pequenos defeitos.



A rede converge em 18 horas em quatro GeForce 1080Ti. A inferência leva 290 ms. É bastante longo, mas esse é o preço do nosso desempenho melhor que o padrão. A validação DICE é igual a 0,35 e ROCAUC - 0,93.

Pintura de imagem

O mesmo com a tarefa de segmentação que usamos Unet. Para fazer a pintura, carregamos uma imagem original e uma máscara em que marcamos toda a área limpa com uma e com zeros - todos os pixels que queremos pintar. É assim que estávamos coletando dados: para qualquer foto de um conjunto de dados de imagens de código-fonte aberto, por exemplo, o OpenImagesV4, adicionamos defeitos semelhantes aos que vemos na vida real. Depois treinamos a rede para restaurar as peças que faltavam.

Como podemos modificar o Unet para esta tarefa?

Podemos usar convolução parcial em vez de uma original. A idéia é que, quando envolvemos uma área com algum kernel, não levamos em consideração os valores dos pixels de defeito. Isso torna a pintura mais precisa. Mostramos um exemplo do recente artigo da NVIDIA . Eles usaram o Unet com uma convolução bidimensional padrão na imagem do meio e uma convolução parcial - na imagem à direita.



Treinamos a rede por cinco dias. No último dia, congelamos o BatchNorms para tornar as bordas da parte pintada menos visíveis.

São necessários 50 ms líquidos para processar uma imagem de 512 x 512. Validação PSNR é igual a 26.4. No entanto, você não pode confiar totalmente nas métricas nesta tarefa. Para escolher o melhor modelo, executamos vários bons modelos em imagens de avaliação, anonimizamos os resultados e votamos nos que mais gostamos. Foi assim que escolhemos nosso modelo final.

Mencionei anteriormente que adicionamos artificialmente alguns defeitos às imagens limpas. Você deve sempre acompanhar o tamanho máximo de defeitos adicionados durante o treinamento; em um caso em que você alimenta uma imagem com um defeito muito grande na rede, que nunca é tratada na fase de treinamento, a rede fica solta e produz um resultado inaplicável. Portanto, se você precisar corrigir grandes defeitos, aumente seu conjunto de treinamento com eles.

Aqui está o exemplo de como nosso algoritmo funciona:

Colorização

Nós segmentamos os defeitos e os pintamos; o terceiro passo - reconstrução de cores. Como eu disse antes, existem muitos retratos individuais e em grupo entre as fotos do Regimento Imortal. Queríamos que nossa rede funcionasse bem com eles. Decidimos criar nossa própria coloração, pois nenhum dos serviços existentes conseguia colorir os retratos com rapidez e eficiência. Queremos que nossas fotos coloridas sejam mais confiáveis.



O GitHub possui um repositório popular para colorização de fotos. Faz um bom trabalho, mas ainda tem alguns problemas. Por exemplo, tende a pintar as roupas de azul. Por isso também a rejeitamos.

Então, decidimos criar um algoritmo para a coloração da imagem. A idéia mais óbvia: tire uma imagem em preto e branco e preveja três canais: vermelho, verde e azul. No entanto, podemos facilitar nosso trabalho: trabalhe não com a representação de cores RGB, mas com a representação de cores YCbCr. O componente Y é o brilho (luma). Uma imagem em preto e branco carregada é o canal Y e vamos reutilizá-la. Agora precisamos prever Cb e Cr: Cb é a diferença de cor e brilho azul e Cr - a diferença de cor e brilho vermelho.



Por que escolhemos a representação YCbCr? Um olho humano é mais sensível às mudanças de brilho do que às mudanças de cor. É por isso que reutilizamos o componente Y (brilho) com o qual o olho humano é mais sensível e prevemos Cb e Cr com os quais podemos cometer um erro, pois não podemos perceber muito bem a falsidade das cores. Essa característica específica foi amplamente usada no início da televisão em cores quando a capacidade do canal não era suficiente para transmitir todas as cores. A imagem foi transmitida em YCbCr, inalterada para o componente Y, e Cb e Cr foram reduzidos pela metade.

Como criar uma linha de base

Podemos usar o Unet com um codificador pré-treinado e minimizar a perda de L1 entre os valores de CbCr existentes e os previstos. Queremos pintar retratos e, portanto, além das fotos do OpenImages, precisamos de mais fotos específicas da tarefa.

Onde podemos obter fotos coloridas de pessoas vestidas com uniforme militar? Existem pessoas na internet que colorem fotos antigas por hobby ou por um preço. Eles fazem isso com muito cuidado, tentando ser muito preciso. Quando colorem um uniforme, ombreiras e medalhas, eles se referem aos materiais de arquivo, para que os resultados de seu trabalho sejam confiáveis. No total, usamos 200 fotos coloridas manualmente com pessoas de uniforme militar.

A outra fonte de dados útil é o site do Exército Vermelho dos Trabalhadores e Camponeses . Um de seus fundadores teve sua foto tirada em praticamente todos os uniformes soviéticos da Segunda Guerra Mundial disponíveis.



Em algumas fotos, ele imitou as poses de pessoas das famosas fotos de arquivo. É bom que suas fotos tenham fundo branco: isso nos permitiu aumentar muito bem os dados adicionando vários objetos naturais ao fundo. Também usamos alguns retratos regulares, complementando-os com insígnias e outros atributos de tempo de guerra.

Treinamos o AlbuNet-50 - é um Unet que usa o ResNet-50 pré-treinado como um codificador. A rede começou a dar resultados adequados: a pele era rosada, os olhos - cinza-azulados, as ombreiras - amarelados. No entanto, o problema é que ele deixa algumas áreas na foto intocadas. Isso foi causado pelo fato de que, de acordo com o erro L1, é tão ideal que é melhor não fazer nada do que tentar prever alguma cor.


Estamos comparando nosso resultado com uma foto Ground Truth - uma coloração manual feita por Klimbim

Como podemos resolver esse problema? Precisamos de um discriminador: uma rede neural que receba uma imagem e nos diga se ela parece realista ou não. Uma das figuras abaixo é colorida manualmente e a outra - pelo nosso gerador, AlbuNet-50. Como o ser humano distingue fotos coloridas manual e automaticamente? Ao olhar para os detalhes. Você pode dizer onde está a foto colorida automaticamente por nossa solução de linha de base?




Utilizamos o discriminador do artigo da GAN de auto-atenção . É uma pequena rede de convolução com a chamada Auto-Atenção criada nas camadas superiores. Isso nos permite "prestar mais atenção" aos detalhes da imagem. Também usamos normalização espectral. Você pode encontrar mais informações no documento acima. Treinamos a rede com uma combinação de perda de L1 e perda do discriminador. Agora, a rede coloriza melhor os detalhes da imagem e o fundo parece mais consistente. Mais um exemplo: à esquerda, o trabalho da rede treinado apenas com perda de L1; à direita - com uma combinação de perdas discriminadoras de L1.



O processo de treinamento levou dois dias em quatro GeForce 1080Ti. São necessários 30 ms para processar uma imagem de 512 x 512. Validação MSE - 34.4. Assim como na pintura, métricas que você não deseja confiar nas métricas. Por isso, escolhemos seis modelos com as melhores métricas de validação e votamos cegamente no melhor modelo.

Quando já criamos um sistema de produção e lançamos um site, continuamos experimentando e concluímos que é melhor minimizar a perda de L1 por pixel, mas a perda de percepção. Para calculá-lo, alimentamos as previsões líquidas e uma foto de base verdadeira à rede VGG-16, pegamos os mapas de recursos nas camadas inferiores e comparamos com o MSE. Essa abordagem pinta mais áreas e fornece resultados mais coloridos.

Recapitular

Unet é um modelo bem legal. Na primeira tarefa de segmentação, enfrentamos um problema durante o treinamento e trabalhamos com imagens de alta resolução e é por isso que usamos o In-Place BatchNorm. Em nossa segunda tarefa (Inpainting), usamos Convolução Parcial em vez de uma padrão, e isso nos permitiu obter melhores resultados. Ao trabalhar na colorização, adicionamos uma pequena rede discriminadora que penaliza o gerador por imagens irreais. Também usamos uma perda perceptiva.

Segunda conclusão - os avaliadores são essenciais. E não apenas durante o estágio de criação de máscaras de segmentação, mas também para a validação do resultado final. No final, fornecemos ao usuário três fotos: uma imagem original com defeitos pintados, uma foto colorida com defeitos pintados e uma simplesmente colorida, caso o algoritmo para pesquisa de defeitos e pintura incorreta.

Tiramos algumas fotos do projeto War Album e as processamos por essas neuronets. Aqui estão os resultados que obtivemos:



Além disso, aqui você pode ver as imagens originais e todas as etapas do processamento.