from typing import Any, Optional import cv2 import modules.globals import logging logger = logging.getLogger(__name__) def get_video_frame(video_path: str, frame_number: int = 0) -> Optional[Any]: """ Extract a specific frame from a video file with proper color handling. Args: video_path: Path to the video file frame_number: Frame number to extract (defaults to first frame) Returns: Video frame as numpy array or None if frame extraction fails """ try: capture = cv2.VideoCapture(video_path) if not capture.isOpened(): logger.error(f"Failed to open video: {video_path}") return None # Set MJPEG format to ensure correct color space handling capture.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG')) # Configure color conversion based on setting if modules.globals.color_correction: capture.set(cv2.CAP_PROP_CONVERT_RGB, 1) else: capture.set(cv2.CAP_PROP_CONVERT_RGB, 0) # Explicitly disable if not needed 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() # Only convert manually if color_correction is enabled but capture didn't handle it if has_frame and modules.globals.color_correction and frame is not None: frame_channels = frame.shape[2] if len(frame.shape) == 3 else 1 if frame_channels == 3: # Only convert if we have a color image frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) capture.release() return frame if has_frame else None except Exception as e: logger.error(f"Error processing video frame: {str(e)}") return None def get_video_frame_total(video_path: str) -> int: """ Get the total number of frames in a video file. Args: video_path: Path to the video file Returns: Total number of frames in the video """ try: capture = cv2.VideoCapture(video_path) if not capture.isOpened(): logger.error(f"Failed to open video for frame counting: {video_path}") return 0 video_frame_total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT)) capture.release() return video_frame_total except Exception as e: logger.error(f"Error counting video frames: {str(e)}") return 0