52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
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
|