Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
macacon python list
import tensorflow as tf
import numpy as np
import cv2
import argparse
import time
start_time = time.time()
parser = argparse.ArgumentParser()
parser.add_argument("inputMovie", type=file)
parser.add_argument("outputMovie")
parser.add_argument("-m", const=1, nargs="?")
parser.add_argument("-c", default=0, type=float, nargs="?")
args = parser.parse_args()
inputMovie = args.inputMovie
outputMovie = args.outputMovie
cap = cv2.VideoCapture(inputMovie.name)
heightInputMovie = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
widthInputMovie = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
FPS = cap.get(cv2.CAP_PROP_FPS)
allFrames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
print ("Width =", widthInputMovie)
print ("Height =", heightInputMovie)
print ("FPS =", FPS)
print ("Frames =", allFrames)
fourcc = "XVID"
out = cv2.VideoWriter(outputMovie, cv2.VideoWriter_fourcc(*fourcc), FPS, (widthInputMovie,heightInputMovie))
with open("colorize.tfmodel", mode='rb') as f:
fileContent = f.read()
graph_def = tf.GraphDef()
graph_def.ParseFromString(fileContent)
grayscale = tf.placeholder("float", [1, 224, 224, 1])
tf.import_graph_def(graph_def, input_map={ "grayscale": grayscale }, name='')
with tf.Session() as sess:
for frameNo in range(1, allFrames+1):
print ("Frame no =", frameNo)
ret, frame = cap.read()
if ret == True:
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame_rgb = (frame_rgb / 255.).astype(np.float32)
frame_lab = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2Lab)
frame_lab_l = frame_lab[:,:,0]
frame_rgb_vgg = cv2.resize(frame_rgb, (224, 224))
frame_rgb_vgg_gray = cv2.cvtColor(frame_rgb_vgg, cv2.COLOR_RGB2GRAY).reshape(1, 224, 224, 1)
inferred_rgb = sess.graph.get_tensor_by_name("inferred_rgb:0")
inferred_batch = sess.run(inferred_rgb, feed_dict={ grayscale: frame_rgb_vgg_gray })
foward_rgb = cv2.resize(inferred_batch[0], (widthInputMovie, heightInputMovie), interpolation=cv2.INTER_CUBIC)
frame_lab = cv2.cvtColor(foward_rgb, cv2.COLOR_RGB2Lab)
if args.c:
frame_lab *= args.c
frame_lab_out = np.concatenate((frame_lab_l[:,:,np.newaxis], frame_lab[:,:,1,np.newaxis], frame_lab[:,:,2,np.newaxis]), axis=2)
frame_rgb_out = cv2.cvtColor(frame_lab_out, cv2.COLOR_Lab2RGB)
frame_rgb_out = (frame_rgb_out * 255).astype(np.uint8)
cvFrame = cv2.cvtColor(frame_rgb_out, cv2.COLOR_RGB2BGR)
out.write(cvFrame)
if args.m:
cv2.imshow('frame', cvFrame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
out.release()
cv2.destroyAllWindows()
print ("Elapsed time:{0:.2f}".format(time.time() - start_time)) + "sec"
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment