| 1 | """Utilities for selecting paired audio/video memory windows.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import math |
| 6 | import random |
| 7 | from typing import Literal |
| 8 | |
| 9 | import torch |
| 10 | from torch import Tensor |
| 11 | |
| 12 | |
| 13 | def latent_window_size_to_pixel_window_size( |
| 14 | latent_window_size: int, |
| 15 | *, |
| 16 | downsample_factor: int, |
| 17 | is_causal: bool = True, |
| 18 | ) -> int: |
| 19 | if latent_window_size <= 0: |
| 20 | raise ValueError(f"latent_window_size must be positive, got {latent_window_size}") |
| 21 | if downsample_factor <= 0: |
| 22 | raise ValueError(f"downsample_factor must be positive, got {downsample_factor}") |
| 23 | |
| 24 | pixel_window_size = int(latent_window_size) * int(downsample_factor) |
| 25 | if is_causal: |
| 26 | pixel_window_size = max(pixel_window_size - (int(downsample_factor) - 1), 1) |
| 27 | return pixel_window_size |
| 28 | |
| 29 | |
| 30 | def select_max_response_audio_window_with_bounds( |
| 31 | segment: Tensor, |
| 32 | window_size: int, |
| 33 | ) -> tuple[Tensor, Tensor, Tensor]: |
| 34 | if segment.dim() != 4: |
| 35 | raise ValueError(f"Expected segment shape [B, C, T, F], got {tuple(segment.shape)}") |
| 36 | if window_size <= 0: |
| 37 | raise ValueError(f"window_size must be positive, got {window_size}") |
| 38 | |
| 39 | num_time_steps = segment.shape[2] |
| 40 | if num_time_steps <= 0: |
| 41 | raise ValueError("Cannot select from an empty audio segment") |
| 42 | |
| 43 | scan_stride = max(1, window_size // 4) |
| 44 | offsets = torch.arange(window_size, device=segment.device) |
| 45 | max_start_idx = num_time_steps - window_size if num_time_steps >= window_size else num_time_steps - 1 |
| 46 | candidate_start_indices = list(range(0, max_start_idx + 1, scan_stride)) |
| 47 | if candidate_start_indices[-1] != max_start_idx: |
| 48 | candidate_start_indices.append(max_start_idx) |
| 49 | |
| 50 | candidate_windows = [] |
| 51 | candidate_scores = [] |
| 52 | candidate_start_indices_tensor = torch.tensor(candidate_start_indices, device=segment.device, dtype=torch.long) |
| 53 | for start_idx in candidate_start_indices: |
| 54 | gather_indices = (start_idx + offsets).clamp(0, num_time_steps - 1).long() |
| 55 | window = segment.index_select(dim=2, index=gather_indices) |
| 56 | candidate_windows.append(window) |
| 57 | candidate_scores.append(window.float().exp().sum(dim=(1, 2, 3))) |
| 58 | |
| 59 | scores = torch.stack(candidate_scores, dim=1) |
| 60 | best_window_indices = scores.argmax(dim=1) |
| 61 | best_start_indices = candidate_start_indices_tensor[best_window_indices] |
| 62 | best_end_indices = torch.clamp(best_start_indices + window_size - 1, max=num_time_steps - 1) |
| 63 | selected_windows = torch.cat( |
| 64 | [ |
| 65 | candidate_windows[int(best_window_indices[batch_index])][batch_index : batch_index + 1] |
| 66 | for batch_index in range(segment.shape[0]) |
| 67 | ], |
| 68 | dim=0, |
| 69 | ) |
| 70 | return selected_windows, best_start_indices, best_end_indices |
| 71 | |
| 72 | |
| 73 | def select_random_audio_window_with_bounds( |
| 74 | segment: Tensor, |
| 75 | window_size: int, |
| 76 | *, |
| 77 | rng: random.Random | None = None, |
| 78 | ) -> tuple[Tensor, Tensor, Tensor]: |
| 79 | if segment.dim() != 4: |
| 80 | raise ValueError(f"Expected segment shape [B, C, T, F], got {tuple(segment.shape)}") |
| 81 | if window_size <= 0: |
| 82 | raise ValueError(f"window_size must be positive, got {window_size}") |
| 83 | |
| 84 | batch_size = segment.shape[0] |
| 85 | num_time_steps = segment.shape[2] |
| 86 | if num_time_steps <= 0: |
| 87 | raise ValueError("Cannot select from an empty audio segment") |
| 88 | |
| 89 | if num_time_steps <= window_size: |
| 90 | start_indices = torch.zeros(batch_size, device=segment.device, dtype=torch.long) |
| 91 | else: |
| 92 | max_start_idx = num_time_steps - window_size |
| 93 | if rng is None: |
| 94 | start_indices = torch.randint( |
| 95 | low=0, |
| 96 | high=max_start_idx + 1, |
| 97 | size=(batch_size,), |
| 98 | device=segment.device, |
| 99 | ) |
| 100 | else: |
| 101 | start_indices = torch.tensor( |
| 102 | [rng.randint(0, max_start_idx) for _ in range(batch_size)], |
| 103 | device=segment.device, |
| 104 | dtype=torch.long, |
| 105 | ) |
| 106 | |
| 107 | offsets = torch.arange(window_size, device=segment.device, dtype=torch.long) |
| 108 | selected_windows = torch.cat( |
| 109 | [ |
| 110 | segment[ |
| 111 | batch_index : batch_index + 1, |
| 112 | :, |
| 113 | (start_indices[batch_index] + offsets).clamp(0, num_time_steps - 1), |
| 114 | :, |
| 115 | ] |
| 116 | for batch_index in range(batch_size) |
| 117 | ], |
| 118 | dim=0, |
| 119 | ) |
| 120 | end_indices = torch.clamp(start_indices + window_size - 1, max=num_time_steps - 1) |
| 121 | return selected_windows, start_indices, end_indices |
| 122 | |
| 123 | |
| 124 | def select_audio_window_with_bounds( |
| 125 | segment: Tensor, |
| 126 | window_size: int, |
| 127 | *, |
| 128 | mode: Literal["max_response", "random"] = "random", |
| 129 | rng: random.Random | None = None, |
| 130 | ) -> tuple[Tensor, Tensor, Tensor]: |
| 131 | if mode == "max_response": |
| 132 | return select_max_response_audio_window_with_bounds(segment, window_size) |
| 133 | if mode == "random": |
| 134 | return select_random_audio_window_with_bounds(segment, window_size, rng=rng) |
| 135 | raise ValueError(f"Unsupported audio window selection mode: {mode}") |
| 136 | |
| 137 | |
| 138 | def mel_window_bounds_to_seconds( |
| 139 | start_index: int, |
| 140 | end_index: int, |
| 141 | *, |
| 142 | hop_length: int, |
| 143 | sample_rate: int, |
| 144 | ) -> tuple[float, float]: |
| 145 | if start_index < 0: |
| 146 | raise ValueError(f"start_index must be non-negative, got {start_index}") |
| 147 | if end_index < start_index: |
| 148 | raise ValueError(f"end_index must be >= start_index, got start={start_index}, end={end_index}") |
| 149 | if hop_length <= 0: |
| 150 | raise ValueError(f"hop_length must be positive, got {hop_length}") |
| 151 | if sample_rate <= 0: |
| 152 | raise ValueError(f"sample_rate must be positive, got {sample_rate}") |
| 153 | |
| 154 | start_time_sec = float(start_index * hop_length) / float(sample_rate) |
| 155 | end_time_sec = float((end_index + 1) * hop_length) / float(sample_rate) |
| 156 | return start_time_sec, end_time_sec |
| 157 | |
| 158 | |
| 159 | def select_video_frame_indices_from_time_range( |
| 160 | *, |
| 161 | num_frames: int, |
| 162 | fps: float, |
| 163 | start_time_sec: float, |
| 164 | end_time_sec: float, |
| 165 | count: int = 1, |
| 166 | mode: Literal["first", "random", "center"] = "center", |
| 167 | rng: random.Random | None = None, |
| 168 | ) -> list[int]: |
| 169 | if num_frames <= 0: |
| 170 | raise ValueError(f"num_frames must be positive, got {num_frames}") |
| 171 | if fps <= 0: |
| 172 | raise ValueError(f"fps must be positive, got {fps}") |
| 173 | if count <= 0: |
| 174 | raise ValueError(f"count must be positive, got {count}") |
| 175 | if end_time_sec < start_time_sec: |
| 176 | raise ValueError(f"end_time_sec must be >= start_time_sec, got {start_time_sec}, {end_time_sec}") |
| 177 | |
| 178 | mode = mode.lower() |
| 179 | if mode not in {"first", "random", "center"}: |
| 180 | raise ValueError(f"Unsupported frame selection mode: {mode}") |
| 181 | |
| 182 | start_frame = int(math.ceil(start_time_sec * fps)) |
| 183 | end_frame = int(math.ceil(end_time_sec * fps)) - 1 |
| 184 | start_frame = max(0, min(start_frame, num_frames - 1)) |
| 185 | end_frame = max(0, min(end_frame, num_frames - 1)) |
| 186 | |
| 187 | if end_frame < start_frame: |
| 188 | center_time_sec = max(0.0, 0.5 * (start_time_sec + end_time_sec)) |
| 189 | center_frame = int(round(center_time_sec * fps)) |
| 190 | candidate_frames = [max(0, min(center_frame, num_frames - 1))] |
| 191 | else: |
| 192 | candidate_frames = list(range(start_frame, end_frame + 1)) |
| 193 | |
| 194 | if mode == "first": |
| 195 | selected = candidate_frames[:count] |
| 196 | elif mode == "center": |
| 197 | if len(candidate_frames) <= count: |
| 198 | selected = candidate_frames[:] |
| 199 | else: |
| 200 | center_offset = max(0, (len(candidate_frames) - count) // 2) |
| 201 | selected = candidate_frames[center_offset : center_offset + count] |
| 202 | else: |
| 203 | rng = rng or random |
| 204 | selected = ( |
| 205 | candidate_frames[:] |
| 206 | if len(candidate_frames) <= count |
| 207 | else sorted(rng.sample(candidate_frames, count)) |
| 208 | ) |
| 209 | |
| 210 | if len(selected) < count: |
| 211 | selected.extend([selected[-1]] * (count - len(selected))) |
| 212 | return selected |
| 213 |