mediaref 0.1.0__py3-none-any.whl

mediaref/__init__.py

@@ -0,0 +1,27 @@
+"""MediaRef - Lightweight media reference management for images and videos."""
+
+from .core import MediaRef
+
+__version__ = "0.1.0"
+__all__ = ["MediaRef"]
+
+# Optional loader module (requires extra dependencies)
+try:
+    from .loader import cleanup_cache, load_batch  # noqa: F401
+    from .video_decoder import BaseVideoDecoder, FrameBatch, PyAVVideoDecoder, TorchCodecVideoDecoder  # noqa: F401
+    from .video_decoder.types import BatchDecodingStrategy, VideoStreamMetadata  # noqa: F401
+
+    __all__.extend(
+        [
+            "load_batch",
+            "cleanup_cache",
+            "BatchDecodingStrategy",
+            "VideoStreamMetadata",
+            "BaseVideoDecoder",
+            "FrameBatch",
+            "PyAVVideoDecoder",
+            "TorchCodecVideoDecoder",
+        ]
+    )
+except ImportError:
+    pass

mediaref/_internal.py

@@ -0,0 +1,254 @@
+"""Internal loading and encoding utilities."""
+
+import base64
+import gc
+import os
+from fractions import Fraction
+from pathlib import Path
+from typing import Literal, Optional, Union
+from urllib.parse import urlparse
+
+import av
+import cv2
+import numpy as np
+import PIL.Image
+import PIL.ImageOps
+import requests
+
+from . import cached_av
+
+# Constants
+REQUEST_TIMEOUT = 60  # HTTP request timeout in seconds
+NANOSECOND = 1_000_000_000  # 1 second in nanoseconds
+
+# Garbage collection for PyAV reference cycles
+_CALLED_TIMES = 0
+GC_COLLECTION_INTERVAL = 10
+
+
+# ============================================================================
+# Image Loading
+# ============================================================================
+
+
+def load_image_as_bgra(path_or_uri: str) -> np.ndarray:
+    """Load image from any source and return as BGRA numpy array.
+
+    Args:
+        path_or_uri: File path, URL, or data URI
+
+    Returns:
+        BGRA numpy array
+
+    Raises:
+        ValueError: If loading fails
+        FileNotFoundError: If local file doesn't exist
+    """
+    try:
+        if path_or_uri.startswith("data:"):
+            return _load_from_data_uri(path_or_uri)
+        else:
+            # Load as PIL image and convert to BGRA
+            pil_image = _load_pil_image(path_or_uri)
+            return _pil_to_bgra_array(pil_image)
+    except FileNotFoundError:
+        raise
+    except Exception as e:
+        raise ValueError(f"Failed to load image from {path_or_uri}: {e}") from e
+
+
+def _load_pil_image(
+    image: Union[str, PIL.Image.Image],
+) -> PIL.Image.Image:
+    """Load image to PIL Image.
+
+    Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/utils/loading_utils.py
+    """
+    if isinstance(image, str):
+        if image.startswith("http://") or image.startswith("https://"):
+            image = PIL.Image.open(requests.get(image, stream=True, timeout=REQUEST_TIMEOUT).raw)
+        elif os.path.isfile(image):
+            image = PIL.Image.open(image)
+        else:
+            raise ValueError(
+                f"Incorrect path or URL. URLs must start with `http://` or `https://`, "
+                f"and {image} is not a valid path."
+            )
+    elif isinstance(image, PIL.Image.Image):
+        pass
+    else:
+        raise ValueError(
+            "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image."
+        )
+
+    # Handle EXIF orientation
+    image = PIL.ImageOps.exif_transpose(image)
+
+    # Convert to RGB
+    image = image.convert("RGB")
+
+    return image
+
+
+def _pil_to_bgra_array(pil_image: PIL.Image.Image) -> np.ndarray:
+    """Convert PIL image to BGRA numpy array."""
+    # Ensure image is in RGB mode
+    if pil_image.mode != "RGB":
+        pil_image = pil_image.convert("RGB")
+
+    # Convert to numpy array and then to BGRA
+    rgb_array = np.array(pil_image)
+    return cv2.cvtColor(rgb_array, cv2.COLOR_RGB2BGRA)
+
+
+# ============================================================================
+# Video Loading
+# ============================================================================
+
+
+def load_video_frame_as_bgra(
+    path_or_url: str,
+    pts_ns: int,
+    *,
+    keep_av_open: bool = False,
+) -> np.ndarray:
+    """Load video frame and return as BGRA numpy array.
+
+    Args:
+        path_or_url: File path or URL to video
+        pts_ns: Presentation timestamp in nanoseconds
+        keep_av_open: Keep AV container open in cache
+
+    Returns:
+        BGRA numpy array
+
+    Raises:
+        ValueError: If loading fails
+        FileNotFoundError: If local file doesn't exist
+    """
+    global _CALLED_TIMES
+    _CALLED_TIMES += 1
+    if _CALLED_TIMES % GC_COLLECTION_INTERVAL == 0:
+        gc.collect()
+
+    try:
+        # Validate local file exists
+        if not path_or_url.startswith(("http://", "https://")):
+            if not Path(path_or_url).exists():
+                raise FileNotFoundError(f"Video file not found: {path_or_url}")
+
+        # Convert nanoseconds to fraction
+        pts_fraction = Fraction(pts_ns, NANOSECOND)
+
+        # Open video and read frame
+        container = cached_av.open(path_or_url, "r", keep_av_open=keep_av_open)
+        try:
+            frame = _read_frame_at_pts(container, pts_fraction)
+            rgb_array = frame.to_ndarray(format="rgb24")
+            return cv2.cvtColor(rgb_array, cv2.COLOR_RGB2BGRA)
+        finally:
+            if not keep_av_open:
+                container.close()
+    except FileNotFoundError:
+        raise
+    except Exception as e:
+        pts_seconds = pts_ns / NANOSECOND
+        raise ValueError(f"Failed to load frame at {pts_seconds:.3f}s from {path_or_url}: {e}") from e
+
+
+def _read_frame_at_pts(
+    container: av.container.InputContainer,
+    pts: Fraction,
+) -> av.VideoFrame:
+    """Read single frame at or after given timestamp."""
+    if not container.streams.video:
+        raise ValueError("No video streams found")
+
+    stream = container.streams.video[0]
+
+    # Seek to the timestamp
+    container.seek(int(pts / stream.time_base), stream=stream)
+
+    # Decode frames until we find the right one
+    for frame in container.decode(stream):
+        if frame.time >= float(pts):
+            return frame
+
+    raise ValueError(f"Frame not found at {float(pts):.2f}s")
+
+
+# ============================================================================
+# Data URI Handling
+# ============================================================================
+
+
+def _load_from_data_uri(data_uri: str) -> np.ndarray:
+    """Load image from data URI."""
+    parsed = urlparse(data_uri)
+    if parsed.scheme != "data":
+        raise ValueError(f"Invalid data URI scheme: {parsed.scheme}")
+
+    try:
+        # Extract base64 data from data URI
+        data_part = parsed.path.split(",", 1)[1]
+        return _decode_from_base64(data_part)
+    except (IndexError, ValueError) as e:
+        raise ValueError(f"Invalid data URI format: {e}") from e
+
+
+def _decode_from_base64(data: str) -> np.ndarray:
+    """Decode base64 string to BGRA numpy array."""
+    try:
+        image_bytes = base64.b64decode(data)
+        nparr = np.frombuffer(image_bytes, np.uint8)
+        bgr_array = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+        if bgr_array is None:
+            raise ValueError("Failed to decode base64 image data")
+
+        return cv2.cvtColor(bgr_array, cv2.COLOR_BGR2BGRA)
+    except Exception as e:
+        raise ValueError(f"Failed to decode base64 data: {e}") from e
+
+
+# ============================================================================
+# Encoding
+# ============================================================================
+
+
+def encode_array_to_base64(
+    array: np.ndarray,
+    format: Literal["png", "jpeg", "bmp"],
+    quality: Optional[int] = None,
+) -> str:
+    """Encode BGRA numpy array to base64 string.
+
+    Args:
+        array: BGRA numpy array
+        format: Output format ('png', 'jpeg', or 'bmp')
+        quality: JPEG quality (1-100), ignored for PNG and BMP
+
+    Returns:
+        Base64 encoded string
+    """
+    # Convert BGRA to BGR for cv2 encoding
+    bgr_array = cv2.cvtColor(array, cv2.COLOR_BGRA2BGR)
+
+    # Encode based on format
+    if format == "png":
+        success, encoded = cv2.imencode(".png", bgr_array)
+    elif format == "jpeg":
+        if quality is None:
+            quality = 85
+        if not (1 <= quality <= 100):
+            raise ValueError("JPEG quality must be between 1 and 100")
+        success, encoded = cv2.imencode(".jpg", bgr_array, [cv2.IMWRITE_JPEG_QUALITY, quality])
+    elif format == "bmp":
+        success, encoded = cv2.imencode(".bmp", bgr_array)
+    else:
+        raise ValueError(f"Unsupported format: {format}")
+
+    if not success:
+        raise ValueError(f"Failed to encode image as {format}")
+
+    return base64.b64encode(encoded.tobytes()).decode("utf-8")

mediaref/_typing.py

@@ -0,0 +1,5 @@
+import os
+from typing import Union
+
+# NOTE: pathlib.Path is an instance of os.PathLike, so we don't need to include it separately
+PathLike = Union[str, bytes, os.PathLike]

mediaref/cached_av/__init__.py

@@ -0,0 +1,72 @@
+import os
+from typing import Literal, overload
+
+import av
+import av.container
+
+from .._typing import PathLike
+from ..resource_cache import ResourceCache
+from .input_container_mixin import InputContainerMixin
+
+DEFAULT_CACHE_SIZE = int(os.environ.get("AV_CACHE_SIZE", 10))
+
+# Global container cache for efficient video file access
+_container_cache: ResourceCache["MockedInputContainer"] = ResourceCache(max_size=DEFAULT_CACHE_SIZE)
+
+
+@overload
+def open(
+    file: PathLike, mode: Literal["r"], *, keep_av_open: bool = False, **kwargs
+) -> av.container.InputContainer: ...
+
+
+@overload
+def open(file: PathLike, mode: Literal["w"], **kwargs) -> av.container.OutputContainer: ...
+
+
+def open(file: PathLike, mode: Literal["r", "w"], *, keep_av_open: bool = False, **kwargs):
+    """
+    Open video container with optional caching for read operations.
+
+    Args:
+        file: Video file path or URL
+        mode: Access mode ('r' for read, 'w' for write)
+        keep_av_open: Enable caching for read containers
+        **kwargs: Additional arguments passed to av.open
+    """
+    if mode == "r":
+        if not keep_av_open:
+            # Direct access without caching
+            return av.open(file, "r", **kwargs)
+
+        # Use cached container when keep_av_open=True
+        cache_key = str(file)
+        if cache_key not in _container_cache:
+            return MockedInputContainer(file, **kwargs)
+        return _container_cache[cache_key].obj
+    else:
+        return av.open(file, mode, **kwargs)
+
+
+def cleanup_cache():
+    """Clear all cached video containers from memory."""
+    _container_cache.clear()
+
+
+class MockedInputContainer(InputContainerMixin):
+    """Cached wrapper for PyAV InputContainer with reference counting."""
+
+    def __init__(self, file: PathLike, **kwargs):
+        self._cache_key = str(file)
+        self._container: av.container.InputContainer = av.open(file, "r", **kwargs)
+        _container_cache.add_entry(self._cache_key, self)
+
+    def __enter__(self) -> "MockedInputContainer":
+        return self
+
+    def close(self):
+        """Release container reference and cleanup when no longer needed."""
+        _container_cache.release_entry(self._cache_key)
+
+
+__all__ = ["open", "cleanup_cache"]

mediaref/cached_av/container_mixin.py

@@ -0,0 +1,109 @@
+from fractions import Fraction
+from types import TracebackType
+from typing import Any, Optional, Type, Union
+
+import av.container
+from av.container.streams import StreamContainer
+from av.format import ContainerFormat
+
+Real = Union[int, float, Fraction]
+
+
+class ContainerMixin:
+    """Mixin that provides Container interface by delegating to _container.
+
+    Methods and properties copied from PyAV container core interface:
+    https://github.com/PyAV-Org/PyAV/blob/main/av/container/core.pyi#L70-L98
+    """
+
+    _container: av.container.Container
+
+    # Properties from Container (copied from PyAV container core interface)
+    @property
+    def writeable(self) -> bool:
+        return self._container.writeable
+
+    @property
+    def name(self) -> str:
+        return self._container.name
+
+    @property
+    def metadata_encoding(self) -> str:
+        return self._container.metadata_encoding
+
+    @property
+    def metadata_errors(self) -> str:
+        return self._container.metadata_errors
+
+    @property
+    def file(self) -> Any:
+        return self._container.file
+
+    @property
+    def buffer_size(self) -> int:
+        return self._container.buffer_size
+
+    @property
+    def input_was_opened(self) -> bool:
+        return self._container.input_was_opened
+
+    @property
+    def io_open(self) -> Any:
+        return self._container.io_open
+
+    @property
+    def open_files(self) -> Any:
+        return self._container.open_files
+
+    @property
+    def format(self) -> ContainerFormat:
+        return self._container.format
+
+    @property
+    def options(self) -> dict[str, str]:
+        return self._container.options
+
+    @property
+    def container_options(self) -> dict[str, str]:
+        return self._container.container_options
+
+    @property
+    def stream_options(self) -> list[dict[str, str]]:
+        return self._container.stream_options
+
+    @property
+    def streams(self) -> StreamContainer:
+        return self._container.streams
+
+    @property
+    def metadata(self) -> dict[str, str]:
+        return self._container.metadata
+
+    @property
+    def open_timeout(self) -> Optional[Real]:
+        return self._container.open_timeout
+
+    @property
+    def read_timeout(self) -> Optional[Real]:
+        return self._container.read_timeout
+
+    @property
+    def flags(self) -> int:
+        return self._container.flags
+
+    def __enter__(self) -> "ContainerMixin":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> bool:
+        return self._container.__exit__(exc_type, exc_val, exc_tb)
+
+    def set_timeout(self, timeout: Optional[Real]) -> None:
+        return self._container.set_timeout(timeout)
+
+    def start_timeout(self) -> None:
+        return self._container.start_timeout()

mediaref/cached_av/input_container_mixin.py

@@ -0,0 +1,92 @@
+from typing import Any, Iterator, Optional, Union, overload
+
+import av.container
+from av.audio.frame import AudioFrame
+from av.audio.stream import AudioStream
+from av.packet import Packet
+from av.stream import Stream
+from av.subtitles.stream import SubtitleStream
+from av.subtitles.subtitle import SubtitleSet
+from av.video.frame import VideoFrame
+from av.video.stream import VideoStream
+
+from .container_mixin import ContainerMixin
+
+
+class InputContainerMixin(ContainerMixin):
+    """Mixin that provides InputContainer interface by delegating to _container.
+
+    Methods and properties copied from PyAV input container interface:
+    https://github.com/PyAV-Org/PyAV/blob/main/av/container/input.pyi#L14-L49
+    """
+
+    _container: av.container.InputContainer  # type: ignore[override]
+
+    # InputContainer-specific properties (copied from PyAV input container interface)
+    @property
+    def start_time(self) -> int:
+        return self._container.start_time
+
+    @property
+    def duration(self) -> Optional[int]:
+        return self._container.duration
+
+    @property
+    def bit_rate(self) -> int:
+        return self._container.bit_rate
+
+    @property
+    def size(self) -> int:
+        return self._container.size
+
+    def __enter__(self) -> "InputContainerMixin":
+        return self
+
+    def close(self) -> None:
+        return self._container.close()
+
+    def demux(self, *args: Any, **kwargs: Any) -> Iterator[Packet]:
+        return self._container.demux(*args, **kwargs)
+
+    @overload
+    def decode(self, video: int) -> Iterator[VideoFrame]: ...
+
+    @overload
+    def decode(self, audio: int) -> Iterator[AudioFrame]: ...
+
+    @overload
+    def decode(self, subtitles: int) -> Iterator[SubtitleSet]: ...
+
+    @overload
+    def decode(self, *args: VideoStream) -> Iterator[VideoFrame]: ...
+
+    @overload
+    def decode(self, *args: AudioStream) -> Iterator[AudioFrame]: ...
+
+    @overload
+    def decode(self, *args: SubtitleStream) -> Iterator[SubtitleSet]: ...
+
+    def decode(self, *args: Any, **kwargs: Any) -> Iterator[Union[VideoFrame, AudioFrame, SubtitleSet]]:
+        return self._container.decode(*args, **kwargs)
+
+    def seek(
+        self,
+        offset: int,
+        *,
+        backward: bool = True,
+        any_frame: bool = False,
+        stream: Optional[Union[Stream, VideoStream, AudioStream]] = None,
+        unsupported_frame_offset: bool = False,
+        unsupported_byte_offset: bool = False,
+    ) -> None:
+        return self._container.seek(
+            offset,
+            backward=backward,
+            any_frame=any_frame,
+            stream=stream,
+            unsupported_frame_offset=unsupported_frame_offset,
+            unsupported_byte_offset=unsupported_byte_offset,
+        )
+
+    def flush_buffers(self) -> None:
+        return self._container.flush_buffers()