Cutter works? 0.2

cc_scene_change.py

@@ -0,0 +1,151 @@
+import subprocess
+import numpy as np
+import cv2
+import torch
+import torch.nn.functional as F
+from scipy import signal
+
+TARGET_FPS = 30
+
+def get_device():
+
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+
+        print("GPU detected")
+        print("Device:", torch.cuda.get_device_name(0))
+        print("ROCm build:", torch.version.hip is not None)
+
+        return device
+
+    print("GPU not available, falling back to CPU")
+    return torch.device("cpu")
+
+
+device = get_device()
+
+
+def ffmpeg_frame_stream(path):
+
+    cmd = [
+        "ffmpeg",
+        "-i", path,
+        "-vf", "fps=30,scale=64:36,format=bgr24",
+        "-f", "rawvideo",
+        "-pix_fmt", "bgr24",
+        "-"
+    ]
+
+    proc = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        bufsize=10**8
+    )
+
+    frame_size = 64 * 36 * 3
+
+    while True:
+        raw = proc.stdout.read(frame_size)
+        if len(raw) < frame_size:
+            break
+
+        frame = np.frombuffer(raw, np.uint8).reshape((36, 64, 3))
+        yield frame
+
+    proc.stdout.close()
+    proc.wait()
+
+
+def compute_signal(video_path):
+
+    prev_blocks = None
+    signal_out = []
+
+    for frame in ffmpeg_frame_stream(video_path):
+
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+
+        t = torch.from_numpy(gray).float().to(device) / 255.0
+
+        # reshape for conv
+        t = t.unsqueeze(0).unsqueeze(0)
+
+        # block averages using avg pool
+        blocks = F.avg_pool2d(t, kernel_size=8, stride=8)
+
+        blocks = blocks.squeeze()
+
+        if prev_blocks is None:
+            prev_blocks = blocks
+            signal_out.append(0)
+            continue
+
+        # block difference energy
+        diff = torch.abs(blocks - prev_blocks).sum()
+
+        # edge energy
+        sobel_x = torch.tensor(
+            [[-1,0,1],[-2,0,2],[-1,0,1]],
+            dtype=torch.float32, device=device
+        ).view(1,1,3,3)
+
+        sobel_y = torch.tensor(
+            [[-1,-2,-1],[0,0,0],[1,2,1]],
+            dtype=torch.float32, device=device
+        ).view(1,1,3,3)
+
+        gx = F.conv2d(t, sobel_x, padding=1)
+        gy = F.conv2d(t, sobel_y, padding=1)
+
+        edge_energy = torch.mean(torch.sqrt(gx**2 + gy**2))
+
+        value = diff + 0.5 * edge_energy
+
+        signal_out.append(value.item())
+
+        prev_blocks = blocks
+
+    signal_out = np.array(signal_out)
+
+    # temporal derivative
+    signal_out = np.diff(signal_out, prepend=signal_out[0])
+
+    return signal_out
+
+
+def normalize(sig):
+    return (sig - np.mean(sig)) / (np.std(sig) + 1e-6)
+
+
+def find_offset(video_a, video_b):
+
+    sig_a = compute_signal(video_a)
+    sig_b = compute_signal(video_b)
+
+    sig_a = normalize(sig_a)
+    sig_b = normalize(sig_b)
+
+    corr = signal.correlate(sig_a, sig_b, mode="full")
+    lags = signal.correlation_lags(len(sig_a), len(sig_b), mode="full")
+
+    best = lags[np.argmax(corr)]
+
+    seconds = best / TARGET_FPS
+
+    return best, seconds
+
+
+def main():
+
+    video_a = "/home/honney/mount/Media/tmp/serien/ger_low_quality/Black Bullet/Black Bullet Season 01 E03 Kinder des Schicksals.mp4"
+    video_b = "/home/honney/mount/Media/tmp/serien/high_quality/Black Bullet (Season 1)/Season 01/Black Bullet - S01E03.mkv"
+
+    frames, seconds = find_offset(video_a, video_b)
+
+    print("offset frames:", frames)
+    print("offset seconds:", seconds)
+
+
+if __name__ == "__main__":
+    main()