Spaces:
Running on Zero
Running on Zero
File size: 8,701 Bytes
ee69ab7 | 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | 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)
|