| import argparse | |
| from pathlib import Path | |
| import cv2 | |
| def get_parser(**parser_kwargs): | |
| parser = argparse.ArgumentParser(**parser_kwargs) | |
| parser.add_argument( | |
| "--vid_dir", | |
| required=True, | |
| type=str, | |
| nargs="?", | |
| const=True, | |
| help="directory with videos to extract frames", | |
| ) | |
| parser.add_argument( | |
| "--out_dir", | |
| type=str, | |
| nargs="?", | |
| const=True, | |
| help="directory to put extracted images", | |
| ) | |
| parser.add_argument( | |
| "--format", | |
| type=str, | |
| nargs="?", | |
| const=True, | |
| default="png", | |
| choices=["png", "jpg"], | |
| help="image file format of the extracted frames", | |
| ) | |
| parser.add_argument( | |
| "--interval", | |
| type=int, | |
| nargs="?", | |
| const=True, | |
| default=10, | |
| help="number of seconds between frame captures", | |
| ) | |
| return parser | |
| def get_videos(input_dir): | |
| for f in input_dir.iterdir(): | |
| file_path = Path(f) | |
| if file_path.suffix in [".mp4", ".avi", ".mov", ".mpeg", ".mpg", ".mkv"]: | |
| yield file_path | |
| def capture_frames(input_dir, output_dir): | |
| print (f'Capturing video frames in {opt.interval} second intervals.\n') | |
| for video_path in get_videos(input_dir): | |
| print(f'Extracting {video_path}') | |
| cap = cv2.VideoCapture(str(video_path)) | |
| if not cap.isOpened(): | |
| print(f'Could not open video') | |
| continue | |
| output = output_dir / video_path.stem | |
| output.mkdir(exist_ok=True, parents=True) | |
| current_frame = 0 | |
| count = 0 | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if ret: | |
| count_str = str(count).zfill(4) | |
| cv2.imwrite(str(output / f'frame_{count_str}.{opt.format}'), frame) | |
| current_frame += fps * opt.interval | |
| cap.set(cv2.CAP_PROP_POS_FRAMES, current_frame) | |
| count += 1 | |
| else: | |
| cap.release() | |
| break | |
| print(f'\nFinished extracting frames to {output_dir}\n') | |
| if __name__ == "__main__": | |
| parser = get_parser() | |
| opt = parser.parse_args() | |
| if (not Path(opt.vid_dir).exists): | |
| print("Video directory does not exist.") | |
| exit(1) | |
| if (opt.out_dir is None): | |
| output = Path(opt.vid_dir) / "output" | |
| print(f"No output directory specified, using default: {output}") | |
| else: | |
| output = Path(opt.out_dir) | |
| capture_frames(Path(opt.vid_dir), output) | |