shared/frame_utils.py

"""Utilities for extracting and displaying frames from images and videos.

The ``gradio`` import is deferred so that Gradio-free functions
(``load_media_frames_raw``) can be used inside worker subprocesses
without pulling in the full Gradio dependency.
"""

import base64

import cv2
import numpy as np

_VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm")


def load_media_frames_raw(media_path: str) -> list[np.ndarray]:
    """Load BGR frames from a media file (image or video).

    Gradio-free — raises ``RuntimeError`` on failure.  Safe to call inside
    worker subprocesses that do not have Gradio installed / imported.

    Args:
        media_path: Path to an image or video file.

    Returns:
        List of BGR numpy arrays (one per frame).

    Raises:
        RuntimeError: If the file cannot be opened or contains no frames.
    """
    if media_path.lower().endswith(_VIDEO_EXTENSIONS):
        cap = cv2.VideoCapture(media_path)
        if not cap.isOpened():
            raise RuntimeError(f"Could not open video: {media_path}")
        frames: list[np.ndarray] = []
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frames.append(frame)
        cap.release()
        if not frames:
            raise RuntimeError(f"No frames read from video: {media_path}")
        return frames

    img = cv2.imread(media_path)
    if img is None:
        raise RuntimeError(f"Could not read image: {media_path}")
    return [img]


def extract_frames(image_path: str | None, video_path: str | None) -> list[np.ndarray]:
    """Extract BGR frames from an image or video.

    Args:
        image_path: Path to an image file, or None.
        video_path: Path to a video file, or None.

    Returns:
        List of BGR numpy arrays.

    Raises:
        gr.Error: If neither input is provided or the media cannot be read.
    """
    import gradio as gr

    if image_path is None and video_path is None:
        raise gr.Error("Please upload an image or video.")

    media_path = image_path or video_path
    assert media_path is not None  # guaranteed by the check above
    try:
        return load_media_frames_raw(media_path)
    except RuntimeError as exc:
        raise gr.Error(str(exc)) from exc


def load_media_frames(media_path: str) -> list[np.ndarray]:
    """Load BGR frames from a stored media file (image or video).

    Gradio-aware wrapper around :func:`load_media_frames_raw` that converts
    ``RuntimeError`` to ``gr.Error`` for UI display.

    Args:
        media_path: Path to the media file.

    Returns:
        List of BGR numpy arrays.
    """
    import gradio as gr

    try:
        return load_media_frames_raw(media_path)
    except RuntimeError as exc:
        raise gr.Error(str(exc)) from exc


def get_thumbnail(media_path: str) -> np.ndarray | None:
    """Get an RGB thumbnail from a media file for display.

    For videos, returns the first frame. For images, returns the image itself.

    Args:
        media_path: Path to the media file.

    Returns:
        RGB numpy array suitable for gr.Image display, or None on failure.
    """
    if media_path.lower().endswith(_VIDEO_EXTENSIONS):
        cap = cv2.VideoCapture(media_path)
        ret, frame = cap.read()
        cap.release()
        if ret:
            return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        return None
    img = cv2.imread(media_path)
    if img is not None:
        return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return None


def get_thumbnail_base64(media_path: str) -> str | None:
    """Return a base64-encoded JPEG string of the first frame from a media file.

    Args:
        media_path: Path to an image or video file.

    Returns:
        Base64-encoded JPEG string, or None on failure.
    """
    thumb = get_thumbnail(media_path)
    if thumb is None:
        return None
    bgr = cv2.cvtColor(thumb, cv2.COLOR_RGB2BGR)
    ok, buf = cv2.imencode(".jpg", bgr, [cv2.IMWRITE_JPEG_QUALITY, 80])
    if not ok:
        return None
    return base64.b64encode(buf.tobytes()).decode()


def get_video_frame_count(video_path: str) -> int:
    """Return the total number of frames in a video file.

    Reads ``CAP_PROP_FRAME_COUNT`` from the container metadata.  Falls back
    to a sequential scan when the metadata is missing or unreliable (common
    with browser-recorded WebM files).

    Args:
        video_path: Path to a video file.

    Returns:
        Total frame count (always >= 1).

    Raises:
        RuntimeError: If the video cannot be opened or contains no frames.
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"Could not open video: {video_path}")

    count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if count > 0:
        cap.release()
        return count

    # Metadata unreliable — count by reading through
    count = 0
    while True:
        ret, _ = cap.read()
        if not ret:
            break
        count += 1
    cap.release()
    if count == 0:
        raise RuntimeError(f"No frames read from video: {video_path}")
    return count


def extract_frame_at_index(video_path: str, frame_index: int) -> np.ndarray:
    """Extract a single BGR frame at the given index from a video.

    Uses ``CAP_PROP_POS_FRAMES`` to seek directly to the requested frame,
    avoiding the need to decode the entire video.

    Args:
        video_path: Path to a video file.
        frame_index: Zero-based frame index.

    Returns:
        BGR numpy array of the requested frame.

    Raises:
        RuntimeError: If the video cannot be opened or the frame cannot be read.
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"Could not open video: {video_path}")
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
    ret, frame = cap.read()
    cap.release()
    if not ret:
        raise RuntimeError(f"Could not read frame {frame_index} from: {video_path}")
    return frame