Spaces:
Running on Zero
Running on Zero
File size: 2,012 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 | from __future__ import annotations
import torch
def target_frame_count(frames: int, temporal_factor: int) -> int:
if frames <= 0 or temporal_factor <= 0:
raise ValueError("frames and temporal_factor must be positive")
return ((frames - 1 + temporal_factor - 1) // temporal_factor) * temporal_factor + 1
def valid_frame_count(frames: int, temporal_factor: int) -> int:
"""Largest ``temporal_factor * n + 1`` frame count not exceeding input."""
if frames <= 0 or temporal_factor <= 0:
raise ValueError("frames and temporal_factor must be positive")
return ((frames - 1) // temporal_factor) * temporal_factor + 1
def truncate_video_frames(
video: torch.Tensor, temporal_factor: int
) -> tuple[torch.Tensor, int]:
count = valid_frame_count(video.shape[2], temporal_factor)
return video[:, :, :count], count
def repeat_last_frame(video: torch.Tensor, target_frames: int) -> torch.Tensor:
if target_frames < video.shape[2]:
return video[:, :, :target_frames]
missing = target_frames - video.shape[2]
if missing == 0:
return video
return torch.cat(
[video, video[:, :, -1:].repeat(1, 1, missing, 1, 1)], dim=2
)
def pad_video_frames(video: torch.Tensor, temporal_factor: int, mode: str) -> tuple[torch.Tensor, int]:
original = video.shape[2]
target = target_frame_count(original, temporal_factor)
missing = target - original
if missing == 0:
return video, original
if mode == "repeat":
padding = video[:, :, -1:].repeat(1, 1, missing, 1, 1)
elif mode == "pingpong":
if original == 1:
padding = video.repeat(1, 1, missing, 1, 1)
else:
indices = list(range(original - 2, -1, -1))
indices = (indices * ((missing + len(indices) - 1) // len(indices)))[:missing]
padding = video[:, :, indices]
else:
raise ValueError(f"Unknown padding mode: {mode}")
return torch.cat([video, padding], dim=2), original
|