Neural colorization

import sys
import numpy
import caffe
import os
import skimage.color
import scipy.ndimage.interpolation
import scipy.misc

if (len(sys.argv) != 2):
        sys.exit("Usage: python2 main.py <target_image>")

filename = sys.argv[1]
gpu_id = 0
caffe.set_mode_gpu()
caffe.set_device(gpu_id)
net = caffe.Net("colorization_deploy_v0.prototxt", "colorization_release_v0.caffemodel", caffe.TEST)

(H_in, W_in) = net.blobs["data_l"].data.shape[2:]
(H_out, W_out) = net.blobs["class8_ab"].data.shape[2:]
net.blobs["Trecip"].data[...] = 6 / numpy.log(10)

img_rgb = caffe.io.load_image(filename)
img_lab = skimage.color.rgb2lab(img_rgb)
img_l = img_lab[:,:,0]
(H_orig,W_orig) = img_rgb.shape[:2]

img_lab_bw = img_lab.copy()
img_lab_bw[:,:,1:] = 0
img_rgb_bw = skimage.color.lab2rgb(img_lab_bw)

img_rs = caffe.io.resize_image(img_rgb,(H_in, W_in))
img_lab_rs = skimage.color.rgb2lab(img_rs)
img_l_rs = img_lab_rs[:,:,0]

img_pad = numpy.ones((H_orig, W_orig / 10, 3))

net.blobs["data_l"].data[0, 0,:,:] = img_l_rs - 50
net.forward()

ab_dec = net.blobs['class8_ab'].data[0,:,:,:].transpose((1, 2, 0))
ab_dec_us = scipy.ndimage.interpolation.zoom(ab_dec, (1.0 * H_orig / H_out, 1.0 * W_orig / W_out, 1))
img_lab_out = numpy.concatenate((img_l[:,:,numpy.newaxis], ab_dec_us), axis = 2)
img_rgb_out = numpy.clip(skimage.color.lab2rgb(img_lab_out), 0, 1)

splitted_filename = filename.split("/")
scipy.misc.imsave("./" + "".join(splitted_filename[0:-1]) + "/out_" + splitted_filename[-1], img_rgb_out)