<reponame>SharmaAjay19/VideoImageColorization
import numpy as np
import os
import time
import skimage.color as color
import matplotlib.pyplot as plt
import scipy.ndimage.interpolation as sni
import caffe
import argparse
import cv2
import shutil

def parse_args():
    parser = argparse.ArgumentParser(description='iColor: deep interactive colorization')
    parser.add_argument('-vid_in',dest='vid_in',help='grayscale image to read in', type=str)
    parser.add_argument('-vid_out',dest='vid_out',help='colorized image to save off', type=str)
    parser.add_argument('--gpu', dest='gpu', help='gpu id', type=int, default=0)
    parser.add_argument('--prototxt',dest='prototxt',help='prototxt filepath', type=str, default='./models/colorization_deploy_v2.prototxt')
    parser.add_argument('--caffemodel',dest='caffemodel',help='caffemodel filepath', type=str, default='./models/colorization_release_v2.caffemodel')

    args = parser.parse_args()
    return args

def readVideoFrames(vid_path):
	vidcap = cv2.VideoCapture(vid_path)
	success,image = vidcap.read()
	print("Read video")
	tempdir = os.getcwd()+"/video_temp" + str(time.clock())
	if not os.path.exists(tempdir):
		os.makedirs(tempdir)
	count = 0
	while success:
		cv2.imwrite(tempdir + "/frame%d.jpg" % count, image)     # save frame as JPEG file      
		success,image = vidcap.read()
		count += 1
	print("Read %d frames"%count)
	return count, tempdir

def createVideoFromFrames(dirpath, vid_path, nframes):
	frame = cv2.imread(os.path.join(dirpath, "frame0_1.jpg"))
	height, width, layers = frame.shape
	fourcc = cv2.VideoWriter_fourcc(*'mp4v')
	video = cv2.VideoWriter(vid_path, fourcc, 20, (width,height))
	for frame in range(nframes):
		video.write(cv2.imread(os.path.join(dirpath, "frame%d_1.jpg"%frame)))
	cv2.destroyAllWindows()
	video.release()
	shutil.rmtree(dirpath)

if __name__ == '__main__':
	args = parse_args()

	caffe.set_mode_cpu()
	#caffe.set_device(args.gpu)

	# Select desired model
	net = caffe.Net(args.prototxt, args.caffemodel, caffe.TEST)

	(H_in,W_in) = net.blobs['data_l'].data.shape[2:] # get input shape
	(H_out,W_out) = net.blobs['class8_ab'].data.shape[2:] # get output shape

	pts_in_hull = np.load('./resources/pts_in_hull.npy') # load cluster centers
	net.params['class8_ab'][0].data[:,:,0,0] = pts_in_hull.transpose((1,0)) # populate cluster centers as 1x1 convolution kernel
	# print 'Annealed-Mean Parameters populated'

	nframes, dirpath = readVideoFrames(args.vid_in)
	for frame in range(nframes):
		# load the original image
		img_rgb = caffe.io.load_image(dirpath+("/frame%d.jpg"%frame))

		img_lab = color.rgb2lab(img_rgb) # convert image to lab color space
		img_l = img_lab[:,:,0] # pull out L channel
		(H_orig,W_orig) = img_rgb.shape[:2] # original image size

		# create grayscale version of image (just for displaying)
		img_lab_bw = img_lab.copy()
		img_lab_bw[:,:,1:] = 0
		img_rgb_bw = color.lab2rgb(img_lab_bw)

		# resize image to network input size
		img_rs = caffe.io.resize_image(img_rgb,(H_in,W_in)) # resize image to network input size
		img_lab_rs = color.rgb2lab(img_rs)
		img_l_rs = img_lab_rs[:,:,0]

		net.blobs['data_l'].data[0,0,:,:] = img_l_rs-50 # subtract 50 for mean-centering
		net.forward() # run network

		ab_dec = net.blobs['class8_ab'].data[0,:,:,:].transpose((1,2,0)) # this is our result
		ab_dec_us = sni.zoom(ab_dec,(1.*H_orig/H_out,1.*W_orig/W_out,1)) # upsample to match size of original image L
		img_lab_out = np.concatenate((img_l[:,:,np.newaxis],ab_dec_us),axis=2) # concatenate with original image L
		img_rgb_out = (255*np.clip(color.lab2rgb(img_lab_out),0,1)).astype('uint8') # convert back to rgb

		plt.imsave(dirpath+("/frame%d_1.jpg"%frame), img_rgb_out)
		print("Processed %d out of %d frames" % (frame+1, nframes))
	createVideoFromFrames(dirpath, args.vid_out, nframes)
import numpy as np
import os
import time
import skimage.color as color
import matplotlib.pyplot as plt
import scipy.ndimage.interpolation as sni
import caffe
import argparse
import cv2
import shutil

def parse_args():
    parser = argparse.ArgumentParser(description='iColor: deep interactive colorization')
    parser.add_argument('-vid_in',dest='vid_in',help='grayscale image to read in', type=str)
    parser.add_argument('-vid_out',dest='vid_out',help='colorized image to save off', type=str)
    parser.add_argument('--gpu', dest='gpu', help='gpu id', type=int, default=0)
    parser.add_argument('--prototxt',dest='prototxt',help='prototxt filepath', type=str, default='./models/colorization_deploy_v2.prototxt')
    parser.add_argument('--caffemodel',dest='caffemodel',help='caffemodel filepath', type=str, default='./models/colorization_release_v2.caffemodel')

    args = parser.parse_args()
    return args

def readVideoFrames(vid_path):
	vidcap = cv2.VideoCapture(vid_path)
	success,image = vidcap.read()
	print("Read video")
	tempdir = os.getcwd()+"/video_temp" + str(time.clock())
	if not os.path.exists(tempdir):
		os.makedirs(tempdir)
	count = 0
	while success:
		cv2.imwrite(tempdir + "/frame%d.jpg" % count, image)     # save frame as JPEG file      
		success,image = vidcap.read()
		count += 1
	print("Read %d frames"%count)
	return count, tempdir

def createVideoFromFrames(dirpath, vid_path, nframes):
	frame = cv2.imread(os.path.join(dirpath, "frame0_1.jpg"))
	height, width, layers = frame.shape
	fourcc = cv2.VideoWriter_fourcc(*'mp4v')
	video = cv2.VideoWriter(vid_path, fourcc, 20, (width,height))
	for frame in range(nframes):
		video.write(cv2.imread(os.path.join(dirpath, "frame%d_1.jpg"%frame)))
	cv2.destroyAllWindows()
	video.release()
	shutil.rmtree(dirpath)

if __name__ == '__main__':
	args = parse_args()

	caffe.set_mode_cpu()
	#caffe.set_device(args.gpu)

	# Select desired model
	net = caffe.Net(args.prototxt, args.caffemodel, caffe.TEST)

	(H_in,W_in) = net.blobs['data_l'].data.shape[2:] # get input shape
	(H_out,W_out) = net.blobs['class8_ab'].data.shape[2:] # get output shape

	pts_in_hull = np.load('./resources/pts_in_hull.npy') # load cluster centers
	net.params['class8_ab'][0].data[:,:,0,0] = pts_in_hull.transpose((1,0)) # populate cluster centers as 1x1 convolution kernel
	# print 'Annealed-Mean Parameters populated'

	nframes, dirpath = readVideoFrames(args.vid_in)
	for frame in range(nframes):
		# load the original image
		img_rgb = caffe.io.load_image(dirpath+("/frame%d.jpg"%frame))

		img_lab = color.rgb2lab(img_rgb) # convert image to lab color space
		img_l = img_lab[:,:,0] # pull out L channel
		(H_orig,W_orig) = img_rgb.shape[:2] # original image size

		# create grayscale version of image (just for displaying)
		img_lab_bw = img_lab.copy()
		img_lab_bw[:,:,1:] = 0
		img_rgb_bw = color.lab2rgb(img_lab_bw)

		# resize image to network input size
		img_rs = caffe.io.resize_image(img_rgb,(H_in,W_in)) # resize image to network input size
		img_lab_rs = color.rgb2lab(img_rs)
		img_l_rs = img_lab_rs[:,:,0]

		net.blobs['data_l'].data[0,0,:,:] = img_l_rs-50 # subtract 50 for mean-centering
		net.forward() # run network

		ab_dec = net.blobs['class8_ab'].data[0,:,:,:].transpose((1,2,0)) # this is our result
		ab_dec_us = sni.zoom(ab_dec,(1.*H_orig/H_out,1.*W_orig/W_out,1)) # upsample to match size of original image L
		img_lab_out = np.concatenate((img_l[:,:,np.newaxis],ab_dec_us),axis=2) # concatenate with original image L
		img_rgb_out = (255*np.clip(color.lab2rgb(img_lab_out),0,1)).astype('uint8') # convert back to rgb

		plt.imsave(dirpath+("/frame%d_1.jpg"%frame), img_rgb_out)
		print("Processed %d out of %d frames" % (frame+1, nframes))
	createVideoFromFrames(dirpath, args.vid_out, nframes)
# save frame as JPEG file #caffe.set_device(args.gpu) # Select desired model # get input shape # get output shape # load cluster centers # populate cluster centers as 1x1 convolution kernel # print 'Annealed-Mean Parameters populated' # load the original image # convert image to lab color space # pull out L channel # original image size # create grayscale version of image (just for displaying) # resize image to network input size # resize image to network input size # subtract 50 for mean-centering # run network # this is our result # upsample to match size of original image L # concatenate with original image L # convert back to rgb