''' 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()