1# Script is based on https://github.com/richzhang/colorization/blob/master/colorization/colorize.py
2# To download the caffemodel and the prototxt, see: https://github.com/richzhang/colorization/tree/master/colorization/models
3# To download pts_in_hull.npy, see: https://github.com/richzhang/colorization/blob/master/colorization/resources/pts_in_hull.npy
4import numpy as np
5import argparse
6import cv2 as cv
7
8def parse_args():
9    parser = argparse.ArgumentParser(description='iColor: deep interactive colorization')
10    parser.add_argument('--input', help='Path to image or video. Skip to capture frames from camera')
11    parser.add_argument('--prototxt', help='Path to colorization_deploy_v2.prototxt', required=True)
12    parser.add_argument('--caffemodel', help='Path to colorization_release_v2.caffemodel', required=True)
13    parser.add_argument('--kernel', help='Path to pts_in_hull.npy', required=True)
14
15    args = parser.parse_args()
16    return args
17
18if __name__ == '__main__':
19    W_in = 224
20    H_in = 224
21    imshowSize = (640, 480)
22
23    args = parse_args()
24
25    # Select desired model
26    net = cv.dnn.readNetFromCaffe(args.prototxt, args.caffemodel)
27
28    pts_in_hull = np.load(args.kernel) # load cluster centers
29
30    # populate cluster centers as 1x1 convolution kernel
31    pts_in_hull = pts_in_hull.transpose().reshape(2, 313, 1, 1)
32    net.getLayer(net.getLayerId('class8_ab')).blobs = [pts_in_hull.astype(np.float32)]
33    net.getLayer(net.getLayerId('conv8_313_rh')).blobs = [np.full([1, 313], 2.606, np.float32)]
34
35    if args.input:
36        cap = cv.VideoCapture(args.input)
37    else:
38        cap = cv.VideoCapture(0)
39
40    while cv.waitKey(1) < 0:
41        hasFrame, frame = cap.read()
42        if not hasFrame:
43            cv.waitKey()
44            break
45
46        img_rgb = (frame[:,:,[2, 1, 0]] * 1.0 / 255).astype(np.float32)
47
48        img_lab = cv.cvtColor(img_rgb, cv.COLOR_RGB2Lab)
49        img_l = img_lab[:,:,0] # pull out L channel
50        (H_orig,W_orig) = img_rgb.shape[:2] # original image size
51
52        # resize image to network input size
53        img_rs = cv.resize(img_rgb, (W_in, H_in)) # resize image to network input size
54        img_lab_rs = cv.cvtColor(img_rs, cv.COLOR_RGB2Lab)
55        img_l_rs = img_lab_rs[:,:,0]
56        img_l_rs -= 50 # subtract 50 for mean-centering
57
58        net.setInput(cv.dnn.blobFromImage(img_l_rs))
59        ab_dec = net.forward()[0,:,:,:].transpose((1,2,0)) # this is our result
60
61        (H_out,W_out) = ab_dec.shape[:2]
62        ab_dec_us = cv.resize(ab_dec, (W_orig, H_orig))
63        img_lab_out = np.concatenate((img_l[:,:,np.newaxis],ab_dec_us),axis=2) # concatenate with original image L
64        img_bgr_out = np.clip(cv.cvtColor(img_lab_out, cv.COLOR_Lab2BGR), 0, 1)
65
66        frame = cv.resize(frame, imshowSize)
67        cv.imshow('origin', frame)
68        cv.imshow('gray', cv.cvtColor(frame, cv.COLOR_RGB2GRAY))
69        cv.imshow('colorized', cv.resize(img_bgr_out, imshowSize))
70