Olá pessoal, trabalho como programador de pesquisa na equipe de visão computacional do Mail.ru Group. Para o Dia da Vitória deste ano, decidimos fazer um projeto para a restauração de fotografias militares . O que é restauração de fotos? Consiste em três etapas:

• encontramos todos os defeitos da imagem: rupturas, arranhões, buracos;
  
• pinte os defeitos encontrados com base nos valores de pixel ao seu redor;
  
• colorir a imagem.
  

Neste artigo, analisarei cada um dos estágios da restauração em detalhes e mostrarei como e para onde levamos os dados, quais redes aprendemos, o que fizemos e em que etapas pisamos.

Pesquisa de defeitos

Queremos encontrar todos os pixels relacionados a defeitos na foto carregada. Primeiro, precisamos entender que tipo de fotografias dos anos de guerra as pessoas enviarão. Nos voltamos para os organizadores do projeto Immortal Regiment, que compartilharam dados conosco. Depois de analisá-los, percebemos que as pessoas frequentemente enviam retratos, únicos ou em grupo, com um número moderado ou grande de defeitos.

Então foi necessário coletar uma amostra de treinamento. A amostra de treinamento para a tarefa de segmentação é uma imagem e uma máscara na qual todos os defeitos são marcados. A maneira mais fácil é dar fotos aos marcadores. Obviamente, as pessoas são boas em encontrar defeitos, mas o problema é que a marcação é um processo muito longo.



Pode levar de uma hora a um dia útil inteiro para marcar os pixels relacionados aos defeitos em uma foto, por isso é difícil coletar uma amostra de mais de 100 fotos em poucas semanas. Portanto, tentamos de alguma forma complementar nossos dados e escrevemos defeitos: tiramos uma foto limpa, aplicamos defeitos artificiais e obtivemos uma máscara nos mostrando quais partes específicas da imagem estavam danificadas. A parte principal de nossa amostra de treinamento foi de 79 fotos marcadas manualmente, das quais 11 foram transferidas para a amostra de teste.

A abordagem mais popular para o problema de segmentação: leve o Unet com um codificador pré-treinado e minimize a quantidade $$$Bce$( entropia cruzada binária ) e $$$DICE$( Sørensen - coeficiente de dados ).

Que problemas surgem com essa abordagem no problema de segmentação de defeitos?

• Mesmo que pareça que há muitos defeitos na foto, que ela está muito suja e muito esfarrapada pelo tempo, a área ocupada por defeitos ainda é muito menor que a parte não danificada da imagem. Para resolver esse problema, você pode aumentar o peso da classe positiva em $$$Bce$, e o peso ideal será a proporção entre o número de todos os pixels puros e o número de pixels pertencentes aos defeitos.
  
• O segundo problema é que, se usarmos o Unet imediatamente com um codificador pré-treinado, por exemplo Albunet-18, perderemos muitas informações posicionais. A primeira camada de Albunet-18 consiste em uma convolução com um núcleo de 5 e passo igual a dois. Isso permite que a rede trabalhe rapidamente. Sacrificamos o tempo de operação da rede para melhor localização de defeitos: removemos o pool máximo após a primeira camada, reduzimos o passo para 1 e reduzimos o núcleo de convolução para 3.
  
• Se trabalharmos com imagens pequenas, por exemplo, compactando uma imagem para 256 x 256 ou 512 x 512, pequenos defeitos simplesmente desaparecerão devido à interpolação. Portanto, você precisa trabalhar com uma imagem grande. Agora, na produção, estamos segmentando defeitos em fotografias de 1024 x 1024. Portanto, era necessário treinar a rede neural em grandes quantidades de imagens grandes. E por isso, há problemas com o tamanho pequeno do lote em uma placa de vídeo.
  
• Durante o treinamento, temos cerca de 20 fotos colocadas em um cartão. Por esse motivo, a estimativa da média e da variação nas camadas BatchNorm é imprecisa. O BatchNorm no local nos ajuda a resolver esse problema, que primeiro economiza memória e, em segundo lugar, possui uma versão do BatchNorm Sincronizado, que sincroniza as estatísticas entre todos os cartões. Agora, consideramos a média e a variação não em 20 fotos em um cartão, mas em 80 fotos em 4 cartões. Isso melhora a convergência de rede.
  

No final, aumentando o peso $$$Bce$Alterando a arquitetura e usando o In-place BatchNorm, começamos a procurar defeitos na foto. Mas, de maneira barata, você pode se sair ainda melhor adicionando o aumento do tempo de teste. Podemos executar a rede uma vez na imagem de entrada, espelhá-la e executar a rede novamente, isso pode nos ajudar a encontrar pequenos defeitos.



Como resultado, nossa rede convergiu em quatro GeForce 1080Ti em 18 horas. A inferência leva 290 ms. Acontece por tempo suficiente, mas é uma taxa pelo fato de estarmos procurando por pequenos defeitos. Validação $$$DICE$igual a 0,35, e $$$ROCAUC$- 0,93.

Restauração de fragmentos

A Unet nos ajudou a resolver esse problema novamente. Na entrada, demos a ele a imagem original e uma máscara na qual marcamos espaços limpos com unidades e os pixels que queremos pintar com zeros. Coletamos os dados da seguinte maneira: tiramos da Internet um grande conjunto de dados com imagens, por exemplo, o OpenImagesV4, e adicionamos defeitos artificialmente semelhantes em forma aos encontrados na vida real. E depois disso, eles treinaram a rede para reparar as peças ausentes.

Como podemos modificar o Unet para esta tarefa?

Você pode usar Convolução Parcial em vez da convolução usual. A ideia dela é que, quando reduzimos uma região de uma imagem com um kernel, não levamos em conta os valores de pixel relacionados aos defeitos. Isso ajuda a tornar a pintura mais precisa. Um exemplo de um artigo da NVIDIA . Na imagem central, eles usaram Unet com a convolução usual e, à direita - com Convolução Parcial:



Treinamos a rede por 5 dias. No último dia, congelamos o BatchNorm, o que ajudou a tornar menos perceptíveis as bordas da parte pintada da imagem.

A rede processa uma imagem de 512 x 512 em 50 ms. A validação PSNR é 26.4. No entanto, as métricas não podem ser confiáveis ​​incondicionalmente nesta tarefa. Portanto, primeiro executamos vários bons modelos em nossos dados, anonimizamos os resultados e depois votamos nos que mais gostamos. Então escolhemos o modelo final.

Mencionei que adicionamos defeitos artificialmente para limpar imagens. Ao treinar, você precisa monitorar cuidadosamente o tamanho máximo dos defeitos sobrepostos, porque, com defeitos muito grandes que a rede nunca viu no processo de aprendizado, ela fantasiará muito e fornecerá um resultado absolutamente inaplicável. Portanto, se você precisar pintar sobre grandes defeitos, aplique também grandes defeitos durante o treinamento.

Aqui está um exemplo do algoritmo:

Coloração

Nós segmentamos os defeitos e os pintamos, o terceiro passo é a reconstrução da cor. Deixe-me lembrá-lo de que, entre as fotografias do "Regimento Imortal", existem muitos retratos individuais ou em grupo. E queríamos que nossa rede funcionasse bem com eles. Decidimos fazer nossa própria coloração, porque nenhum dos serviços conhecidos por nós retrata retratos de maneira rápida e eficaz.



O GitHub tem um repositório popular para colorir fotos. Em média, ele faz bem esse trabalho, mas ele tem vários problemas. Por exemplo, ele gosta de pintar roupas de azul. Portanto, também a rejeitamos.

Então, decidimos fazer uma rede neural para colorir. A idéia mais óbvia: tire uma imagem em preto e branco e preveja três canais, vermelho, verde e azul. Mas, de um modo geral, podemos simplificar nosso trabalho. Podemos trabalhar não com a representação RGB da cor, mas com a representação YCbCr. O componente Y é o brilho (luma). A imagem em preto e branco baixada é o canal Y, vamos reutilizá-la. Restava prever Cb e Cr: Cb é a diferença na cor azul e brilho, e Cr é a diferença na cor vermelha e brilho.



Por que escolhemos a visualização YCbCr? O olho humano é mais suscetível a mudanças no brilho do que a mudanças de cor. Portanto, reutilizamos o componente Y (brilho), algo ao qual o olho é inicialmente bem receptivo, e prevemos Cb e Cr, nos quais podemos cometer um pouco mais de erro, já que as pessoas percebem menos cores "falsas". Esse recurso começou a ser usado ativamente no início da televisão em cores, quando a largura de banda do canal não era suficiente para transmitir todas as cores na íntegra. A imagem foi transferida para YCbCr, transferida para o componente Y inalterada e Cb e Cr foram compactados duas vezes.

Como montar a linha de base

Você pode novamente usar o Unet com um codificador pré-treinado e minimizar a perda de L1 entre o CbCr real e o previsto. Queremos retratos em cores, portanto, além das fotos do OpenImages, precisamos adicionar fotos específicas à nossa tarefa.

Onde posso obter fotografias coloridas de pessoas em uniforme militar? Existem pessoas na Internet que pintam fotografias antigas como hobby ou por encomenda. Eles fazem isso com muito cuidado, tentando cumprir totalmente todas as nuances. Pintando o uniforme, dragonas, medalhas, eles recorrem a materiais de arquivo, para que o resultado de seu trabalho seja confiável. No total, foram utilizadas 200 fotografias pintadas à mão. A segunda fonte de dados útil é o site do Exército Vermelho dos Trabalhadores e Camponeses . Um de seus criadores foi fotografado em quase todas as variantes possíveis de um uniforme militar durante a Grande Guerra Patriótica.



Em algumas fotografias, ele repetiu as poses de pessoas de famosas fotografias de arquivo. É especialmente bom que ele tenha filmado em um fundo branco. Isso nos permitiu aumentar muito bem os dados, adicionando vários objetos naturais ao fundo. Também usamos retratos modernos comuns de pessoas, complementando-os com insígnias e outros atributos de roupas de guerra.

Nós treinamos o AlbuNet-50 - este é o Unet, no qual o AlbuNet-50 é usado como um codificador. A rede começou a dar resultados adequados: a pele é rosada, os olhos são azul acinzentados, as alças são amareladas. Mas o problema é que ela pintou as figuras com manchas. Isso se deve ao fato de que, do ponto de vista do erro L1, às vezes é mais lucrativo não fazer nada do que tentar prever alguma cor.


Estamos comparando nosso resultado com uma foto de Ground Truth - coloração manual do artista sob o apelido Klimbim

Como resolver este problema? Precisamos de um discriminador: uma rede neural, para a qual forneceremos imagens à entrada, e dirá quão realista essa imagem parece. Abaixo, uma fotografia é pintada à mão e a segunda por uma rede neural. Qual você acha?




Como discriminador, usamos o discriminador do artigo Self-Attention GAN . Trata-se de uma pequena rede convolucional, nas últimas camadas, na qual está embutida a chamada Auto-Atenção. Ele permite que você "preste atenção" aos detalhes da imagem. Também usamos normalização espectral. A explicação exata e a motivação podem ser encontradas no artigo. Nós treinamos uma rede com uma combinação de perda de L1 e o erro retornado pelo discriminador. Agora a rede pinta os detalhes da imagem melhor e o fundo é mais consistente. Outro exemplo: à esquerda é o resultado da rede treinada apenas com perda L1, à direita - com perda L1 e um erro discriminador.



Em quatro Geforce 1080Ti, o treinamento levou dois dias. A rede funcionou em 30 ms na imagem 512 x 512. A validação MSE foi 34.4. Como no problema da pintura, as métricas não podem ser totalmente confiáveis. Portanto, selecionamos 6 modelos que tinham as melhores métricas para validação e votamos cegamente no melhor modelo.

Depois de lançar o modelo em produção, continuamos os experimentos e chegamos à conclusão de que é melhor minimizar a perda de L1 por pixel, mas a perda de percepção. Para calculá-lo, é necessário executar a previsão de rede e a foto de origem através da rede VGG-16, pegar os mapas de atributos nas camadas inferiores e compará-los de acordo com o MSE. Essa abordagem pinta mais áreas e ajuda a obter uma imagem mais colorida.

Conclusões e Conclusão

Unet é um modelo legal. No primeiro problema de segmentação, encontramos um problema no treinamento e no trabalho com imagens de alta resolução; portanto, usamos o In-Place BatchNorm. Na segunda tarefa (Pintura), em vez da convolução usual, usamos Convolução Parcial, que ajudou a obter melhores resultados. No problema de coloração da Unet, adicionamos uma pequena rede de discriminadores que multou o gerador por uma imagem com aparência irreal e usou perda perceptiva.

A segunda conclusão é que os acessadores são importantes. E não apenas na etapa de marcar as figuras antes do treinamento, mas também na validação do resultado final, porque em problemas de defeitos ou coloração da pintura, você ainda precisa validar o resultado com a ajuda de uma pessoa. Damos ao usuário três fotos: a original com os defeitos removidos, colorida com os defeitos removidos e apenas a foto colorida, caso o algoritmo de pesquisa e pintura de defeitos esteja errado.

Tiramos algumas fotos do projeto Álbum Militar e as processamos com nossas redes neurais. Aqui estão os resultados obtidos:



E aqui você pode vê-los na resolução original e em cada estágio do processamento.