diff --git a/data/vg-30.pb b/data/vg-30.pb new file mode 100644 index 0000000..2f57642 Binary files /dev/null and b/data/vg-30.pb differ diff --git a/input/jaguar.mp4 b/input/jaguar.mp4 new file mode 100644 index 0000000..89c6806 Binary files /dev/null and b/input/jaguar.mp4 differ diff --git a/main.py b/main.py new file mode 100644 index 0000000..0643e53 --- /dev/null +++ b/main.py @@ -0,0 +1,74 @@ +import logging +import sys +from datetime import datetime + +import tensorflow as tf +import os +os.environ["TF_CPP_MIN_LOG_LEVEL"]="3" + +''' +1. import `Video` and `save_video` from the correct module of package "styler" +''' +from styler.video import Video +from styler.utils import save_video + + +model_file = 'data/vg-30.pb' +model_name = 'vg-30' +logging.basicConfig( + stream=sys.stdout, + format='%(asctime)s %(levelname)s:%(message)s', + level=logging.INFO, + datefmt='%I:%M:%S') + + +def main(): + + with open(model_file, 'rb') as f: + graph_def = tf.GraphDef() + graph_def.ParseFromString(f.read()) + tf.import_graph_def(graph_def) + graph = tf.get_default_graph() + + with tf.Session(config=tf.ConfigProto( + intra_op_parallelism_threads=4)) as session: + + logging.info("Initializing graph") + session.run(tf.global_variables_initializer()) + + image = graph.get_tensor_by_name("import/%s/image_in:0" % model_name) + out = graph.get_tensor_by_name("import/%s/output:0" % model_name) + shape = image.get_shape().as_list() + + ''' + 2. set the `path` to your input + ''' + with Video('input/jaguar.mp4') as v: + frames = v.read_frames(image_h=shape[1], image_w=shape[2]) + + logging.info("Processing image") + start_time = datetime.now() + + ''' + 3. Write a list comprehension to iterate through all frames, + and make it be processed by Tensorflow. + ''' + processed = [ + session.run(out, feed_dict={image: [frame]}) + for frame in frames + ] + + ''' + 4. Pass the results as a argument into function + ''' + save_video('result.avi', + fps=30, h=shape[1], w=shape[2], + frames=processed) + + logging.info("Processing took %f" % ( + (datetime.now() - start_time).total_seconds())) + logging.info("Done") + + +if __name__ == '__main__': + main() diff --git a/result.avi b/result.avi new file mode 100644 index 0000000..e8f3116 Binary files /dev/null and b/result.avi differ diff --git a/styler/__pycache__/utils.cpython-36.pyc b/styler/__pycache__/utils.cpython-36.pyc new file mode 100644 index 0000000..054778f Binary files /dev/null and b/styler/__pycache__/utils.cpython-36.pyc differ diff --git a/styler/__pycache__/video.cpython-36.pyc b/styler/__pycache__/video.cpython-36.pyc new file mode 100644 index 0000000..caf714f Binary files /dev/null and b/styler/__pycache__/video.cpython-36.pyc differ diff --git a/styler/utils.py b/styler/utils.py new file mode 100644 index 0000000..c4dcb5c --- /dev/null +++ b/styler/utils.py @@ -0,0 +1,48 @@ +import cv2 +import numpy as np + +MEAN_VALUES = np.array([123, 117, 104]).reshape((1,1,3)) + + +def resize(img, image_h, image_w, zoom=False): + # crop image from center + ratio = float(image_h) / image_w + height, width = int(img.shape[0]), int(img.shape[1]) + yy, xx = 0, 0 + if height > width * ratio: #too tall + yy = int(height - width * ratio) // 2 + height = int(width * ratio) + else: # too wide + xx = int(width - height / ratio) // 2 + width = int(height / ratio) + if zoom: + yy += int(height / 6) + xx += int(height / 6) + height = int(height * 2 / 3) + width = int(width * 2 / 3) + crop_img = img[yy:yy + height, xx:xx + width] + # resize + resized_img = cv2.resize(crop_img, (image_h, image_w)) + centered_img = resized_img - MEAN_VALUES + + return centered_img + + +def save_video(filepath, fps, w, h, frames): + codecs = ['WMV1', 'MJPG', 'XVID', 'PIM1'] + ''' + If you cannot write video file, you may change the used codec + ''' + used_codec = codecs[2] # change the index from codecs + fourcc = cv2.VideoWriter_fourcc(*used_codec) + out = cv2.VideoWriter(filepath, fourcc, fps, (w, h)) + for frame in frames: + f = frame[0, :, :, :] + out.write(post_process(f)) + out.release() + + +def post_process(image): + img = image + MEAN_VALUES + img = np.clip(img, 0, 255).astype('uint8') + return cv2.cvtColor(img, cv2.COLOR_RGB2BGR) diff --git a/styler/video.py b/styler/video.py new file mode 100644 index 0000000..2497dfe --- /dev/null +++ b/styler/video.py @@ -0,0 +1,47 @@ +import cv2 +import time + +from styler.utils import resize + + +class Video: + + def __init__(self, path): + self.path = path + self.cap = cv2.VideoCapture(self.path) + self.frames = [] + + def __enter__(self): + if not self.cap.isOpened(): + raise Exception('Cannot open video: {}'.format(self.path)) + return self + + def __len__(self): + return len(self.frames) + + def read_frames(self, image_h, image_w): + ''' + 5. + - Read video frames from `self.cap` and collect frames into list + - Apply `resize()` on each frame before add it to list + - Also assign frames to "self" object + - Return your results + ''' + frames = [] + while self.cap.isOpened(): + ret, cap_frame = self.cap.read() + if ret == True: + out = resize(cap_frame, image_h=int(image_h), image_w=int(image_w)) + frames.append(out) + # 5-1 /5-2 Read video and collect them + + # Break the loop + else: + break + + print(len(frames)) + self.frames = frames # 5-3 let object have the result + return frames # return your results + + def __exit__(self, exc_type, exc_val, exc_tb): + self.cap.release()