Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2025 The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16import os17import warnings18from collections.abc import Iterable, Mapping19from contextlib import redirect_stdout20from dataclasses import dataclass, fields21from io import BytesIO22from typing import Callable, NewType, Optional, Union23from urllib.parse import urlparse24 25import numpy as np26import requests27 28from .image_transforms import PaddingMode, to_channel_dimension_format29from .image_utils import ChannelDimension, infer_channel_dimension_format, is_valid_image30from .utils import (31 is_av_available,32 is_cv2_available,33 is_decord_available,34 is_numpy_array,35 is_torch_available,36 is_torch_tensor,37 is_torchcodec_available,38 is_torchvision_available,39 is_vision_available,40 is_yt_dlp_available,41 logging,42 requires_backends,43)44 45 46if is_vision_available():47 import PIL.Image48 import PIL.ImageOps49 50 if is_torchvision_available():51 from torchvision import io as torchvision_io52 53if is_torch_available():54 import torch55 56 57logger = logging.get_logger(__name__)58 59URL = NewType("URL", str)60Path = NewType("Path", str)61 62VideoInput = Union[63 list["PIL.Image.Image"],64 np.ndarray,65 "torch.Tensor",66 list[np.ndarray],67 list["torch.Tensor"],68 list[list["PIL.Image.Image"]],69 list[list[np.ndarray]],70 list[list["torch.Tensor"]],71 URL,72 list[URL],73 list[list[URL]],74 Path,75 list[Path],76 list[list[Path]],77]78 79 80@dataclass81class VideoMetadata(Mapping):82 total_num_frames: int83 fps: Optional[float] = None84 width: Optional[int] = None85 height: Optional[int] = None86 duration: Optional[float] = None87 video_backend: Optional[str] = None88 frames_indices: Optional[list[int]] = None89 90 def __iter__(self):91 return (f.name for f in fields(self))92 93 def __len__(self):94 return len(fields(self))95 96 def __getitem__(self, item):97 return getattr(self, item)98 99 def __setitem__(self, key, value):100 return setattr(self, key, value)101 102 @property103 def timestamps(self) -> list[float]:104 "Timestamps of the sampled frames in seconds."105 if self.fps is None or self.frames_indices is None:106 raise ValueError("Cannot infer video `timestamps` when `fps` or `frames_indices` is None.")107 return [frame_idx / self.fps for frame_idx in self.frames_indices]108 109 def update(self, dictionary):110 for key, value in dictionary.items():111 if hasattr(self, key):112 setattr(self, key, value)113 114 115def is_valid_video_frame(frame):116 return isinstance(frame, PIL.Image.Image) or (117 (is_numpy_array(frame) or is_torch_tensor(frame)) and frame.ndim == 3118 )119 120 121def is_valid_video(video):122 if not isinstance(video, (list, tuple)):123 return (is_numpy_array(video) or is_torch_tensor(video)) and video.ndim == 4124 return video and all(is_valid_video_frame(frame) for frame in video)125 126 127def valid_videos(videos):128 # If we have a list of videos, it could be either one video as list of frames or a batch129 if isinstance(videos, (list, tuple)):130 for video_or_frame in videos:131 if not (is_valid_video(video_or_frame) or is_valid_video_frame(video_or_frame)):132 return False133 # If not a list, then we have a single 4D video or 5D batched tensor134 elif not is_valid_video(videos) or videos.ndim == 5:135 return False136 return True137 138 139def is_batched_video(videos):140 if isinstance(videos, (list, tuple)):141 return is_valid_video(videos[0])142 elif (is_numpy_array(videos) or is_torch_tensor(videos)) and videos.ndim == 5:143 return True144 return False145 146 147def is_scaled_video(video: np.ndarray) -> bool:148 """149 Checks to see whether the pixel values have already been rescaled to [0, 1].150 """151 # It's possible the video has pixel values in [0, 255] but is of floating type152 return np.min(video) >= 0 and np.max(video) <= 1153 154 155def convert_pil_frames_to_video(videos: list[VideoInput]) -> list[Union[np.ndarray, "torch.Tensor"]]:156 """157 Given a batch of videos, converts each video to a 4D array. If video is already in array type,158 it is simply returned. We assume that all inputs in the list are in the same format, based on the type of the first element.159 160 Args:161 videos (`VideoInput`):162 Video inputs to turn into a list of videos.163 """164 165 if not (isinstance(videos[0], (list, tuple)) and is_valid_image(videos[0][0])):166 return videos167 168 video_converted = []169 for video in videos:170 video = [np.array(frame) for frame in video]171 video = np.stack(video)172 video_converted.append(video)173 return video_converted174 175 176def make_batched_videos(videos) -> list[Union[np.ndarray, "torch.Tensor", "URL", "Path"]]:177 """178 Ensure that the input is a list of videos. If the input is a single video, it is converted to a list of length 1.179 If the input is a batch of videos, it is converted to a list of 4D video arrays. Videos passed as list `PIL.Image`180 frames are converted to 4D arrays.181 182 We assume that all inputs in the list are in the same format, based on the type of the first element.183 184 Args:185 videos (`VideoInput`):186 Video inputs to turn into a list of videos.187 """188 # Early exit for deeply nested list of image frame paths. We shouldn't flatten them189 try:190 if isinstance(videos[0][0], list) and isinstance(videos[0][0][0], str):191 return [image_paths for sublist in videos for image_paths in sublist]192 except (IndexError, TypeError):193 pass194 195 if isinstance(videos, str) or is_valid_video(videos):196 return convert_pil_frames_to_video([videos])197 # only one frame passed, thus we unsqueeze time dim198 elif is_valid_image(videos):199 if isinstance(videos, PIL.Image.Image):200 videos = np.array(videos)201 return [videos[None, ...]]202 elif not isinstance(videos, list):203 raise ValueError(204 f"Invalid video input. Expected either a list of video frames or an input of 4 or 5 dimensions, but got"205 f" type {type(videos)}."206 )207 208 # Recursively flatten any nested structure209 flat_videos_list = []210 for item in videos:211 if isinstance(item, str) or is_valid_video(item):212 flat_videos_list.append(item)213 elif isinstance(item, list) and item:214 flat_videos_list.extend(make_batched_videos(item))215 216 flat_videos_list = convert_pil_frames_to_video(flat_videos_list)217 return flat_videos_list218 219 220def make_batched_metadata(videos: VideoInput, video_metadata: Union[VideoMetadata, dict]):221 if video_metadata is None:222 # Create default metadata and fill attributes we can infer from given video223 video_metadata = [224 {225 "total_num_frames": len(video),226 "fps": None,227 "duration": None,228 "frames_indices": list(range(len(video))),229 "height": get_video_size(video)[0] if is_valid_video(video) else None,230 "width": get_video_size(video)[1] if is_valid_video(video) else None,231 }232 for video in videos233 ]234 235 if isinstance(video_metadata, list):236 # Flatten if nested list237 if isinstance(video_metadata[0], list):238 video_metadata = [239 VideoMetadata(**metadata) for metadata_list in video_metadata for metadata in metadata_list240 ]241 # Simply wrap in VideoMetadata if simple dict242 elif isinstance(video_metadata[0], dict):243 video_metadata = [VideoMetadata(**metadata) for metadata in video_metadata]244 else:245 # Create a batched list from single object246 video_metadata = [VideoMetadata(**video_metadata)]247 return video_metadata248 249 250def get_video_size(video: np.ndarray, channel_dim: Optional[ChannelDimension] = None) -> tuple[int, int]:251 """252 Returns the (height, width) dimensions of the video.253 254 Args:255 video (`np.ndarray`):256 The video to get the dimensions of.257 channel_dim (`ChannelDimension`, *optional*):258 Which dimension the channel dimension is in. If `None`, will infer the channel dimension from the video.259 260 Returns:261 A tuple of the video's height and width.262 """263 if channel_dim is None:264 channel_dim = infer_channel_dimension_format(video, num_channels=(1, 3, 4))265 266 if channel_dim == ChannelDimension.FIRST:267 return video.shape[-2], video.shape[-1]268 elif channel_dim == ChannelDimension.LAST:269 return video.shape[-3], video.shape[-2]270 else:271 raise ValueError(f"Unsupported data format: {channel_dim}")272 273 274def get_uniform_frame_indices(total_num_frames: int, num_frames: Optional[int] = None):275 """276 Creates a numpy array for uniform sampling of `num_frame` frames from `total_num_frames`277 when loading a video.278 279 Args:280 total_num_frames (`int`):281 Total number of frames that a video has.282 num_frames (`int`, *optional*):283 Number of frames to sample uniformly. If not specified, all frames are sampled.284 285 Returns:286 np.ndarray: np array of frame indices that will be sampled.287 """288 if num_frames is not None:289 indices = np.arange(0, total_num_frames, total_num_frames / num_frames).astype(int)290 else:291 indices = np.arange(0, total_num_frames).astype(int)292 return indices293 294 295def default_sample_indices_fn(metadata: VideoMetadata, num_frames=None, fps=None, **kwargs):296 """297 A default sampling function that replicates the logic used in get_uniform_frame_indices,298 while optionally handling `fps` if `num_frames` is not provided.299 300 Args:301 metadata (`VideoMetadata`):302 `VideoMetadata` object containing metadata about the video, such as "total_num_frames" or "fps".303 num_frames (`int`, *optional*):304 Number of frames to sample uniformly.305 fps (`int` or `float`, *optional*):306 Desired frames per second. Takes priority over num_frames if both are provided.307 308 Returns:309 `np.ndarray`: Array of frame indices to sample.310 """311 total_num_frames = metadata.total_num_frames312 video_fps = metadata.fps313 314 # If num_frames is not given but fps is, calculate num_frames from fps315 if num_frames is None and fps is not None:316 num_frames = int(total_num_frames / video_fps * fps)317 if num_frames > total_num_frames:318 raise ValueError(319 f"When loading the video with fps={fps}, we computed num_frames={num_frames} "320 f"which exceeds total_num_frames={total_num_frames}. Check fps or video metadata."321 )322 323 if num_frames is not None:324 indices = np.arange(0, total_num_frames, total_num_frames / num_frames, dtype=int)325 else:326 indices = np.arange(0, total_num_frames, dtype=int)327 return indices328 329 330def read_video_opencv(331 video_path: Union["URL", "Path"],332 sample_indices_fn: Callable,333 **kwargs,334) -> tuple[np.ndarray, VideoMetadata]:335 """336 Decode a video using the OpenCV backend.337 338 Args:339 video_path (`str`):340 Path to the video file.341 sample_indices_fn (`Callable`):342 A callable function that will return indices at which the video should be sampled. If the video has to be loaded using343 by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`.344 If not provided, simple uniform sampling with fps is performed.345 Example:346 def sample_indices_fn(metadata, **kwargs):347 return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int)348 349 Returns:350 tuple[`np.ndarray`, `VideoMetadata`]: A tuple containing:351 - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).352 - `VideoMetadata` object.353 """354 # Lazy import cv2355 requires_backends(read_video_opencv, ["cv2"])356 import cv2357 358 video = cv2.VideoCapture(video_path)359 total_num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))360 video_fps = video.get(cv2.CAP_PROP_FPS)361 duration = total_num_frames / video_fps if video_fps else 0362 metadata = VideoMetadata(363 total_num_frames=int(total_num_frames),364 fps=float(video_fps),365 duration=float(duration),366 video_backend="opencv",367 height=int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)),368 width=int(video.get(cv2.CAP_PROP_FRAME_WIDTH)),369 )370 indices = sample_indices_fn(metadata=metadata, **kwargs)371 372 index = 0373 frames = []374 while video.isOpened():375 success, frame = video.read()376 if not success:377 break378 if index in indices:379 height, width, channel = frame.shape380 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)381 frames.append(frame[0:height, 0:width, 0:channel])382 if success:383 index += 1384 if index >= total_num_frames:385 break386 387 video.release()388 metadata.frames_indices = indices389 return np.stack(frames), metadata390 391 392def read_video_decord(393 video_path: Union["URL", "Path"],394 sample_indices_fn: Callable,395 **kwargs,396):397 """398 Decode a video using the Decord backend.399 400 Args:401 video_path (`str`):402 Path to the video file.403 sample_indices_fn (`Callable`):404 A callable function that will return indices at which the video should be sampled. If the video has to be loaded using405 by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`.406 If not provided, simple uniform sampling with fps is performed.407 Example:408 def sample_indices_fn(metadata, **kwargs):409 return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int)410 411 Returns:412 tuple[`np.array`, `VideoMetadata`]: A tuple containing:413 - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).414 - `VideoMetadata` object.415 """416 # Lazy import from decord417 requires_backends(read_video_decord, ["decord"])418 from decord import VideoReader, cpu419 420 vr = VideoReader(uri=video_path, ctx=cpu(0)) # decord has problems with gpu421 video_fps = vr.get_avg_fps()422 total_num_frames = len(vr)423 duration = total_num_frames / video_fps if video_fps else 0424 metadata = VideoMetadata(425 total_num_frames=int(total_num_frames),426 fps=float(video_fps),427 duration=float(duration),428 video_backend="decord",429 )430 431 indices = sample_indices_fn(metadata=metadata, **kwargs)432 video = vr.get_batch(indices).asnumpy()433 434 metadata.update(435 {436 "frames_indices": indices,437 "height": video.shape[1],438 "width": video.shape[2],439 }440 )441 return video, metadata442 443 444def read_video_pyav(445 video_path: Union["URL", "Path"],446 sample_indices_fn: Callable,447 **kwargs,448):449 """450 Decode the video with PyAV decoder.451 452 Args:453 video_path (`str`):454 Path to the video file.455 sample_indices_fn (`Callable`, *optional*):456 A callable function that will return indices at which the video should be sampled. If the video has to be loaded using457 by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`.458 If not provided, simple uniform sampling with fps is performed.459 Example:460 def sample_indices_fn(metadata, **kwargs):461 return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int)462 463 Returns:464 tuple[`np.array`, `VideoMetadata`]: A tuple containing:465 - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).466 - `VideoMetadata` object.467 """468 # Lazy import av469 requires_backends(read_video_pyav, ["av"])470 import av471 472 container = av.open(video_path)473 total_num_frames = container.streams.video[0].frames474 video_fps = container.streams.video[0].average_rate # should we better use `av_guess_frame_rate`?475 duration = total_num_frames / video_fps if video_fps else 0476 metadata = VideoMetadata(477 total_num_frames=int(total_num_frames),478 fps=float(video_fps),479 duration=float(duration),480 video_backend="pyav",481 height=container.streams.video[0].height,482 width=container.streams.video[0].width,483 )484 indices = sample_indices_fn(metadata=metadata, **kwargs)485 486 frames = []487 container.seek(0)488 end_index = indices[-1]489 for i, frame in enumerate(container.decode(video=0)):490 if i > end_index:491 break492 if i >= 0 and i in indices:493 frames.append(frame)494 495 video = np.stack([x.to_ndarray(format="rgb24") for x in frames])496 metadata.frames_indices = indices497 return video, metadata498 499 500def read_video_torchvision(501 video_path: Union["URL", "Path"],502 sample_indices_fn: Callable,503 **kwargs,504):505 """506 Decode the video with torchvision decoder.507 508 Args:509 video_path (`str`):510 Path to the video file.511 sample_indices_fn (`Callable`, *optional*):512 A callable function that will return indices at which the video should be sampled. If the video has to be loaded using513 by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`.514 If not provided, simple uniform sampling with fps is performed.515 Example:516 def sample_indices_fn(metadata, **kwargs):517 return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int)518 519 Returns:520 tuple[`torch.Tensor`, `VideoMetadata`]: A tuple containing:521 - Torch tensor of frames in RGB (shape: [num_frames, height, width, 3]).522 - `VideoMetadata` object.523 """524 warnings.warn(525 "Using `torchvision` for video decoding is deprecated and will be removed in future versions. "526 "Please use `torchcodec` instead."527 )528 video, _, info = torchvision_io.read_video(529 video_path,530 start_pts=0.0,531 end_pts=None,532 pts_unit="sec",533 output_format="TCHW",534 )535 video_fps = info["video_fps"]536 total_num_frames = video.size(0)537 duration = total_num_frames / video_fps if video_fps else 0538 metadata = VideoMetadata(539 total_num_frames=int(total_num_frames),540 fps=float(video_fps),541 duration=float(duration),542 video_backend="torchvision",543 )544 545 indices = sample_indices_fn(metadata=metadata, **kwargs)546 547 video = video[indices].contiguous()548 metadata.update(549 {550 "frames_indices": indices,551 "height": video.shape[2],552 "width": video.shape[3],553 }554 )555 return video, metadata556 557 558def read_video_torchcodec(559 video_path: Union["URL", "Path"],560 sample_indices_fn: Callable,561 **kwargs,562):563 """564 Decode the video with torchcodec decoder.565 566 Args:567 video_path (`str`):568 Path to the video file.569 sample_indices_fn (`Callable`):570 A callable function that will return indices at which the video should be sampled. If the video has to be loaded using571 by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`.572 If not provided, simple uniform sampling with fps is performed.573 Example:574 def sample_indices_fn(metadata, **kwargs):575 return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int)576 577 Returns:578 Tuple[`torch.Tensor`, `VideoMetadata`]: A tuple containing:579 - Torch tensor of frames in RGB (shape: [num_frames, height, width, 3]).580 - `VideoMetadata` object.581 """582 # Lazy import torchcodec583 requires_backends(read_video_torchcodec, ["torchcodec"])584 from torchcodec.decoders import VideoDecoder585 586 decoder = VideoDecoder(587 video_path,588 # Interestingly `exact` mode takes less than approximate when we load the whole video589 seek_mode="exact",590 # Allow FFmpeg decide on the number of threads for efficiency591 num_ffmpeg_threads=0,592 device=kwargs.get("device"),593 )594 metadata = VideoMetadata(595 total_num_frames=decoder.metadata.num_frames,596 fps=decoder.metadata.average_fps,597 duration=decoder.metadata.duration_seconds,598 video_backend="torchcodec",599 height=decoder.metadata.height,600 width=decoder.metadata.width,601 )602 indices = sample_indices_fn(metadata=metadata, **kwargs)603 604 video = decoder.get_frames_at(indices=indices).data.contiguous()605 metadata.frames_indices = indices606 return video, metadata607 608 609VIDEO_DECODERS = {610 "decord": read_video_decord,611 "opencv": read_video_opencv,612 "pyav": read_video_pyav,613 "torchvision": read_video_torchvision,614 "torchcodec": read_video_torchcodec,615}616 617 618def load_video(619 video: VideoInput,620 num_frames: Optional[int] = None,621 fps: Optional[Union[int, float]] = None,622 backend: str = "pyav",623 sample_indices_fn: Optional[Callable] = None,624 **kwargs,625) -> np.ndarray:626 """627 Loads `video` to a numpy array.628 629 Args:630 video (`VideoInput`):631 The video to convert to the numpy array format. Can be a link to video or local path.632 num_frames (`int`, *optional*):633 Number of frames to sample uniformly. If not passed, the whole video is loaded.634 fps (`int` or `float`, *optional*):635 Number of frames to sample per second. Should be passed only when `num_frames=None`.636 If not specified and `num_frames==None`, all frames are sampled.637 backend (`str`, *optional*, defaults to `"pyav"`):638 The backend to use when loading the video. Can be any of ["decord", "pyav", "opencv", "torchvision", "torchcodec"]. Defaults to "pyav".639 sample_indices_fn (`Callable`, *optional*):640 A callable function that will return indices at which the video should be sampled. If the video has to be loaded using641 by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`.642 If not provided, simple uniformt sampling with fps is performed, otherwise `sample_indices_fn` has priority over other args.643 The function expects at input the all args along with all kwargs passed to `load_video` and should output valid644 indices at which the video should be sampled. For example:645 646 Example:647 def sample_indices_fn(metadata, **kwargs):648 return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int)649 650 Returns:651 tuple[`np.ndarray`, Dict]: A tuple containing:652 - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).653 - Metadata dictionary.654 """655 656 # If `sample_indices_fn` is given, we can accept any args as those might be needed by custom `sample_indices_fn`657 if fps is not None and num_frames is not None and sample_indices_fn is None:658 raise ValueError(659 "`num_frames`, `fps`, and `sample_indices_fn` are mutually exclusive arguments, please use only one!"660 )661 662 # If user didn't pass a sampling function, create one on the fly with default logic663 if sample_indices_fn is None:664 665 def sample_indices_fn_func(metadata, **fn_kwargs):666 return default_sample_indices_fn(metadata, num_frames=num_frames, fps=fps, **fn_kwargs)667 668 sample_indices_fn = sample_indices_fn_func669 670 # Early exit if provided an array or `PIL` frames671 if not isinstance(video, str):672 metadata = [None] * len(video)673 return video, metadata674 675 if urlparse(video).netloc in ["www.youtube.com", "youtube.com"]:676 if not is_yt_dlp_available():677 raise ImportError("To load a video from YouTube url you have to install `yt_dlp` first.")678 # Lazy import from yt_dlp679 requires_backends(load_video, ["yt_dlp"])680 from yt_dlp import YoutubeDL681 682 buffer = BytesIO()683 with redirect_stdout(buffer), YoutubeDL() as f:684 f.download([video])685 bytes_obj = buffer.getvalue()686 file_obj = BytesIO(bytes_obj)687 elif video.startswith("http://") or video.startswith("https://"):688 file_obj = BytesIO(requests.get(video).content)689 elif os.path.isfile(video):690 file_obj = video691 else:692 raise TypeError("Incorrect format used for video. Should be an url linking to an video or a local path.")693 694 # can also load with decord, but not cv2/torchvision695 # both will fail in case of url links696 video_is_url = video.startswith("http://") or video.startswith("https://")697 if video_is_url and backend == "opencv":698 raise ValueError("If you are trying to load a video from URL, you cannot use 'opencv' as backend")699 700 if (701 (not is_decord_available() and backend == "decord")702 or (not is_av_available() and backend == "pyav")703 or (not is_cv2_available() and backend == "opencv")704 or (not is_torchvision_available() and backend == "torchvision")705 or (not is_torchcodec_available() and backend == "torchcodec")706 ):707 raise ImportError(708 f"You chose backend={backend} for loading the video but the required library is not found in your environment "709 f"Make sure to install {backend} before loading the video."710 )711 712 video_decoder = VIDEO_DECODERS[backend]713 video, metadata = video_decoder(file_obj, sample_indices_fn, **kwargs)714 return video, metadata715 716 717def convert_to_rgb(718 video: np.ndarray,719 input_data_format: Optional[Union[str, ChannelDimension]] = None,720) -> np.ndarray:721 """722 Convert video to RGB by blending the transparency layer if it's in RGBA format, otherwise simply returns it.723 724 Args:725 video (`np.ndarray`):726 The video to convert.727 input_data_format (`ChannelDimension`, *optional*):728 The channel dimension format of the input video. If unset, will use the inferred format from the input.729 """730 if not isinstance(video, np.ndarray):731 raise TypeError(f"Video has to be a numpy array to convert to RGB format, but found {type(video)}")732 733 # np.array usually comes with ChannelDimension.LAST so let's convert it734 if input_data_format is None:735 input_data_format = infer_channel_dimension_format(video)736 video = to_channel_dimension_format(video, ChannelDimension.FIRST, input_channel_dim=input_data_format)737 738 # 3 channels for RGB already739 if video.shape[-3] == 3:740 return video741 742 # Grayscale video so we repeat it 3 times for each channel743 if video.shape[-3] == 1:744 return video.repeat(3, -3)745 746 if not (video[..., 3, :, :] < 255).any():747 return video748 749 # There is a transparency layer, blend it with a white background.750 # Calculate the alpha proportion for blending.751 alpha = video[..., 3, :, :] / 255.0752 video = (1 - alpha[..., None, :, :]) * 255 + alpha[..., None, :, :] * video[..., 3, :, :]753 return video754 755 756def pad(757 video: np.ndarray,758 padding: Union[int, tuple[int, int], Iterable[tuple[int, int]]],759 mode: PaddingMode = PaddingMode.CONSTANT,760 constant_values: Union[float, Iterable[float]] = 0.0,761 data_format: Optional[Union[str, ChannelDimension]] = None,762 input_data_format: Optional[Union[str, ChannelDimension]] = None,763) -> np.ndarray:764 """765 Pads the `video` with the specified (height, width) `padding` and `mode`.766 767 Args:768 video (`np.ndarray`):769 The video to pad.770 padding (`int` or `tuple[int, int]` or `Iterable[tuple[int, int]]`):771 Padding to apply to the edges of the height, width axes. Can be one of three formats:772 - `((before_height, after_height), (before_width, after_width))` unique pad widths for each axis.773 - `((before, after),)` yields same before and after pad for height and width.774 - `(pad,)` or int is a shortcut for before = after = pad width for all axes.775 mode (`PaddingMode`):776 The padding mode to use. Can be one of:777 - `"constant"`: pads with a constant value.778 - `"reflect"`: pads with the reflection of the vector mirrored on the first and last values of the779 vector along each axis.780 - `"replicate"`: pads with the replication of the last value on the edge of the array along each axis.781 - `"symmetric"`: pads with the reflection of the vector mirrored along the edge of the array.782 constant_values (`float` or `Iterable[float]`, *optional*):783 The value to use for the padding if `mode` is `"constant"`.784 data_format (`str` or `ChannelDimension`, *optional*):785 The channel dimension format for the output video. Can be one of:786 - `"channels_first"` or `ChannelDimension.FIRST`: video in (num_frames, num_channels, height, width) format.787 - `"channels_last"` or `ChannelDimension.LAST`: video in (num_frames, height, width, num_channels) format.788 If unset, will use same as the input video.789 input_data_format (`str` or `ChannelDimension`, *optional*):790 The channel dimension format for the input video. Can be one of:791 - `"channels_first"` or `ChannelDimension.FIRST`: video in (num_frames, num_channels, height, width) format.792 - `"channels_last"` or `ChannelDimension.LAST`: video in (num_frames, height, width, num_channels) format.793 If unset, will use the inferred format of the input video.794 795 Returns:796 `np.ndarray`: The padded video.797 798 """799 if input_data_format is None:800 input_data_format = infer_channel_dimension_format(video)801 802 def _expand_for_data_format(values):803 """804 Convert values to be in the format expected by np.pad based on the data format.805 """806 if isinstance(values, (int, float)):807 values = ((values, values), (values, values))808 elif isinstance(values, tuple) and len(values) == 1:809 values = ((values[0], values[0]), (values[0], values[0]))810 elif isinstance(values, tuple) and len(values) == 2 and isinstance(values[0], int):811 values = (values, values)812 elif isinstance(values, tuple) and len(values) == 2 and isinstance(values[0], tuple):813 pass814 else:815 raise ValueError(f"Unsupported format: {values}")816 817 # add 0 for channel dimension818 values = (819 ((0, 0), (0, 0), *values) if input_data_format == ChannelDimension.FIRST else ((0, 0), *values, (0, 0))820 )821 822 # Add additional padding if there's a batch dimension823 values = (0, *values) if video.ndim == 5 else values824 return values825 826 padding_map = {827 PaddingMode.CONSTANT: "constant",828 PaddingMode.REFLECT: "reflect",829 PaddingMode.REPLICATE: "replicate",830 PaddingMode.SYMMETRIC: "symmetric",831 }832 padding = _expand_for_data_format(padding)833 834 pad_kwargs = {}835 if mode not in padding_map:836 raise ValueError(f"Invalid padding mode: {mode}")837 elif mode == PaddingMode.CONSTANT:838 pad_kwargs["constant_values"] = _expand_for_data_format(constant_values)839 840 video = np.pad(video, padding, mode=padding_map[mode], **pad_kwargs)841 video = to_channel_dimension_format(video, data_format, input_data_format) if data_format is not None else video842 return video843 844 845def group_videos_by_shape(846 videos: list["torch.Tensor"],847) -> tuple[dict[tuple[int, int], "torch.Tensor"], dict[int, tuple[tuple[int, int], int]]]:848 """849 Groups videos by shape.850 Returns a dictionary with the shape as key and a list of videos with that shape as value,851 and a dictionary with the index of the video in the original list as key and the shape and index in the grouped list as value.852 """853 grouped_videos = {}854 grouped_videos_index = {}855 for i, video in enumerate(videos):856 shape = video.shape[-2::]857 num_frames = video.shape[-4] # video format BTCHW858 shape = (num_frames, *shape)859 if shape not in grouped_videos:860 grouped_videos[shape] = []861 grouped_videos[shape].append(video)862 grouped_videos_index[i] = (shape, len(grouped_videos[shape]) - 1)863 # stack videos with the same size and number of frames864 grouped_videos = {shape: torch.stack(videos, dim=0) for shape, videos in grouped_videos.items()}865 return grouped_videos, grouped_videos_index866 867 868def reorder_videos(869 processed_videos: dict[tuple[int, int], "torch.Tensor"],870 grouped_videos_index: dict[int, tuple[tuple[int, int], int]],871) -> list["torch.Tensor"]:872 """873 Reconstructs a list of videos in the original order.874 """875 return [876 processed_videos[grouped_videos_index[i][0]][grouped_videos_index[i][1]]877 for i in range(len(grouped_videos_index))878 ]879