Update roop/capturer.py

roop/capturer.py

@@ -1,20 +1,51 @@
-from typing import Any
+from typing import Any, Optional
 import cv2
 
+def get_video_frame(video_path: str, frame_number: int = 0) -> Optional[Any]:
+    """Retrieve a specific frame from a video.
 
-def get_video_frame(video_path: str, frame_number: int = 0) -> Any:
+    Args:
+        video_path (str): Path to the video file.
+        frame_number (int): Frame number to retrieve (0-based index). Default is 0.
+
+    Returns:
+        Optional[Any]: The requested frame as a numpy array if available, otherwise None.
+    """
     capture = cv2.VideoCapture(video_path)
+    if not capture.isOpened():
+        print(f"Error: Could not open video file {video_path}")
+        return None
+
     frame_total = capture.get(cv2.CAP_PROP_FRAME_COUNT)
-    capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 1))
+    if frame_number < 0 or frame_number >= frame_total:
+        print(f"Error: Frame number {frame_number} is out of range. Total frames: {frame_total}")
+        capture.release()
+        return None
+
+    capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
     has_frame, frame = capture.read()
     capture.release()
+    
     if has_frame:
         return frame
-    return None
-
+    else:
+        print(f"Error: Could not read frame {frame_number}")
+        return None
 
 def get_video_frame_total(video_path: str) -> int:
+    """Get the total number of frames in a video.
+
+    Args:
+        video_path (str): Path to the video file.
+
+    Returns:
+        int: Total number of frames in the video.
+    """
     capture = cv2.VideoCapture(video_path)
+    if not capture.isOpened():
+        print(f"Error: Could not open video file {video_path}")
+        return 0
+
     video_frame_total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
     capture.release()
     return video_frame_total