File size: 1,400 Bytes
df67f77 | 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 | """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)
|