Filters/video_processor.py at main · AJ0617/Filters · GitHub

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
76
77
78
79
80
81
82
83
84
85
import os
import subprocess
import tempfile

import cv2
import imageio_ffmpeg

from hand_pipeline import create_detector, get_hand_points
from hand_tracking import INDEX_TIP, THUMB_TIP
from geometry import render_portal


def _mux_audio(silent_video_path, source_path, output_path):
    ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
    cmd = [
        ffmpeg_exe, "-y",
        "-i", silent_video_path,
        "-i", source_path,
        "-map", "0:v:0",
        "-map", "1:a:0?",
        "-c:v", "libx264",
        "-pix_fmt", "yuv420p",
        "-preset", "fast",
        "-crf", "23",
        "-c:a", "aac",
        "-shortest",
        output_path,
    ]
    result = subprocess.run(cmd, capture_output=True)
    return result.returncode == 0


def process_video(input_path, output_path, filter_func, progress_callback=None):
    detector = create_detector()

    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        raise RuntimeError(f"Could not open the video: {input_path}")

    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    silent_fd, silent_path = tempfile.mkstemp(suffix=".mp4")
    os.close(silent_fd)

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(silent_path, fourcc, fps, (w, h))
    if not writer.isOpened():
        cap.release()
        os.remove(silent_path)
        raise RuntimeError(f"Could not create the output video: {output_path}")

    frame_idx = 0
    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                break

            timestamp_ms = int((frame_idx / fps) * 1000)
            left_hand, right_hand = get_hand_points(detector, frame, timestamp_ms)

            if left_hand is not None and right_hand is not None:
                p1 = (left_hand[INDEX_TIP].x * w, left_hand[INDEX_TIP].y * h)
                p2 = (left_hand[THUMB_TIP].x * w, left_hand[THUMB_TIP].y * h)
                p3 = (right_hand[INDEX_TIP].x * w, right_hand[INDEX_TIP].y * h)
                p4 = (right_hand[THUMB_TIP].x * w, right_hand[THUMB_TIP].y * h)
                frame = render_portal(frame, p1, p2, p3, p4, filter_func)

            writer.write(frame)
            frame_idx += 1
            if progress_callback:
                progress_callback(frame_idx, total_frames)
    finally:
        cap.release()
        writer.release()

    if _mux_audio(silent_path, input_path, output_path):
        os.remove(silent_path)
    else:
        os.replace(silent_path, output_path)

    return output_path