91 lines
2.5 KiB
Python
91 lines
2.5 KiB
Python
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)
|