File size: 2,573 Bytes
ddb7519 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
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)
|