Merge pull request #11 from mstevenson/main

Added video extractor arguments and documentation

doc/VIDEO_EXTRACTOR.md

@@ -0,0 +1,27 @@
+# Video frame extractor
+
+## Usage
+
+Place video files into the top level of a directory.
+
+Execute `python scripts/extract_video_frames.py --vid_dir path/to/videos` to iterate over all files, extract frames at regular intervals, and save full resolution frame images to disk.
+
+This tool supports a wide variety of input video containers and codecs (via OpenCV), and exports jpg or png files.
+
+## Arguments
+
+### --vid_dir
+
+Required directory path for input video files.
+
+### --out_dir
+
+Optional directory path in which to store extracted frame images. Defaults to a directory named 'output' that will be created inside the specified videos directory.
+
+### --format
+
+The format for image files saved to disk. Defaults to `png`, or optionally `jpg`.
+
+### --interval
+
+The number of seconds between frame captures. Defaults to 10 seconds.

scripts/extract_video_frames.py

@@ -1,42 +1,90 @@
-# extract frames from videos every n seconds
+import argparse
-
-import cv2
-import os
 from pathlib import Path
+import cv2
 
-video_base_path = Path('../movies')
+def get_parser(**parser_kwargs):
-frames_root_path = Path('../frames')
+    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
 
-interval = 30 #seconds
+def capture_frames(input_dir, output_dir):
+    print (f'Capturing video frames in {opt.interval} second intervals.\n')
 
-if not video_base_path.exists():
+    for video_path in get_videos(input_dir):
-    print('Video base path does not exist')
+        print(f'Extracting {video_path}')
-    exit(1)
+        cap = cv2.VideoCapture(str(video_path))
+        if not cap.isOpened():
+            print(f'Could not open video')
+            continue
 
-if not frames_root_path.exists():
+        output = output_dir / video_path.stem
-    os.mkdir(frames_root_path)
+        output.mkdir(exist_ok=True, parents=True)
 
-# loop over all files in the base path
+        current_frame = 0
-for video_path in video_base_path.iterdir():
+        count = 0
-    frames_folder = frames_root_path / f'frames-{video_path.stem}'
+        fps = cap.get(cv2.CAP_PROP_FPS)
-    if not frames_folder.exists():
+        while cap.isOpened():
-        os.mkdir(frames_folder)
+            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'Loading video {video_path}')
+    print(f'\nFinished extracting frames to {output_dir}\n')
-    print(f'Saving frames to {frames_folder}')
 
-    cap = cv2.VideoCapture(str(video_path))
+if __name__ == "__main__":
-    fps = cap.get(cv2.CAP_PROP_FPS)
+    parser = get_parser()
-    count = 0
+    opt = parser.parse_args()
 
-    while(cap.isOpened()):
+    if (not Path(opt.vid_dir).exists):
-        ret, frame = cap.read()
+        print("Video directory does not exist.")
-        if ret:
+        exit(1)
-            cv2.imwrite(str(frames_folder / f'frame{count}.jpg'), frame)
-            count += fps * interval
-            cap.set(cv2.CAP_PROP_POS_FRAMES, count)
-        else:
-            cap.release()
-            break
 
-print('done')
+    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)