Spaces:
Sleeping
Sleeping
| # Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved | |
| # pyre-unsafe | |
| """ | |
| This file extracts the frames for the frame datasets in SA-CO/Gold and Silver. | |
| Call like: | |
| > python extract_frames.py <dataset_name> | |
| """ | |
| import json | |
| import os | |
| import shutil | |
| import sys | |
| from multiprocessing import Pool | |
| from PIL import Image | |
| from tqdm import tqdm | |
| from utils import ( | |
| annotation_files, | |
| config, | |
| get_frame_from_video, | |
| is_valid_image, | |
| update_annotations, | |
| ) | |
| def extract_frame(path_video, global_frame_idx, path_frame, image_size, file_name): | |
| frame = get_frame_from_video(path_video, global_frame_idx) | |
| os.makedirs(os.path.dirname(path_frame), exist_ok=True) | |
| img = Image.fromarray(frame) | |
| if frame.shape[:2] != image_size: | |
| print(f"Resizing image {file_name} from {frame.shape[:2]} to {image_size}") | |
| height, width = image_size | |
| img = img.resize((width, height)) # Uses Image.NEAREST by default | |
| img.save(path_frame) | |
| def process_image(args): | |
| image, dataset_name, config = args | |
| original_video, global_frame_idx, file_name, image_size = image | |
| extra_subpath = "" | |
| if dataset_name == "ego4d": | |
| extra_subpath = "v1/clips" | |
| elif dataset_name == "yt1b": | |
| original_video = f"video_{original_video}.mp4" | |
| elif dataset_name == "sav": | |
| extra_subpath = "videos_fps_6" | |
| path_video = os.path.join( | |
| config[f"{dataset_name}_path"], | |
| "downloaded_videos", | |
| extra_subpath, | |
| original_video, | |
| ) | |
| path_frame = os.path.join(config[f"{dataset_name}_path"], "frames", file_name) | |
| to_return = file_name | |
| try: | |
| extract_frame(path_video, global_frame_idx, path_frame, image_size, file_name) | |
| if not is_valid_image(path_frame): | |
| print(f"Invalid image in {path_frame}") | |
| to_return = None | |
| except: | |
| print(f"Invalid image in {path_frame}") | |
| to_return = None | |
| return to_return | |
| def main(): | |
| assert len(sys.argv) > 1, "You have to provide the name of the dataset" | |
| dataset_name = sys.argv[1] | |
| assert dataset_name in annotation_files, ( | |
| f"The dataset can be one of {list(annotation_files.keys())}" | |
| ) | |
| all_outputs = [] | |
| for file in annotation_files[dataset_name]: | |
| with open(os.path.join(config["path_annotations"], file), "r") as f: | |
| annotation = json.load(f) | |
| images = annotation["images"] | |
| images = set( | |
| ( | |
| image["original_video"], | |
| image["global_frame_idx"], | |
| image["file_name"], | |
| tuple(image["image_size"]), | |
| ) | |
| for image in images | |
| ) | |
| args_list = [(image, dataset_name, config) for image in images] | |
| with Pool(os.cpu_count()) as pool: | |
| outputs = list( | |
| tqdm(pool.imap_unordered(process_image, args_list), total=len(images)) | |
| ) | |
| all_outputs.extend(outputs) | |
| if any(out is None for out in outputs): | |
| update_annotations(dataset_name, all_outputs, key="file_name") | |
| if config[f"remove_downloaded_videos_{dataset_name}"]: | |
| shutil.rmtree(os.path.join(config[f"{dataset_name}_path"], "downloaded_videos")) | |
| if __name__ == "__main__": | |
| main() | |