Update services/video_service.py

services/video_service.py

@@ -1,6 +1,7 @@
 import cv2
-import os
 import logging
+from typing import Optional
+import numpy as np
 
 # Setup logging
 logging.basicConfig(
@@ -9,98 +10,78 @@ logging.basicConfig(
     format="%(asctime)s - %(levelname)s - %(message)s"
 )
 
-# Global state
-VIDEO_DIR = "data"
-video_files = []
-video_index = 0
-cap = None
-FRAME_SKIP = 3
-PRELOADED_FRAMES = []
-
-def update_video_files(video_path: str = None):
-    global video_files
-    if video_path and os.path.exists(video_path) and video_path.endswith((".mp4", ".avi")):
-        video_files = [video_path]
-    else:
-        if os.path.exists(VIDEO_DIR):
-            video_files = [
-                os.path.join(VIDEO_DIR, file)
-                for file in sorted(os.listdir(VIDEO_DIR))
-                if file.endswith((".mp4", ".avi"))
-            ]
-        else:
-            video_files = []
-    if not video_files:
-        logging.warning("No video files found.")
-    else:
-        logging.info(f"Found video files: {video_files}")
-
-def preload_video(video_path: str = None):
-    global cap, video_index, PRELOADED_FRAMES, video_files
-    update_video_files(video_path)
-    if not video_files:
-        logging.error("No video files available to preload.")
-        raise RuntimeError("No video files found in the 'data' directory or provided path.")
+# Global video capture object
+cap: Optional[cv2.VideoCapture] = None
+FRAME_SKIP: int = 5
+current_frame_idx: int = 0
 
+def preload_video(video_path: str) -> None:
+    """
+    Load a video file for processing.
+    Args:
+        video_path: Path to the video file.
+    Raises:
+        RuntimeError: If the video cannot be loaded.
+    """
+    global cap, current_frame_idx
+    release_video()
     try:
-        video_index = 0
-        cap = cv2.VideoCapture(video_files[video_index])
+        cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
-            raise RuntimeError(f"Failed to open video {video_files[video_index]}")
-        PRELOADED_FRAMES = []
-        for _ in range(5):
-            ret, frame = cap.read()
-            if ret:
-                frame = cv2.resize(frame, (320, 240))
-                PRELOADED_FRAMES.append(frame)
-            else:
-                break
-        cap.release()
-        cap = None
-        logging.info(f"Preloaded {len(PRELOADED_FRAMES)} frames from {video_files[video_index]}")
+            raise RuntimeError(f"Failed to open video: {video_path}. Ensure the file exists and is a supported format (.mp4, .avi).")
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        logging.info(f"Loaded video: {video_path}, Total Frames: {total_frames}, FPS: {fps}")
+        current_frame_idx = 0
     except Exception as e:
-        logging.error(f"Error preloading video: {str(e)}")
-        raise RuntimeError(f"Failed to preload video: {str(e)}")
+        logging.error(f"Error loading video {video_path}: {str(e)}")
+        raise RuntimeError(f"Error loading video {video_path}: {str(e)}")
 
-def get_next_video_frame():
-    global cap, video_index, PRELOADED_FRAMES
-    if not video_files:
-        raise RuntimeError("No video files found.")
-    
-    if PRELOADED_FRAMES:
-        frame = PRELOADED_FRAMES.pop(0)
-        logging.debug(f"Retrieved preloaded frame {len(PRELOADED_FRAMES) + 1}")
-        return frame
+def release_video() -> None:
+    """
+    Release the video capture object.
+    """
+    global cap
+    if cap is not None:
+        cap.release()
+    cap = None
 
+def reset_video_index() -> None:
+    """
+    Reset the video frame index to the beginning.
+    """
+    global current_frame_idx
+    if cap is not None:
+        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
+        current_frame_idx = 0
+        logging.info("Reset video to start.")
+
+def get_next_video_frame() -> Optional[np.ndarray]:
+    """
+    Retrieve the next frame from the video, skipping frames as specified.
+    Returns:
+        The next frame as a numpy array, or None if no frame is available.
+    """
+    global cap, current_frame_idx
     if cap is None or not cap.isOpened():
-        cap = cv2.VideoCapture(video_files[video_index])
+        logging.error("Video capture not initialized.")
+        return None
 
     try:
-        for _ in range(FRAME_SKIP):
-            ret, frame = cap.read()
-            if not ret:
-                cap.release()
-                video_index = (video_index + 1) % len(video_files)
-                cap = cv2.VideoCapture(video_files[video_index])
-                ret, frame = cap.read()
-                if not ret:
-                    raise RuntimeError(f"Cannot read video {video_files[video_index]}")
-                break
-        frame = cv2.resize(frame, (320, 240))
-        logging.debug(f"Retrieved frame from {video_files[video_index]}")
+        # Skip frames to reduce processing load
+        target_frame = current_frame_idx + FRAME_SKIP
+        cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame)
+        success, frame = cap.read()
+        if not success or frame is None:
+            # Video ended, reset to start
+            reset_video_index()
+            success, frame = cap.read()
+            if not success or frame is None:
+                logging.warning("Reached end of video and failed to reset.")
+                return None
+        current_frame_idx = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
+        logging.info(f"Retrieved frame {current_frame_idx}")
         return frame
     except Exception as e:
         logging.error(f"Error retrieving frame: {str(e)}")
-        raise RuntimeError(f"Failed to retrieve frame: {str(e)}")
-
-def reset_video_index():
-    global video_index
-    video_index = 0
-    logging.info("Reset video index to 0.")
-
-def release_video():
-    global cap
-    if cap is not None:
-        cap.release()
-        cap = None
-        logging.info("Released video capture object.")
+        return None