(
flow_predictor: FlowPredictor,
full_video_path: Path,
subsampled_path: Path,
target_num_frames: int,
flow_resolution: int,
device: torch.device,
)
| 96 | |
| 97 | |
| 98 | def subsample_frames( |
| 99 | flow_predictor: FlowPredictor, |
| 100 | full_video_path: Path, |
| 101 | subsampled_path: Path, |
| 102 | target_num_frames: int, |
| 103 | flow_resolution: int, |
| 104 | device: torch.device, |
| 105 | ) -> None: |
| 106 | # Just symlink the frames if they don't need to be subsampled. |
| 107 | if len(list(full_video_path.iterdir())) <= target_num_frames: |
| 108 | subsampled_path.parent.mkdir(exist_ok=True, parents=True) |
| 109 | shutil.copytree(full_video_path, subsampled_path) |
| 110 | return |
| 111 | |
| 112 | last = None |
| 113 | mean_flows = [] |
| 114 | for image in tqdm( |
| 115 | list(sorted(full_video_path.iterdir())), desc="Computing mean flows" |
| 116 | ): |
| 117 | # Get the current two-frame video segment. |
| 118 | if last is None: |
| 119 | last = resize_to_resolution(load_image(image), flow_resolution) |
| 120 | continue |
| 121 | current = resize_to_resolution(load_image(image), flow_resolution) |
| 122 | videos = torch.stack((last, current))[None].to(device) |
| 123 | |
| 124 | # Crop the video segment. |
| 125 | _, _, _, h, w = videos.shape |
| 126 | new_shape = compute_patch_cropped_shape((h, w), 8) |
| 127 | videos = center_crop_images(videos, new_shape) |
| 128 | |
| 129 | # Compute the mean flows. |
| 130 | mean_flows.append(flow_predictor.forward(videos).norm(dim=-1).mean().item()) |
| 131 | last = current |
| 132 | |
| 133 | flow_step = sum(mean_flows) / target_num_frames |
| 134 | remaining = 0 |
| 135 | subsampled_path.mkdir(exist_ok=True, parents=True) |
| 136 | num_saved = 0 |
| 137 | for mean_flow, frame in zip(mean_flows, sorted(full_video_path.iterdir())): |
| 138 | if remaining <= 0: |
| 139 | shutil.copy(frame, subsampled_path / frame.name) |
| 140 | remaining += flow_step |
| 141 | num_saved += 1 |
| 142 | |
| 143 | remaining -= mean_flow |
| 144 | |
| 145 | # Randomly fill in the remaining frames. |
| 146 | generator = np.random.default_rng(seed=0) |
| 147 | paths = list(full_video_path.iterdir()) |
| 148 | while num_saved < target_num_frames: |
| 149 | # Pick a random frame. |
| 150 | frame = paths[generator.choice(len(paths))] |
| 151 | if (subsampled_path / frame.name).exists(): |
| 152 | continue |
| 153 | shutil.copy(frame, subsampled_path / frame.name) |
| 154 | num_saved += 1 |
| 155 |
no test coverage detected