Como colorir imagens em preto e branco com Pyhton e Deep Learning
Introdução
Colorir imagens em preto e branco é um árduo trabalho que tradicionalmente é feito com softwares específicos, como o Photoshop. Imagine, por exemplo, a complexidade para colorir um rosto, com diversos tons envolvidos.
Dê uma olhada nesse vídeo e contemple a difícil tarefa que diversos artistas tem ao colorir fotos do passado.
O objetivo desse artigo é apresentar uma abordagem, baseada em Deep Learning, para automaticamente colorir imagens. A técnica foi proposta pelos pesquisadores Richard Zhang, Phillip Isola e Alexei Efros, da Universidade da Califórnia, no paper Colorful Image Colorization.
Se você não está familiarizado com o termo Deep Learning, dê uma olhada nesse outro artigo do blog.
A estratégia proposta pelos pesquisadores consiste em, dado uma imagem em preto e branco como input, produzir como resultado do processo de colorização um output que seja realista.
Considere as fotos em preto e branco abaixo. A primeira vista, descobrir as cores reais parece uma tarefa impossível. Entretanto, com um olhar mais atento, percebemos que em muitos casos a cena exposta e sua textura nos mostram diversas regiões de conhecimento geral, como a grama, tipicamente verde, ou o céu, tipicamente azul. Mas é claro que essa linha de raciocínio não funciona para todos os casos, por exemplo, a bola de críquete pode na verdade, não ser vermelha e sim, verde, amarela etc.
Contudo, como os pesquisadores deixam claro no paper, o objetivo do trabalho é produzir um resultado que seja plausível e possa, potencialmente, enganar observadores humanos. Inclusive, um dos testes em que o algoritmo foi avaliado se tratou de uma espécie de teste de Turing, onde diversos participantes humanos escolheram entre uma imagem que passou pelo processo de colorização e a versão original. Cerca de 32% das avaliações foram enganadas, percentual bem superior a alcançada por trabalhos semelhantes.
Definindo o problema de colorização
Normalmente, a forma como uma imagem é codificada digitalmente é feita usando o modelo RGB. Esse modelo de cores funciona de forma aditiva, onde o vermelho (Red), verde (Green) e azul (Blue) são adicionados de várias maneiras, para reproduzir uma ampla variedade de cores.
No entanto, o modelo usado nesse projeto se chama CIELAB. Nesse modelo, as cores são expressas em três valores numéricos.
- O canal L codifica a informação referente a intensidade luminosa.
- O canal A codifica os componentes verde e vermelho.
- O canal B codifica os componentes azul e amarelo.
Esse modelo foi proposto após a teoria das cores opostas, onde duas cores não podem ser verde e vermelho ou amarela e azul ao mesmo tempo. O CIELAB foi projetado para ser perceptivelmente uniforme em relação à visão de cores humana, significando que a mesma quantidade de mudança numérica nesses valores corresponde a aproximadamente a mesma quantidade de mudança percebida visualmente.
Ao contrário do modelo RGB, o CIELAB foi projetado para aproximar a visão humana, visando a uniformidade perceptual, e seu componente L se aproxima muito da percepção humana de leveza.
O L é o canal usado como entrada do modelo de Deep Learning, que foi treinado para estimar os componentes remanescentes, “A” e “B”.
O modelo de Deep Learning
Para o modelo de aprendizagem profunda, os pesquisadores escolheram a rede neural do tipo CNN (Convolutional Neural Networks) . Na imagem abaixo, você pode conferir a arquitetura usada.
A rede é composta por diversas camadas convolucionais, cada uma com 2 ou 3 layers, seguidos pela função de ativação do tipo ReLU (Rectified Linear Unit).
Para fazer o modelo de Deep Learning, os pesquisadores usaram o framework Caffe, desenvolvido pelo grupo de pesquisa em Inteligência Artificial da Universidade da Califórnia.
Para o treinamento do algoritmo, foi usado um conjunto de um milhão e trezentas mil imagens da base de dados ImageNet. O download do modelo pré-treinado pode ser feito através do Github do projeto.
Implementação com Pyhton
Para o projeto, vamos nos basear no código feito pelo Adrian Rosebrock, do excelente blog sobre Visão Computacional, o pyimagesearch. Como primeiro passo, precisamos criar uma pasta para armazenar as imagens de input e os arquivos disponibilizados pelos pesquisadores, que podem ser encontrados nesse link.
- colorization_release_v2.caffemodel
- colorization_release_v2_norebal.caffemodel
- colorization_release_v1.caffemodel.
A primeira etapa no Pyhton será a importação das bibliotecas necessárias.
import numpy as npimport argparseimport cv2
A Open CV será a biblioteca usada para o tratamento das imagens e a argparse para enviar argumentos pela linha de comando. Vamos passar os Paths dos arquivos necessários com o seguinte código.
ap = argparse.ArgumentParser()ap.add_argument("-i", "--image", type=str, required=True, help="imagem desejada")ap.add_argument("-p", "--prototxt", type=str, required=True, help="path para o arquivo Caffe prototxt")ap.add_argument("-m", "--model", type=str, required=True, help="path para o arquivo Caffe pre-trained model")ap.add_argument("-c", "--points", type=str, required=True, help="path para o arquivo cluster center points")args = vars(ap.parse_args())
Com os Paths especificados, o script poderá ser executado com inputs diferentes (imagens diferentes).
Agora vamos usar a Open CV para ler o modelo através da função cv2.dnn.readNetFromCaffe. Com isso, os arquivos conseguiram ser passados através da linha de comando.
print("[INFO] loading model...")net = cv2.dnn.readNetFromCaffe(args["prototxt"], args["model"])
Semelhante ao último passo, vamos usar agora a função np.lood para carregar os cluster center points. Perceba que o formato do arquivo é referente ao NumPy.
pts = np.load(args["points"])As próximas linhas terão a tarefa de tratar cada ponto como uma convolução 1×1 e carregar para o modelo.
class8 = net.getLayerId("class8_ab")conv8 = net.getLayerId("conv8_313_rh")pts = pts.transpose().reshape(2, 313, 1, 1)net.getLayer(class8).blobs = [pts.astype("float32")]net.getLayer(conv8).blobs = [np.full([1, 313], 2.606, dtype="float32")]
Agora, vamos tratar a imagem usando a Open CV.
image = cv2.imread(args["image"])scaled = image.astype("float32") / 255.0lab = cv2.cvtColor(scaled, cv2.COLOR_BGR2LAB)
Na primeira linha, usamos cv2.imread para carregar a imagem. Para o pré-processamento, na segunda linha fizemos um redimensionamento da intensidade dos pixels para o intervalo de [0,1] e na terceira, convertemos o modelo de RGB para LAB.
Seguindo o pre-processamento, faremos um resize da imagem para 224×224, que é a dimensão exigida pelo modelo dos pesquisadores e depois vamos extrair o canal L.
resized = cv2.resize(lab, (224, 224))L = cv2.split(resized)[0]L -= 50
Agora, vamos passar o canal L, extraído, como input para rede, para que ela seja capaz de prever os canais A e B.
'print("[INFO] colorizing image...")'net.setInput(cv2.dnn.blobFromImage(L))ab = net.forward()[0, :, :, :].transpose((1, 2, 0))ab = cv2.resize(ab, (image.shape[1], image.shape[0]))
Perceba que na última linha fizemos um resize mais uma vez, dessa vez para colocar a imagem de volta a dimensão da entrada.
Agora, vamos para a etapa do pós-processamento.
Primeiro, vamos extrair o canal L da imagem original e concatena-lo com os canais A e B, da predição, formando assim a imagem colorida.
L = cv2.split(lab)[0]colorized = np.concatenate((L[:, :, np.newaxis], ab), axis=2)
Segundo, converter modelos, de LAB para RGB, e excluir qualquer pixels que tenha ficado fora do intervalo especificado, [0,1].
colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2BGR)colorized = np.clip(colorized, 0, 1)
Terceiro, colocar as intensidades dos pixels de volta ao intervalo [0,255], pois durante o pré-processamento dividimos por 255 e agora estamos multiplicando por 255.
colorized = (255 * colorized).astype("uint8")E por último, vamos mostrar na tela as duas imagens, a original, em preto e branco , e a versão colorida.
cv2.imshow("Original", image)cv2.imshow("Colorized", colorized)cv2.waitKey(0)
Resultados
Nesta seção, vamos mostrar o resultado do processo para algumas imagens.
Conclusão
Através dos resultados, percebemos que o processo teve um output mais realista nas imagens onde a natureza está em maior evidência. Esse comportamento era esperado por conta da estratégia adotada pelos pesquisadores, explicada na introdução. Mas de qualquer forma, o processo de colorização se mostrou bastante eficaz, dada a dificuldade do problema e a autonomia do sistema.
Esse pequeno projeto é mais uma demonstração da imensa gama de aplicabilidade que as tecnologias envolvendo Deep Learning podem ter na sociedade. Espero que tenha curtido e feedbacks são sempre bem-vindos!
Abraço
Referências
Os seguintes blogs foram usados para a produção deste artigo:
