new main

colorization/main.py

+import colorize
+import os
+import re
+import moviepy.editor as mp
+import cv2
+
+
+def subsample(input_file_path):
+    clip = mp.VideoFileClip(input_file_path)
+    clip_resized = clip.resize(height=360)
+    clip_resized.write_videofile(os.path.join('DATA/videosDS/', input_file_path.replace('DATA/videos/', '')))
+
+
+def video_to_frames(video_path):
+    # create a folder to store extracted images
+    folder_name = 'DATA/frames/' + video_path.replace('DATA/videosDS/', '').replace('.mp4', '/')
+    if not os.path.exists(folder_name):
+        os.makedirs(folder_name)
+
+    # extract frames
+    video_cap = cv2.VideoCapture(video_path)
+
+    fps = video_cap.get(cv2.CAP_PROP_FPS)
+    success, image = video_cap.read()
+    count = 0
+    while True:
+        success, image = video_cap.read()
+        if not success:
+            break
+        cv2.imwrite(os.path.join(folder_name, "{:d}.jpg".format(count)), image) # save frame as JPEG file
+        count += 1
+    print("{} frames are extracted in {}.".format(count, folder_name))
+    return fps
+
+
+def numerical_sort(value): # sort frames from dir
+    numbers = re.compile(r'(\d+)')
+    parts = numbers.split(value)
+    parts[1::2] = map(int, parts[1::2])
+    return parts
+
+
+def colorize_frames(greyscale_frames_folder):
+    colored_frames_folder = 'DATA/colored_frames/' + greyscale_frames_folder.replace('DATA/frames/', '')
+    if not os.path.exists(colored_frames_folder):
+        os.makedirs(colored_frames_folder)
+
+    input_images = [img for img in os.listdir(greyscale_frames_folder) if img.endswith(".jpg")]
+    input_images.sort(key=numerical_sort)
+
+    count = 0
+    for input_image in input_images:
+        colorize.main(os.path.join(greyscale_frames_folder, input_image),
+                      os.path.join(colored_frames_folder, "{:d}.jpg".format(count)))
+        count += 1
+
+
+def frames_to_video(frames_path, fps):
+    video_folder = 'DATA/colored_videos/' + frames_path.replace('DATA/colored_frames/', '') + '.avi'
+
+    images = [img for img in os.listdir(frames_path) if img.endswith(".jpg")]
+    images.sort(key=numerical_sort)
+    frame = cv2.imread(os.path.join(frames_path, images[0]))
+    height, width, layers = frame.shape
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # codec
+
+    video = cv2.VideoWriter(video_folder, fourcc, fps, (width, height))
+
+    for image in images:
+        video.write(cv2.imread(os.path.join(frames_path, image)))
+
+    cv2.destroyAllWindows()
+    video.release()
+
+
+def main():
+
+    input_file_path = "DATA/videos/chaplin2.mp4"   # preserve the path, change only the file name
+
+    subsample(input_file_path)
+    fps = video_to_frames(input_file_path.replace('videos', 'videosDS'))
+    colorize_frames('DATA/frames/' + input_file_path.replace('DATA/videos/', '').replace('.mp4', ''))
+    frames_to_video('DATA/colored_frames/' + input_file_path.replace('DATA/videos/', '').replace('.mp4', ''), fps)
+
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file