| import os | |
| import json | |
| import numpy as np | |
| import argparse | |
| from pytubefix import YouTube | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--root", type=str, help="Directory of WiLoR") | |
| parser.add_argument("--mode", type=str, choices=['train', 'test'], default= 'train', help="Train/Test set") | |
| args = parser.parse_args() | |
| with open(os.path.join(args.root, f'./whim/{args.mode}_video_ids.json')) as f: | |
| video_dict = json.load(f) | |
| Video_IDs = video_dict.keys() | |
| failed_IDs = [] | |
| os.makedirs(os.path.join(args.root, 'Videos'), exist_ok=True) | |
| for Video_ID in Video_IDs: | |
| res = video_dict[Video_ID]['res'][0] | |
| try: | |
| YouTube('https://youtu.be/'+Video_ID).streams.filter(only_video=True, | |
| file_extension='mp4', | |
| res =f'{res}p' | |
| ).order_by('resolution').desc().first().download( | |
| output_path=os.path.join(args.root, 'Videos') , | |
| filename = Video_ID +'.mp4') | |
| except: | |
| print(f'Failed {Video_ID}') | |
| failed_IDs.append(Video_ID) | |
| continue | |
| cap = cv2.VideoCapture(os.path.join(args.root, 'Videos', Video_ID + '.mp4')) | |
| if (cap.isOpened()== False): | |
| print(f"Error opening video stream {os.path.join(args.root, 'Videos', Video_ID + '.mp4')}") | |
| VIDEO_LEN = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| fps_org = video_dict[Video_ID]['fps'] | |
| fps_rate = round(fps / fps_org) | |
| all_frames = os.listdir(os.path.join(args.root, 'WHIM', args.mode, 'anno', Video_ID)) | |
| for frame in all_frames: | |
| frame_gt = int(frame[:-4]) | |
| frame_idx = (frame_gt * fps_rate) | |
| cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) | |
| ret, img_cv2 = cap.read() | |
| cv2.imwrite(os.path.join(args.root, 'WHIM', args.mode, 'anno', Video_ID, frame +'.jpg' ), img_cv2.astype(np.float32)) | |
| np.save(os.path.join(args.root, 'failed_videos.npy'), failed_IDs) | |