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
(
video_path,
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,
)
| 28 | |
| 29 | |
| 30 | def load_video_frames( |
| 31 | video_path, |
| 32 | image_size, |
| 33 | offload_video_to_cpu, |
| 34 | img_mean=(0.485, 0.456, 0.406), |
| 35 | img_std=(0.229, 0.224, 0.225), |
| 36 | async_loading_frames=False, |
| 37 | frame_names=None, |
| 38 | ): |
| 39 | """ |
| 40 | Load the video frames from a directory of JPEG files ("<frame_index>.jpg" format). |
| 41 | |
| 42 | The frames are resized to image_size x image_size and are loaded to GPU if |
| 43 | `offload_video_to_cpu` is `False` and to CPU if `offload_video_to_cpu` is `True`. |
| 44 | |
| 45 | You can load a frame asynchronously by setting `async_loading_frames` to `True`. |
| 46 | """ |
| 47 | if isinstance(video_path, str) and os.path.isdir(video_path): |
| 48 | jpg_folder = video_path |
| 49 | else: |
| 50 | raise NotImplementedError("Only JPEG frames are supported at this moment") |
| 51 | if frame_names is None: |
| 52 | frame_names = [p for p in os.listdir(jpg_folder) if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG", ".png"]] |
| 53 | frame_names.sort(key=lambda p: int(os.path.splitext(p)[0])) |
| 54 | |
| 55 | num_frames = len(frame_names) |
| 56 | if num_frames == 0: |
| 57 | raise RuntimeError(f"no images found in {jpg_folder}") |
| 58 | img_paths = [os.path.join(jpg_folder, frame_name) for frame_name in frame_names] |
| 59 | img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None] |
| 60 | img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None] |
| 61 | |
| 62 | if async_loading_frames: |
| 63 | lazy_images = AsyncVideoFrameLoader(img_paths, image_size, offload_video_to_cpu, img_mean, img_std) |
| 64 | return lazy_images, lazy_images.video_height, lazy_images.video_width |
| 65 | |
| 66 | images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32) |
| 67 | for n, img_path in enumerate(tqdm(img_paths, desc="frame loading (JPEG)")): |
| 68 | images[n], video_height, video_width = _load_img_as_tensor(img_path, image_size) |
| 69 | if not offload_video_to_cpu: |
| 70 | images = images.cuda() |
| 71 | img_mean = img_mean.cuda() |
| 72 | img_std = img_std.cuda() |
| 73 | # normalize by mean and std |
| 74 | images -= img_mean |
| 75 | images /= img_std |
| 76 | return images, video_height, video_width |
| 77 | |
| 78 | |
| 79 | def load_video_frames_v2( |