Em suma, a tarefa deste projeto é colorir e restaurar fotografias antigas. Vou me aprofundar um pouco mais nos detalhes, mas primeiro, vamos ver as fotos! A propósito, a maioria das imagens de origem é tirada do subreddit r / TheWayWeWere, agradeço a todos por fotos grandes e de alta qualidade.

Estes são apenas alguns exemplos e são bastante típicos!

Maria Anderson como a pequena fada e sua página Lyubov Ryabtsova no balé Bela Adormecida no Teatro Imperial, São Petersburgo, Rússia, 1890

Uma mulher relaxa em sua sala de estar (1920, Suécia)

Estudantes de medicina se passando perto de um cadáver, por volta de 1890

Surfista no Havaí, 1890

Cavalo de giro, 1898

O interior do bar Miller e Shoemaker, 1899

Paris na década de 1880

Vista aérea de Edimburgo na década de 1920

Mulher do Texas em 1938

As pessoas na estação Waterloo assistem TV pela primeira vez, Londres, 1936

Lição de Geografia em 1850

Fumantes chineses de ópio em 1880

Observe que mesmo fotos muito antigas e / ou de baixa qualidade ainda são muito legais:

Deadwood, Dakota do Sul, 1877

Irmãos e irmãs em 1877 (Deadwood)

Praça Portsmouth em São Francisco, 1851

Samurai, por volta dos anos 1860

Obviamente, o modelo não é perfeito. Essa mão vermelha me deixa louco, mas, caso contrário, funciona de maneira fantástica:

Menina de Seneca Iroquois, 1908

Ela também pode colorir desenhos em preto e branco:

Detalhes técnicos

Este é um modelo de aprendizado profundo. Em particular, combinei as seguintes abordagens:

Auto-atenção GAN . A única coisa é que o Unet pré-treinado é usado como gerador e eu apenas o troquei para normalização espectral e, de fato, o mecanismo de Auto-Atenção. Esta é uma modificação bastante simples. Vou lhe dizer que a diferença é impressionante em comparação com a versão anterior do Wasserstein GAN, que tentei fazer funcionar. Gostei da teoria de Wasserstein GAN, mas na prática ela não funciona. Mas me apaixonei pela rede GAN de auto-atenção.
Uma estrutura de aprendizado como o crescimento progressivo de um GAN (mas não exatamente o mesmo). A diferença é que o número de camadas permanece constante: alterei o tamanho dos dados de entrada e ajustei a velocidade de aprendizado para que as transições entre os tamanhos fossem bem-sucedidas. Parece que produz o mesmo resultado final, mas aprende mais rápido, é mais estável e executa melhor a generalização.
Regra TTUR ( Regra de atualização em duas escalas de tempo). Aqui está bem claro: apenas uma iteração individual do gerador / discriminador (crítica) e uma maior velocidade de treinamento de discriminador.
A função de perda do gerador consiste em duas partes: uma delas é a principal função da Perda Perceptual (ou Perda de Recurso) baseada no VGG16 - simplesmente pressiona o modelo do gerador para replicar a imagem de entrada. A segunda parte é a estimativa de perdas do discriminador (crítica). Para os curiosos: somente a função Perda Perceptual não é suficiente para um bom resultado. Isso tende a simplesmente incentivar um monte de marrom / verde / azul - você sabe, enganando o teste, em que redes neurais são realmente boas! O ponto principal é que os próprios GANs aprendem essencialmente a função de perda para você, o que é realmente um grande passo em direção ao ideal que estamos buscando no aprendizado de máquina. E, é claro, os resultados melhorarão significativamente quando a própria máquina aprender o que você codificou anteriormente manualmente. Claro, este é o caso aqui.

A beleza deste modelo é que ele é muito bom em uma variedade de modificações de imagem. O que você vê acima são os resultados do modelo de coloração, mas esse é apenas um componente do pipeline que eu quero desenvolver com o mesmo modelo.

Em seguida, tentarei aperfeiçoar as imagens antigas, e o próximo item da agenda é um modelo para melhorar a saturação e a riqueza (desvanecimento). Agora ela está nos estágios iniciais do treinamento. Este é basicamente o mesmo modelo, mas com algumas configurações de contraste / brilho como uma simulação de fotos desbotadas e fotos tiradas com equipamento antigo / ruim. Já recebi alguns resultados encorajadores:

Detalhes do Projeto

Qual é a essência deste projeto? Eu só quero aplicar o GAN para que as fotos antigas pareçam muito, muito boas. E mais importante, tornará o projeto útil . E sim, definitivamente estou interessado em trabalhar com o vídeo, mas primeiro preciso descobrir como controlar esse modelo sob o controle do consumo de memória (este é um verdadeiro animal). Seria bom se os modelos não aprendessem de dois a três dias no 1080Ti (infelizmente, típico da GAN). Embora este seja meu filho, atualizarei e melhorarei ativamente o código no futuro próximo, mas tentarei tornar o programa o mais amigável possível, embora provavelmente haja algumas dificuldades com ele.

E juro que documentarei o código corretamente ... algum dia. É certo que sou uma daquelas pessoas que acredita em "código de auto-documentação" (LOL).

Modelo de auto-lançamento

O projeto é construído na maravilhosa biblioteca Fast.AI. Infelizmente, esta é uma versão antiga e ainda precisa ser atualizada para uma nova (definitivamente esta na ordem do dia). Portanto, os pré-requisitos, em resumo:

Biblioteca antiga Fast.AI. Tendo me enterrado no projeto por dois meses, perdi um pouco o que aconteceu com ele, porque o que agora está marcado como “antigo” não se parece realmente com o que tenho. Tudo mudou nos últimos dois meses ou mais. Portanto, se nada funcionar com outras versões, eu bifurquei aqui . Mais uma vez, a atualização para a versão mais recente está na agenda, peço desculpas antecipadamente.
Todas as dependências do Fast.AI : existem arquivos requirements.txt e environment.yml convenientes.
Pytorch 0.4.1 (spectral_norm é necessário, portanto, você precisa da versão estável mais recente).
JupyterLab .
Tensorboard (ou seja, instalação do Tensorflow) e TensorboardX . Eu acho que isso não é estritamente necessário, mas é muito mais fácil. Para sua conveniência, já forneci todos os ganchos / retornos de chamada necessários no Tensorboard! Existem exemplos de seu uso. Vale ressaltar que, por padrão, as imagens durante o processamento são gravadas no Tensorboard a cada 200 iterações, para que você tenha uma visão constante e conveniente do que o modelo faz.
ImageNet : Um excelente conjunto de dados para treinamento.
Placa gráfica poderosa . Eu realmente gostaria de ter mais memória do que 11 GB na minha GeForce 1080Ti. Se você tem algo mais fraco, será difícil. Unet e Critic são absurdamente grandes, mas quanto maiores, melhores os resultados.

Se você deseja iniciar o processamento de imagens agora mesmo sem treinar o modelo, pode fazer o download de pesos prontos aqui . Em seguida, abra ColorizationVisualization.ipynb no JupyterLab. Verifique se há uma linha com um link para os pesos:

colorizer_path = Path('/path/to/colorizer_gen_192.h5')

Depois, você precisa carregar o modelo do colorizador após a inicialização do netG:

 load_model(netG, colorizer_path)

Em seguida, basta colocar as imagens na pasta / test_images /, de onde você inicia o programa. Você pode visualizar os resultados no Notebook Jupyter com as seguintes linhas:

 vis.plot_transformed_image("test_images/derp.jpg", netG, md.val_ds, tfms=x_tfms, sz=500)

Eu salvaria um tamanho de cerca de 500 px, mais ou menos, se você executar o programa em uma GPU com muita memória (por exemplo, GeForce 1080Ti 11 GB). Se houver menos memória, você deverá reduzir o tamanho das imagens ou tentar executar na CPU. Na verdade, tentei fazer o último, mas por alguma razão o modelo funcionou muito, absurdamente devagar, e não encontrei tempo para investigar o problema. Os especialistas recomendaram a construção de Pytorch a partir das fontes, resultando em um grande aumento de desempenho. Hmm ... Naquele momento não foi antes disso.

Informações Adicionais

A visualização das imagens geradas conforme você aprende também pode ser feita no Jupyter: você só precisa configurá-la como verdadeira ao criar uma instância deste gancho de visualização:

GANVisualizationHook(TENSORBOARD_PATH, trainer, 'trainer', jupyter=True, visual_iters=100

Eu prefiro deixar false e apenas usar o Tensorboard. Acredite, você também quer fazer exatamente isso. Além disso, se você deixar que ele funcione por muito tempo, o Jupyter consumirá muita memória com essas imagens.

Os pesos dos modelos também são salvos automaticamente durante as execuções de treinamento do GANTrainer. Por padrão, eles são salvos a cada 1000 iterações (esta é uma operação cara). Eles são armazenados na pasta raiz que você especificou para treinamento e o nome corresponde a save_base_name especificado na programação de treinamento. Os pesos são armazenados separadamente para cada tamanho de treino.

Eu recomendaria navegar pelo código de cima para baixo, começando com o Jupyter Notebook. Tomo essas anotações simplesmente como uma interface conveniente para prototipagem e visualização; todo o resto irá para os arquivos .py assim que eu encontrar um local para eles. Já tenho exemplos de visualização que você pode ativar e ver com facilidade: basta abrir o xVisualization no Notebook, ele mostra as imagens de teste incluídas no projeto (elas estão em test_images).

Se você vir os Agendamentos da GAN, essa é a coisa mais feia do projeto, apenas a minha versão da implementação da GAN de aprendizado progressivo, adequada para o gerador Unet.

Os pesos pré-treinados para o gerador de coloração também estão aqui . O projeto DeFade ainda está em andamento, tentarei colocar bons pesos em alguns dias.

Geralmente, durante o treinamento, você verá os primeiros bons resultados na metade, ou seja, com um tamanho de 192px (se você usar os exemplos de treinamento fornecidos).

Tenho certeza de que estraguei tudo em algum lugar, então, por favor, deixe-me saber se é assim.

Problemas conhecidos

Você precisa brincar um pouco com o tamanho da imagem para obter o melhor resultado. O modelo sofre claramente de algumas proporções e proporções ao gerar imagens. Costumava ser muito pior, mas a situação melhorou significativamente com o aumento da iluminação / contraste e a introdução da aprendizagem progressiva. Eu quero eliminar completamente esse problema e focar nele, mas até agora não se desespere se a imagem parecer muito saturada ou com falhas estranhas. Provavelmente, tudo ficará normal após um pequeno redimensionamento. Como regra, para imagens super saturadas você precisa aumentar o tamanho.
Além do acima: obter as melhores imagens realmente se resume à arte de escolher os parâmetros ideais . Sim, os resultados são selecionados manualmente. Estou muito satisfeito com a qualidade e o modelo funciona de maneira bastante confiável, mas não perfeitamente. O projeto ainda está em andamento! Eu acho que a ferramenta pode ser usada como um "artista de IA", mas ainda não está pronta para o público em geral. Apenas não é o tempo.
Para complicar a situação: no momento, o modelo está brutalmente consumindo memória ; portanto, no meu cartão 1080Ti, ele processa fotos com um máximo de 500-600px. Aposto que existem muitas opções de otimização aqui, mas ainda não o fiz.
Adicionei zero padding ao gerador Unet para qualquer coisa que não caiba nos tamanhos esperados (é assim que eu posso carregar uma imagem de tamanho arbitrário). Foi um hack muito rápido e leva a estúpidas bordas direita e inferior na saída de imagens de teste de tamanho arbitrário. Estou certo de que existe uma maneira melhor, mas ainda não a encontrei.
Modelo adora roupas azuis. Não sei por que, a solução está na pesquisa!