Load the video frames from a directory of JPEG files (" .jpg" format). The frames are resized to image_size x image_size and are loaded to GPU if `offload_video_to_cpu` is `False` and to CPU if `offload_video_to_cpu` is `True`. You can load a frame asynchronously by se
(
frames,
image_size,
offload_video_to_cpu,
img_mean=(0.485, 0.456, 0.406),
img_std=(0.229, 0.224, 0.225),
async_loading_frames=False,
frame_names=None,
)
| 77 | |
| 78 | |
| 79 | def load_video_frames_v2( |
| 80 | frames, |
| 81 | image_size, |
| 82 | offload_video_to_cpu, |
| 83 | img_mean=(0.485, 0.456, 0.406), |
| 84 | img_std=(0.229, 0.224, 0.225), |
| 85 | async_loading_frames=False, |
| 86 | frame_names=None, |
| 87 | ): |
| 88 | """ |
| 89 | Load the video frames from a directory of JPEG files ("<frame_index>.jpg" format). |
| 90 | |
| 91 | The frames are resized to image_size x image_size and are loaded to GPU if |
| 92 | `offload_video_to_cpu` is `False` and to CPU if `offload_video_to_cpu` is `True`. |
| 93 | |
| 94 | You can load a frame asynchronously by setting `async_loading_frames` to `True`. |
| 95 | """ |
| 96 | num_frames = len(frames) |
| 97 | img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None] |
| 98 | img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None] |
| 99 | |
| 100 | images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32) |
| 101 | for n, frame in enumerate(tqdm(frames, desc="video frame")): |
| 102 | images[n], video_height, video_width = _load_img_v2_as_tensor(frame, image_size) |
| 103 | if not offload_video_to_cpu: |
| 104 | images = images.cuda() |
| 105 | img_mean = img_mean.cuda() |
| 106 | img_std = img_std.cuda() |
| 107 | # normalize by mean and std |
| 108 | images -= img_mean |
| 109 | images /= img_std |
| 110 | return images, video_height, video_width |
| 111 | |
| 112 | |
| 113 | def build_sam2_video_predictor( |
no test coverage detected