| from __future__ import annotations |
|
|
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from PIL import Image |
|
|
|
|
| def _frame_to_tensor(frame: Any) -> torch.Tensor: |
| if isinstance(frame, Image.Image): |
| arr = np.asarray(frame.convert("RGB"), dtype=np.float32) / 255.0 |
| return torch.from_numpy(arr).permute(2, 0, 1) |
| if isinstance(frame, np.ndarray): |
| arr = frame.astype(np.float32, copy=False) |
| if arr.max() > 1.0: |
| arr = arr / 255.0 |
| tensor = torch.from_numpy(arr) |
| if tensor.ndim == 3 and tensor.shape[-1] == 3: |
| tensor = tensor.permute(2, 0, 1) |
| if tensor.ndim != 3: |
| raise ValueError(f"Expected frame array with 3 dims, got {arr.shape}") |
| return tensor.float().contiguous() |
| if torch.is_tensor(frame): |
| tensor = frame.detach().float() |
| if tensor.ndim == 3 and tensor.shape[-1] == 3: |
| tensor = tensor.permute(2, 0, 1) |
| if tensor.ndim != 3: |
| raise ValueError(f"Expected frame tensor with 3 dims, got {tuple(tensor.shape)}") |
| if tensor.max() > 1.0: |
| tensor = tensor / 255.0 |
| return tensor.contiguous() |
| raise TypeError(f"Unsupported frame type: {type(frame)!r}") |
|
|
|
|
| def video_to_t_c_h_w(video: Any) -> torch.Tensor: |
| if torch.is_tensor(video): |
| tensor = video.detach().float() |
| if tensor.ndim == 5 and tensor.shape[0] == 1: |
| tensor = tensor.squeeze(0) |
| if tensor.ndim == 3: |
| tensor = tensor.unsqueeze(0) |
| if tensor.ndim != 4: |
| raise ValueError(f"Expected video tensor with 4 dims, got {tuple(tensor.shape)}") |
| if tensor.shape[-1] == 3: |
| tensor = tensor.permute(0, 3, 1, 2) |
| if tensor.max() > 1.0: |
| tensor = tensor / 255.0 |
| return tensor.contiguous() |
|
|
| if isinstance(video, np.ndarray): |
| arr = video.astype(np.float32, copy=False) |
| if arr.ndim == 5 and arr.shape[0] == 1: |
| arr = arr[0] |
| if arr.ndim == 3: |
| return _frame_to_tensor(arr).unsqueeze(0).contiguous() |
| if arr.ndim != 4: |
| raise ValueError(f"Expected video array with 4 dims, got {arr.shape}") |
| tensor = torch.from_numpy(arr) |
| if tensor.shape[-1] == 3: |
| tensor = tensor.permute(0, 3, 1, 2) |
| if tensor.max() > 1.0: |
| tensor = tensor.float() / 255.0 |
| return tensor.float().contiguous() |
|
|
| if isinstance(video, (list, tuple)): |
| if not video: |
| raise ValueError("Video frame list is empty.") |
| return torch.stack([_frame_to_tensor(frame) for frame in video], dim=0).contiguous() |
|
|
| if isinstance(video, Image.Image): |
| return _frame_to_tensor(video).unsqueeze(0).contiguous() |
|
|
| raise TypeError(f"Unsupported video type: {type(video)!r}") |
|
|
|
|
| def videos_to_b_t_c_h_w(videos: list[Any], *, size: int | tuple[int, int] | None = None, num_frames: int | None = None) -> torch.Tensor: |
| tensors = [video_to_t_c_h_w(video) for video in videos] |
| if num_frames is not None: |
| tensors = [_match_frames(tensor, num_frames) for tensor in tensors] |
| if size is not None: |
| target_size = (size, size) if isinstance(size, int) else tuple(size) |
| tensors = [_resize_video(tensor, target_size) for tensor in tensors] |
| return torch.stack(tensors, dim=0).clamp(0.0, 1.0).contiguous() |
|
|
|
|
| def _match_frames(video: torch.Tensor, num_frames: int) -> torch.Tensor: |
| if video.shape[0] == num_frames: |
| return video |
| if video.shape[0] > num_frames: |
| idx = torch.linspace(0, video.shape[0] - 1, steps=num_frames).round().long() |
| return video[idx] |
| reps = int(np.ceil(num_frames / video.shape[0])) |
| return video.repeat((reps, 1, 1, 1))[:num_frames] |
|
|
|
|
| def _resize_video(video: torch.Tensor, size: tuple[int, int]) -> torch.Tensor: |
| return F.interpolate(video, size=size, mode="bilinear", align_corners=False) |
|
|
