OpenCV를 사용한 흑백 사진 컬러화: Caffe 모델 활용

original

colorized

위 사진들을 보면 이번 학습때 무엇을 배우는지 알 수 있다.

왼쪽의 흑백 사진을 오른쪽의 컬러화 시켰는데,

그 반대로 흑백화를 해야 납득이 갈 수준으로 컬러화를 해놨다.

신기하다.
Colorful Image Colorization 기법 개요

컬러화를 위한 사전 학습된 모델

• 작업에 필요한 특정 Caffe 모델 파일 소개:

1. colorization_deploy_v2.prototxt
   네트워크의 아키텍처를 명확히 정의하여 모델 구조를 쉽게 이해하고 수정
2. colorization_release_v2.caffe
   사전 학습된 가중치를 제공하여 바로 흑백 이미지 컬러화 작업에 사용
3. pts_in_hull.npy
   다양한 색상 예측을 위해 색상 클러스터 정보를 제공하여 컬러화 결과의 품질을 향상

사전 학습된 모델을 OpenCV에서 로드 및 사용하기

# 이미지 컬러화를 위한 스크립트
# 스크립트 출처: https://github.com/richzhang/colorization/blob/master/colorize.py
# caffemodel 및 prototxt 다운로드: https://github.com/richzhang/colorization/tree/master/models
# pts_in_hull.npy 다운로드: https://github.com/richzhang/colorization/blob/master/resources/pts_in_hull.npy

# 이미지 파일 경로 설정
file_path = "colorize/blackandwhite/"
# 흑백 이미지 목록 가져오기
blackandwhite_imgs = [f for f in listdir(file_path) if isfile(join(file_path, f))]
# 컬러 클러스터 정보 경로 설정
kernel = 'colorize/pts_in_hull.npy'

# 주 프로그램 실행
if __name__ == '__main__':

    # 원하는 모델 선택
    net = cv2.dnn.readNetFromCaffe("colorize/colorization_deploy_v2.prototxt",
                               "colorize/colorization_release_v2.caffemodel")

    # 클러스터 센터 불러오기
    pts_in_hull = np.load(kernel) 

    # 클러스터 센터를 1x1 컨볼루션 커널로 변환
    pts_in_hull = pts_in_hull.transpose().reshape(2, 313, 1, 1)
    net.getLayer(net.getLayerId('class8_ab')).blobs = [pts_in_hull.astype(np.float32)]
    net.getLayer(net.getLayerId('conv8_313_rh')).blobs = [np.full([1, 313], 2.606, np.float32)]

    # 각 흑백 이미지에 대해 반복
    for image in blackandwhite_imgs:
        # 이미지 불러오기
        img = cv2.imread(file_path+image)
        
        # 이미지를 RGB로 변환하여 처리
        img_rgb = (img[:,:,[2, 1, 0]] * 1.0 / 255).astype(np.float32)
        img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2LAB)
        
        # L 채널 추출
        img_l = img_lab[:,:,0]
        
        # 원본 이미지 크기 가져오기
        (H_orig,W_orig) = img_rgb.shape[:2] 

        # 네트워크 입력 크기로 이미지 크기 조정
        img_rs = cv2.resize(img_rgb, (224, 224)) 
        
        # 네트워크 입력 크기로 이미지 크기 조정
        img_lab_rs = cv2.cvtColor(img_rs, cv2.COLOR_RGB2Lab)
        img_l_rs = img_lab_rs[:,:,0]
        
        # 평균 중심화를 위해 50 빼기
        img_l_rs -= 50 

        # 네트워크 입력으로 이미지 전달
        net.setInput(cv2.dnn.blobFromImage(img_l_rs))
        
        # 결과 추출
        ab_dec = net.forward('class8_ab')[0,:,:,:].transpose((1,2,0)) 

        (H_out,W_out) = ab_dec.shape[:2]
        ab_dec_us = cv2.resize(ab_dec, (W_orig, H_orig))
        img_lab_out = np.concatenate((img_l[:,:,np.newaxis],ab_dec_us),axis=2) 
        
        # 원본 이미지 L과 결합
        img_bgr_out = np.clip(cv2.cvtColor(img_lab_out, cv2.COLOR_Lab2BGR), 0, 1)

        # 원본 이미지 표시
        imshow('Original', img)
        # 컬러화된 이미지를 원본 차원으로 크기 조정
        img_bgr_out = cv2.resize(img_bgr_out, (W_orig, H_orig), interpolation = cv2.INTER_AREA)
        imshow('Colorized', img_bgr_out)

다른 대상혁도 컬러화 시켜보자.

original

colorized

Holy..