!mkdir models img  # 모델과 이미지 저장 폴더 생성

# 샘플 흑백 이미지 5장 다운로드

!gdown --id 1SGVNbOjBoitZw5nHHJQ5A3LL3uDqeznA --output ./img/sample01.jpg

!gdown --id 11dlWfIkDRnJ0BfojfBYbhS9q3Fl2TT1q --output ./img/sample02.jpg

!gdown --id 1U_StuW-A0dgtmzAoY4tOYHrSj5xROJms --output ./img/sample03.jpg

!gdown --id 1uD0u_9kG_gCI5pq8FBa8W6kmjQjYk1vn --output ./img/sample04.jpg

!gdown --id 1mQdvtEW2pwVslQ7bZb3OLtrr-LzA0cVa --output ./img/sample05.jpg

# 모델 구성 요소 다운로드

!gdown --id 1mLB3UZ0hgVPXguPRedHtb1UWAy2nlUyF --output ./models/pts_in_hull.npy           # ab 색상 포인트 클러스터

!gdown --id 16nCigXxDjH7KKyZxJ-KvE-Zlv79ZO-mz --output ./models/colorization_deploy.prototxt  # 네트워크 구조

!gdown --id 1PgvjpEId1mlX7i4MqNfsJeoiB1rJfxJ3 --output ./models/colorization_release.caffemodel  # 학습된 모델

!gdown --id 1Zk3Tv2PRMhnJQR6sEjla8V4idUt2rJBC --output ./models/colorization_release_norebal.caffemodel  # 리밸런싱 제거된 모델 (미사용)

import cv2

import numpy as np

import matplotlib.pyplot as plt

# 모델 설정

proto = './models/colorization_deploy.prototxt'

weight = './models/colorization_release.caffemodel'

# 학습된 ab 색 공간의 클러스터 (313개)

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)

# Caffe 모델 로드

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

# 모델 내부의 레이어에 클러스터 데이터를 넣어줌

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

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

img_path = './img/sample03.jpg'

img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)  # 흑백 이미지로 읽기

img_input = img.copy()

# 흑백 이미지를 BGR 3채널 이미지로 변환

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

img_rgb = img.copy()

# 0~1 사이로 정규화 후 RGB → LAB 색공간으로 변환

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

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

img_l = img_lab[:,:,0]  # L 채널(밝기)만 추출

# 네트워크 입력용 L 채널 크기 맞춤

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

input_img -= 50  # 네트워크 요구 사항: L 값에서 50 빼기

# 입력 이미지 시각화

plt.axis('off')

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

# 네트워크에 입력

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

# 예측된 ab 채널 출력

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

# ab 채널을 원본 이미지 크기로 리사이즈

prediction_resize = cv2.resize(prediction, (img.shape[1], img.shape[0]))

# LAB 형식으로 병합 (L + 예측된 ab)

prediction_lab = np.concatenate([img_l[:, :, np.newaxis], prediction_resize], axis=2)

# RGB로 변환 및 후처리

prediction_rgb = cv2.cvtColor(prediction_lab, cv2.COLOR_LAB2RGB)

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

prediction_rgb = prediction_rgb.astype(np.uint8)

# 출력 비교: 좌측은 흑백, 우측은 컬러화된 이미지

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

# 원본 L 채널

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

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

# 복원된 RGB 이미지

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

plt.imshow(prediction_rgb)