diff --git a/modules/capturer.py b/modules/capturer.py index a87cf4c..952cb6d 100644 --- a/modules/capturer.py +++ b/modules/capturer.py @@ -2,19 +2,38 @@ 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_path: str, frame_number: int = 0) -> Any: - capture = cv2.VideoCapture(video_path) +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) - - frame_total = capture.get(cv2.CAP_PROP_FRAME_COUNT) - capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 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: @@ -24,8 +43,8 @@ def get_video_frame(video_path: str, frame_number: int = 0) -> Any: 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()