Spaces:
Running
Running
File size: 2,308 Bytes
1b38e97 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
'''
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()
|