(video, frame_indices)
| 104 | |
| 105 | |
| 106 | def pick_video_frame(video, frame_indices): |
| 107 | batch, device = video.shape[0], video.device |
| 108 | video = rearrange(video, "b c f ... -> b f c ...") |
| 109 | batch_indices = torch.arange(batch, device=device) |
| 110 | batch_indices = rearrange(batch_indices, "b -> b 1") |
| 111 | images = video[batch_indices, frame_indices] |
| 112 | images = rearrange(images, "b 1 c ... -> b c ...") |
| 113 | return images |
| 114 | |
| 115 | |
| 116 | # gan related |