240 lines
7.0 KiB
Python
240 lines
7.0 KiB
Python
import logging
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
from tqdm import tqdm
|
|
from time import perf_counter
|
|
from decimal import Decimal
|
|
|
|
from interpolator import get_device
|
|
from interpolator import ImageInterpolator
|
|
from interpolator import ModelRunner, Anchor
|
|
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
|
|
)
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
def move_images(src_dir: str, interpolated_dir: str, output_dir: str):
|
|
src_dir = Path(src_dir)
|
|
interpolated_dir = Path(interpolated_dir)
|
|
output_dir = Path(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
index = 0
|
|
src_frames = sorted(src_dir.glob("img_*.png"))
|
|
interp_frames = sorted(interpolated_dir.glob("img_*.png"))
|
|
for i in range(len(src_frames)):
|
|
output_frame = output_dir / f"img_{index:08d}.png"
|
|
src_frames[i].rename(output_frame)
|
|
index += 1
|
|
|
|
if i < len(interp_frames):
|
|
output_interp = output_dir / f"img_{index:08d}.png"
|
|
interp_frames[i].rename(output_interp)
|
|
index += 1
|
|
|
|
|
|
def build_file_list(moved_dir: str, list_path: str):
|
|
import os
|
|
moved_dir = Path(moved_dir)
|
|
frames = sorted(moved_dir.glob("img_*.png"))
|
|
print(frames[0])
|
|
|
|
with open(list_path, "w") as f:
|
|
for frame in frames:
|
|
f.write(f"file '{os.path.abspath(frame)}'\n")
|
|
|
|
|
|
def build_ffmpeg_file_list(frames_dir: str, interpolated_dir: str, list_path: str):
|
|
frames = sorted(Path(frames_dir).glob("img_*.png"))
|
|
interps = sorted(Path(interpolated_dir).glob("img_*.png"))
|
|
|
|
if len(interps) != len(frames) - 1:
|
|
raise ValueError("Interpolated frames must be N-1")
|
|
|
|
with open(list_path, "w") as f:
|
|
for i in range(len(frames)):
|
|
f.write(f"file '{frames[i].resolve().as_posix()}'\n")
|
|
|
|
if i < len(interps):
|
|
f.write(f"file '{interps[i].resolve().as_posix()}'\n")
|
|
|
|
|
|
def merge_with_ffmpeg(
|
|
original_video: str,
|
|
file_list: str,
|
|
output_video: str,
|
|
):
|
|
cap = cv2.VideoCapture(original_video)
|
|
|
|
if not cap.isOpened():
|
|
raise ValueError("Cannot open original video")
|
|
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
cap.release()
|
|
|
|
new_fps = Decimal(fps * 2)
|
|
|
|
cmd = [
|
|
"ffmpeg",
|
|
"-y",
|
|
"-r", str(new_fps.quantize(Decimal("1.0000000000"))),
|
|
"-f", "concat",
|
|
"-safe", "0",
|
|
"-i", file_list,
|
|
"-c:v", "libx264rgb",
|
|
output_video,
|
|
]
|
|
print("Running ffmpeg command:", " ".join(cmd))
|
|
|
|
subprocess.run(cmd, check=True)
|
|
|
|
|
|
|
|
def video_frames_to_disk_generator(
|
|
video_path: str | Path,
|
|
output_dir: str | Path,
|
|
chunk_seconds: int = 10
|
|
):
|
|
output_dir = Path(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
cap = cv2.VideoCapture(str(video_path))
|
|
|
|
if not cap.isOpened():
|
|
raise ValueError(f"Cannot open video: {video_path}")
|
|
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
frames_per_chunk = int(fps * chunk_seconds)
|
|
|
|
frame_index = 0
|
|
|
|
while True:
|
|
paths = []
|
|
|
|
for _ in range(frames_per_chunk):
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
cap.release()
|
|
return
|
|
|
|
frame_path = output_dir / f"img_{frame_index:08d}.png"
|
|
cv2.imwrite(str(frame_path), frame)
|
|
|
|
paths.append(frame_path)
|
|
frame_index += 1
|
|
|
|
yield tuple(paths)
|
|
|
|
|
|
def main():
|
|
start = perf_counter()
|
|
logging.info("Starting video interpolation process")
|
|
config_path = Path("src/config/AMT-G.yaml")
|
|
ckpt_path = Path("src/pretrained/amt-g.pth")
|
|
video_path = Path("example/video.mp4")
|
|
output_dir = Path("output/frames")
|
|
output_interpolated_dir = Path("output/interpolated")
|
|
output_interpolated_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
device = get_device()
|
|
model_runner = ModelRunner(config_path, ckpt_path, device)
|
|
if device.type in ("cpu", "mps"):
|
|
if device.type == "mps":
|
|
logging.warning(
|
|
"Running on Apple Silicon GPU (MPS) may have limited performance. Consider using a CUDA-enabled GPU for better performance."
|
|
)
|
|
else:
|
|
logging.warning(
|
|
"Running on CPU may be very slow. Consider using a GPU for better performance."
|
|
)
|
|
anchor = Anchor(resolution=8192 * 8192, memory=1, memory_bias=0)
|
|
elif device.type == "cuda":
|
|
anchor = Anchor(
|
|
resolution=1024 * 512, memory=1500 * 1024**2, memory_bias=2500 * 1024**2
|
|
)
|
|
else:
|
|
raise Exception(f"Unsupported device type: {device.type}")
|
|
|
|
interpolator = ImageInterpolator(device, anchor, model_runner)
|
|
|
|
loaded_time = perf_counter() - start
|
|
logging.info(f"Model loaded and initialized in {loaded_time:.2f} seconds")
|
|
|
|
prev_frame_path = None
|
|
frame_count = 0
|
|
for frame_paths in video_frames_to_disk_generator(video_path, output_dir):
|
|
logging.info(f"Processing frames: {len(frame_paths)}")
|
|
|
|
if prev_frame_path is not None:
|
|
img1 = prev_frame_path[-1]
|
|
img2 = frame_paths[0]
|
|
output_path = output_interpolated_dir / f"img_{frame_count:08d}.png"
|
|
interpolator.interpolate(img1, img2, output_path)
|
|
logging.debug(f"Interpolated image saved to: {output_path}")
|
|
frame_count += 1
|
|
for i in tqdm(range(len(frame_paths) - 1), desc="Interpolating frames"):
|
|
img1 = frame_paths[i]
|
|
img2 = frame_paths[i + 1]
|
|
output_path = output_interpolated_dir / f"img_{frame_count:08d}.png"
|
|
interpolator.interpolate(img1, img2, output_path)
|
|
logging.debug(f"Interpolated image saved to: {output_path}")
|
|
frame_count += 1
|
|
prev_frame_path = frame_paths
|
|
total_time = perf_counter() - start
|
|
logging.info(f"Video interpolation completed in {total_time:.2f} seconds")
|
|
|
|
|
|
def builder():
|
|
frames_dir = "output/frames"
|
|
interpolated_dir = "output/interpolated"
|
|
moved_dir = "output/moved"
|
|
video_path = "example/video.mp4"
|
|
output_video = "output/interpolated_video.mp4"
|
|
move_images(frames_dir, interpolated_dir, moved_dir)
|
|
|
|
cap = cv2.VideoCapture(video_path)
|
|
|
|
if not cap.isOpened():
|
|
raise ValueError("Cannot open original video")
|
|
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
cmd = [
|
|
"ffmpeg",
|
|
"-y",
|
|
"-framerate", str(fps * 2),
|
|
"-i", f"{moved_dir}/img_%08d.png",
|
|
"-i", video_path,
|
|
"-c:v", "libx264",
|
|
"-c:a", "copy",
|
|
"-shortest",
|
|
output_video,
|
|
]
|
|
logging.info("Running ffmpeg command to build final video: " + " ".join(cmd))
|
|
subprocess.run(cmd, check=True)
|
|
|
|
|
|
def cleanup():
|
|
import os
|
|
import shutil
|
|
frames_dir = "output/frames"
|
|
interpolated_dir = "output/interpolated"
|
|
moved_dir = "output/moved"
|
|
os.makedirs(frames_dir, exist_ok=True)
|
|
os.makedirs(interpolated_dir, exist_ok=True)
|
|
os.makedirs(moved_dir, exist_ok=True)
|
|
shutil.rmtree(frames_dir)
|
|
shutil.rmtree(interpolated_dir)
|
|
shutil.rmtree(moved_dir)
|
|
|
|
if __name__ == "__main__":
|
|
cleanup()
|
|
main()
|
|
builder()
|
|
cleanup()
|