Upload colorize_video.py

colorize_video.py

@@ -0,0 +1,74 @@
+'''
+python colorize_video.py xiangling_skt.mp4 xiangling_clr.mp4 --model ECCV16
+'''
+
+import os
+import tempfile
+import time
+import argparse
+
+import cv2
+import moviepy.editor as mp
+import numpy as np
+from tqdm import tqdm
+
+from models.deep_colorization.colorizers import eccv16
+from utils import colorize_frame, change_model
+
+def colorize_video(input_path, output_path, model_name):
+    # Load the model
+    loaded_model = change_model("None", model_name)
+
+    # Open the video using cv2.VideoCapture
+    video = cv2.VideoCapture(input_path)
+
+    # Get video information
+    fps = video.get(cv2.CAP_PROP_FPS)
+    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    # Colorize video frames and store in a list
+    output_frames = []
+
+    for _ in tqdm(range(total_frames), unit='frame', desc="Progress"):
+        ret, frame = video.read()
+        if not ret:
+            break
+
+        colorized_frame = colorize_frame(frame, loaded_model)
+        output_frames.append((colorized_frame * 255).astype(np.uint8))
+
+    # Merge frames to video
+    frame_size = output_frames[0].shape[:2]
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Codec for MP4 video
+    out = cv2.VideoWriter(output_path, fourcc, fps, (frame_size[1], frame_size[0]))
+
+    for frame in output_frames:
+        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+        out.write(frame_bgr)
+
+    out.release()
+
+    # Close the video capture
+    video.release()
+    
+    # Convert the output video to a format compatible with Streamlit
+    clip = mp.VideoFileClip(output_path)
+
+    # Write the video file
+    #clip.write_videofile(output_path, codec="libx264", audio=audio is not None)
+    clip.write_videofile(output_path)
+    
+    print(f"Colorized video saved to {output_path}")
+
+def main():
+    parser = argparse.ArgumentParser(description="Colorize a black and white video.")
+    parser.add_argument("input_path", type=str, help="Path to the input video file.")
+    parser.add_argument("output_path", type=str, help="Path to save the colorized video file.")
+    parser.add_argument("--model", type=str, default="ECCV16", choices=["ECCV16", "SIGGRAPH17"], help="Model to use for colorization.")
+
+    args = parser.parse_args()
+
+    colorize_video(args.input_path, args.output_path, args.model)
+
+if __name__ == "__main__":
+    main()