| """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() |
| 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) |
|
|
