Deep-Live-Cam/modules/capturer.py

from typing import Any
import cv2
import modules.globals  # Import the globals to check the color correction toggle

def list_available_cameras(max_tested: int = 10):
    """ List all available camera indices. """
    available_cameras = []
    for i in range(max_tested):
        cap = cv2.VideoCapture(i)
        if cap.isOpened():
            available_cameras.append(i)
            cap.release()
    return available_cameras

def get_video_frame(video_source: Any, frame_number: int = 0, is_camera: bool = False) -> Any:
    """
    Capture a video frame from a camera or video file.
    
    :param video_source: The camera index or video file path.
    :param frame_number: Frame number to retrieve (only applicable for video files).
    :param is_camera: Flag to indicate if the source is a camera.
    :return: The captured frame.
    """
    capture = cv2.VideoCapture(video_source)

    # Set MJPEG format to ensure correct color space handling
    capture.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG'))

    # Only force RGB conversion if color correction is enabled
    if modules.globals.color_correction:
        capture.set(cv2.CAP_PROP_CONVERT_RGB, 1)

    if not is_camera:
        frame_total = capture.get(cv2.CAP_PROP_FRAME_COUNT)
        capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 1))

    has_frame, frame = capture.read()

    if has_frame and modules.globals.color_correction:
        # Convert the frame color if necessary
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    capture.release()
    return frame if has_frame else None

def get_video_frame_total(video_path: str) -> int:
    """ Get total frame count of a video file. """
    capture = cv2.VideoCapture(video_path)
    video_frame_total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    capture.release()
    return video_frame_total