roop/capturer.py

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.

    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)
    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
    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