from __future__ import annotations from pathlib import Path import cv2 import numpy as np import torch import torch.nn.functional as F from PIL import Image from inference.request import InferenceRequest def _resize_tensor_video(value: torch.Tensor, height: int, width: int) -> torch.Tensor: """Center-crop and resize a [B,C,F,H,W] tensor in its existing range.""" if height <= 0 or width <= 0: raise ValueError("target height and width must be positive") source_height, source_width = value.shape[-2:] if (source_height, source_width) == (height, width): return value source_aspect = source_height / source_width target_aspect = height / width if source_aspect < target_aspect: crop_width = max(1, int(source_height / height * width)) left = max(0, (source_width - crop_width) // 2) value = value[..., :, left:left + crop_width] else: crop_height = max(1, int(source_width / width * height)) top = max(0, (source_height - crop_height) // 2) value = value[..., top:top + crop_height, :] batch, channels, frames = value.shape[:3] flat = value.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, *value.shape[-2:]) flat = F.interpolate(flat, size=(height, width), mode="bilinear", align_corners=False) return flat.reshape(batch, frames, channels, height, width).permute(0, 2, 1, 3, 4) def pil_center_crop_resize( frames: np.ndarray, target_height: int, target_width: int, *, resample: Image.Resampling = Image.Resampling.BILINEAR, ) -> np.ndarray: """Center crop to the requested aspect ratio and resize with PIL. This mirrors the geometry used by the original full-video inference path. OpenCV is only responsible for decoding; all spatial image operations are deliberately performed by PIL. """ if frames.ndim != 4 or frames.shape[-1] != 3: raise ValueError(f"Expected [F,H,W,3] RGB frames, got {frames.shape}") if target_height <= 0 or target_width <= 0: raise ValueError("target height and width must be positive") output = [] for frame in frames: image = Image.fromarray(frame, mode="RGB") source_width, source_height = image.size if source_width == target_width and source_height == target_height: output.append(np.asarray(image)) continue source_aspect = source_height / source_width target_aspect = target_height / target_width if source_aspect < target_aspect: crop_width = int(source_height / target_height * target_width) left = (source_width - crop_width) // 2 image = image.crop((left, 0, left + crop_width, source_height)) else: crop_height = int(source_width / target_width * target_height) top = (source_height - crop_height) // 2 image = image.crop((0, top, source_width, top + crop_height)) image = image.resize((target_width, target_height), resample=resample) output.append(np.asarray(image)) return np.ascontiguousarray(np.stack(output)) def load_lq_video( request: InferenceRequest, *, stream: bool = False, ) -> tuple[torch.Tensor, float]: request.validate() if request.lq_video is not None: value = request.lq_video if value.shape[1] != 3 and value.shape[2] == 3: value = value.permute(0, 2, 1, 3, 4) if value.shape[1] != 3: raise ValueError("lq_video channel dimension must be 3") if request.height is not None: value = _resize_tensor_video(value.float(), request.height, request.width) if request.max_frames is not None: value = value[:, :, : request.max_frames] return value.float(), float(request.fps or 24) capture = cv2.VideoCapture(str(Path(request.input_path).expanduser())) if not capture.isOpened(): raise ValueError(f"Could not open input video: {request.input_path}") fps = float(request.fps or capture.get(cv2.CAP_PROP_FPS) or 24) frames = [] try: while request.max_frames is None or len(frames) < request.max_frames: ok, frame = capture.read() if not ok: break frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) finally: capture.release() if not frames: raise ValueError(f"No frames decoded from: {request.input_path}") array = np.stack(frames) if request.height is not None or request.width is not None: if request.height is None or request.width is None: raise ValueError("height and width must be provided together") array = pil_center_crop_resize( array, request.height, request.width, resample=(Image.Resampling.BICUBIC if stream else Image.Resampling.BILINEAR), ) tensor = torch.from_numpy(np.ascontiguousarray(array)).permute(0, 3, 1, 2).float() tensor = tensor.div(127.5).sub(1.0).permute(1, 0, 2, 3).unsqueeze(0) return tensor, fps def get_lq_video_fps(request: InferenceRequest) -> float: request.validate() if request.lq_video is not None: return float(request.fps or 24) capture = cv2.VideoCapture(str(Path(request.input_path).expanduser())) try: if not capture.isOpened(): raise ValueError(f"Could not open input video: {request.input_path}") return float(request.fps or capture.get(cv2.CAP_PROP_FPS) or 24) finally: capture.release() def iter_lq_video( request: InferenceRequest, *, chunk_frames: int, stream: bool = True, first_chunk_frames: int | None = None, ): """Yield bounded ``[1,3,F,H,W]`` CPU windows from a video file. Windows are independent ownership units: callers may release each tensor before requesting the next one. The final short window is yielded as-is; temporal padding remains the runner's responsibility. ``first_chunk_frames`` allows the unique first frame plus one full look-ahead region to be read before the first decode segment. """ request.validate() if chunk_frames <= 0: raise ValueError("chunk_frames must be positive") if first_chunk_frames is not None and first_chunk_frames <= 0: raise ValueError("first_chunk_frames must be positive") window_size = first_chunk_frames or chunk_frames if request.lq_video is not None: value = request.lq_video if value.shape[1] != 3 and value.shape[2] == 3: value = value.permute(0, 2, 1, 3, 4) if value.shape[1] != 3: raise ValueError("lq_video channel dimension must be 3") value = value[:, :, :request.max_frames] if request.max_frames else value if not value.shape[2]: raise ValueError("No frames in lq_video") start = 0 while start < value.shape[2]: clip = value[:, :, start:start + window_size].float() if request.height is not None: clip = _resize_tensor_video(clip, request.height, request.width) yield clip start += window_size window_size = chunk_frames return capture = cv2.VideoCapture(str(Path(request.input_path).expanduser())) if not capture.isOpened(): raise ValueError(f"Could not open input video: {request.input_path}") count = 0 batch = [] try: while request.max_frames is None or count < request.max_frames: ok, frame = capture.read() if not ok: break count += 1 batch.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) if len(batch) == window_size: value = _frames_to_tensor(batch, request, stream=stream) batch.clear() yield value del value window_size = chunk_frames if batch: yield _frames_to_tensor(batch, request, stream=stream) if not count: raise ValueError(f"No frames decoded from: {request.input_path}") finally: capture.release() def _frames_to_tensor(frames, request, *, stream): array = np.stack(frames) if request.height is not None: if request.width is None: raise ValueError("height and width must be provided together") array = pil_center_crop_resize( array, request.height, request.width, resample=Image.Resampling.BICUBIC if stream else Image.Resampling.BILINEAR, ) return torch.from_numpy(np.ascontiguousarray(array)).permute(0, 3, 1, 2).float().div_(127.5).sub_(1.0).permute(1, 0, 2, 3).unsqueeze(0)