Python 흑백 사진 AI로 컬러로 바꾸는 프로그램 모델 파일 추가

python으로 흑백사진을 채색하는 프로그램입니다.

 

제가 만든것은 아니고 다른데서 퍼와서 그냥 간략하게 돌아가도록 한번 해본것인데

 

원 출처는 https://opentutorials.org/module/3811/22946 

흑백 영화 - 컬러 복원 - 파이썬으로 재밌는거 만들기 [파재만]

 

 

이 사이트로 들어가시면 모든 코드들과 정보들이 있습니다.

코드와 설명은 위쪽 사이트를 참고하시기를 바라며, (사실상 저분이 거의 하신거지 저는 그냥 퍼온거에요..)

 

 

 

여기서 제가 하나 조금 힘들었던 부분이 무엇이었냐면

게시글이 몇년전에 올라온것이라 그런지

사용하던 모델이 링크가 전부 만료가 되어 있었습니다.

 

이미 몇년이 지나버려서 

그 모델 파일을 구하는게 힘들었기 때문에 추가적으로 도움이 되고자 작성을 하려고 합니다. 

 

일단 먼저 코드를 보겠습니다.

 

import cv2 # opencv 3.4.2+ required

import os

import numpy as np

import matplotlib.pyplot as plt

proto = './models/colorization_deploy_v2.prototxt'

weights = './models/colorization_release_v2.caffemodel'

# colorization_release_v2_norebal.caffemodel is trained with a classification loss with no class re-balancing term.

# The results are duller but "safer" colorizations

# weights = './models/colorization_release_v2_norebal.caffemodel'

# load cluster centers

pts_in_hull = np.load('./models/pts_in_hull.npy')

pts_in_hull = pts_in_hull.transpose().reshape(2, 313, 1, 1).astype(np.float32)

# load model

net = cv2.dnn.readNetFromCaffe(proto, weights)

# net.getLayerNames()

# populate cluster centers as 1x1 convolution kernel

net.getLayer(net.getLayerId('class8_ab')).blobs = [pts_in_hull]

# scale layer doesn't look work in OpenCV dnn module, we need to fill 2.606 to conv8_313_rh layer manually

net.getLayer(net.getLayerId('conv8_313_rh')).blobs = [np.full((1, 313), 2.606, np.float32)]

img_path = 'sample_23_input.jpg'

img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)

img_input = img.copy()

# convert BGR to RGB

img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)

img_rgb = img.copy()

# normalize input

img_rgb = (img_rgb / 255.).astype(np.float32)

# convert RGB to LAB

img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2Lab)

# only L channel to be used

img_l = img_lab[:, :, 0]

input_img = cv2.resize(img_l, (224, 224))

input_img -= 50 # subtract 50 for mean-centering

# plot images

# fig = plt.figure(figsize=(10, 5))

# fig.add_subplot(1, 2, 1)

# plt.imshow(img_rgb)

# fig.add_subplot(1, 2, 2)

plt.axis('off')

plt.imshow(input_img, cmap='gray')

net.setInput(cv2.dnn.blobFromImage(input_img))

pred = net.forward()[0,:,:,:].transpose((1, 2, 0))

# resize to original image shape

pred_resize = cv2.resize(pred, (img.shape[1], img.shape[0]))

# concatenate with original image L

pred_lab = np.concatenate([img_l[:, :, np.newaxis], pred_resize], axis=2)

# convert LAB to RGB

pred_rgb = cv2.cvtColor(pred_lab, cv2.COLOR_Lab2RGB)

pred_rgb = np.clip(pred_rgb, 0, 1) * 255

pred_rgb = pred_rgb.astype(np.uint8)

# plot prediction result

fig = plt.figure(figsize=(20, 10))

fig.add_subplot(1, 2, 1).axis('off')

plt.imshow(img_l, cmap='gray')

fig.add_subplot(1, 2, 2).axis('off')

plt.imshow(pred_rgb)

# plt.savefig(output_filename)

# save result image file

filename, ext = os.path.splitext(img_path)

input_filename = '%s_input%s' % (filename, ext)

output_filename = '%s_output%s' % (filename, ext)

pred_rgb_output = cv2.cvtColor(pred_rgb, cv2.COLOR_RGB2BGR)

cv2.imwrite(input_filename, img_input)

cv2.imwrite(output_filename, np.concatenate([img, pred_rgb_output], axis=1))

 

 

 

코드를 보시면 위쪽에 이런 코드들이 있습니다.

 

 

proto = './models/colorization_deploy_v2.prototxt'

weights = './models/colorization_release_v2.caffemodel'

pts_in_hull = np.load('./models/pts_in_hull.npy')

 

 

 

이 코드들이 모델 파일을 불러오는 코드였습니다.

근데 pts_in_hull.npy 파일과 colorization_deploy_v2.prototxt 파일 두개는 비교적 구하기가 쉬웠고 용량도 크지 않았습니다.

먼저 두 파일은 다운받으시면 됩니다.

colorization_deploy_v2.prototxt

0.01MB

pts_in_hull.npy

0.00MB

 

 

 

하지만 colorization_release_v2.caffemodel 파일은 구하기도 어렵고 용량도 상당히 컸었습니다.

 

그래서 한참의 구글링 끝에 일단 파일 사이트를 찾아내서 링크를 걸겠습니다.

 

https://huggingface.co/spaces/viveknarayan/Image_Colorization/blob/main/colorization_release_v2.caffemodel

colorization_release_v2.caffemodel · viveknarayan/Image_Colorization at main

 

링크에 들어가셔서 보시면 이런 화면이 나옵니다.

 

 

 

 

빨간색 네모로 표시한 다운로드를 받아주시면 colorization_release_v2.caffemodel 파일이 다운 받아집니다.

 

그럼 이렇게 같은 폴더에 넣어주신 후에

 

 

 

 

 

 

colorize를 실행하여 코드에 약간의 수정을 해줍니다.

 

proto = './models/colorization_deploy_v2.prototxt'

weights = './models/colorization_release_v2.caffemodel'

pts_in_hull = np.load('./models/pts_in_hull.npy')

 

 

 

코드에서 제가 빨간색으로 표시한 경로 부분을 삭제합니다.

 

img_path = 'sample_23_input.jpg'

 

이 부분에 흑백사진으로 쓸 파일명을 넣어주면 됩니다. (확장자 꼭 잊지마세요)

 

 

 

 

아주 잘 됩니다.