Create extract_frames.py

safety_analyzer/extract_frames.py

@@ -0,0 +1,70 @@
+```python
+import cv2
+import os
+import logging
+import argparse
+
+# Setup logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+def extract_frames(video_path, output_dir, frame_interval=30):
+    try:
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Output directory: {output_dir}")
+
+        video = cv2.VideoCapture(video_path)
+        if not video.isOpened():
+            raise ValueError(f"Could not open video: {video_path}")
+
+        fps = video.get(cv2.CAP_PROP_FPS)
+        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+        logger.info(f"Video: {video_path}, FPS: {fps}, Total frames: {total_frames}")
+
+        frame_count = 0
+        saved_count = 0
+        video_name = os.path.splitext(os.path.basename(video_path))[0]
+
+        while True:
+            ret, frame = video.read()
+            if not ret:
+                break
+
+            if frame_count % frame_interval == 0:
+                frame_path = os.path.join(output_dir, f"{video_name}_frame_{frame_count}.jpg")
+                cv2.imwrite(frame_path, frame)
+                saved_count += 1
+                logger.info(f"Saved frame: {frame_path}")
+
+            frame_count += 1
+
+        video.release()
+        logger.info(f"Extracted {saved_count} frames from {video_path}")
+    except Exception as e:
+        logger.error(f"Error extracting frames from {video_path}: {e}")
+        raise
+
+def main():
+    parser = argparse.ArgumentParser(description="Extract frames from videos for training dataset.")
+    parser.add_argument("--video-dir", default="videos", help="Directory containing input videos")
+    parser.add_argument("--output-dir", default="dataset/images/train", help="Directory to save extracted frames")
+    parser.add_argument("--frame-interval", type=int, default=30, help="Extract every nth frame")
+    args = parser.parse_args()
+
+    video_dir = args.video_dir
+    output_dir = args.output_dir
+    frame_interval = args.frame_interval
+
+    if not os.path.exists(video_dir):
+        logger.error(f"Video directory does not exist: {video_dir}")
+        return
+
+    for video_file in os.listdir(video_dir):
+        if video_file.endswith((".mp4", ".avi", ".mov")):
+            video_path = os.path.join(video_dir, video_file)
+            extract_frames(video_path, output_dir, frame_interval)
+
+if __name__ == "__main__":
+    logger.info("Starting frame extraction...")
+    main()
+```