Upload ComfyUI-ShotSplitter/first_frame_reader.py with huggingface_hub

ComfyUI-ShotSplitter/first_frame_reader.py

@@ -0,0 +1,42 @@
+"""Read the first frame of a list of clips into a unified IMAGE batch."""
+
+import logging
+from typing import List
+
+import numpy as np
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def read_first_frames(clip_paths: List[str]) -> torch.Tensor:
+    """Return (N, H, W, 3) float32 tensor in 0..1. Uses max-dim resize to first clip's size.
+
+    On read failure for any clip, substitute a black frame of the reference size.
+    """
+    import decord
+    decord.bridge.set_bridge("native")
+
+    frames: List[np.ndarray] = []
+    ref_h, ref_w = None, None
+    for path in clip_paths:
+        try:
+            vr = decord.VideoReader(path)
+            f = vr[0].asnumpy()  # HxWx3 uint8 RGB
+            if ref_h is None:
+                ref_h, ref_w = f.shape[0], f.shape[1]
+            if f.shape[0] != ref_h or f.shape[1] != ref_w:
+                import cv2
+                f = cv2.resize(f, (ref_w, ref_h), interpolation=cv2.INTER_AREA)
+            frames.append(f)
+        except Exception as e:
+            logger.warning("first-frame read failed for %s: %s", path, e)
+            if ref_h is None:
+                ref_h, ref_w = 64, 64
+            frames.append(np.zeros((ref_h, ref_w, 3), dtype=np.uint8))
+
+    if not frames:
+        return torch.zeros((0, 64, 64, 3), dtype=torch.float32)
+
+    stacked = np.stack(frames, axis=0).astype(np.float32) / 255.0
+    return torch.from_numpy(stacked)