mcap-codec-support 0.2.0__py3-none-any.whl

mcap_codec_support/video/common.py

@@ -0,0 +1,283 @@
+"""Shared encoder types, config, and option builders — no PyAV dependency."""
+
+from __future__ import annotations
+
+import platform
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import numpy as np
+
+    from mcap_codec_support._protocols import RawImageMessage
+
+
+class VideoEncoderError(Exception):
+    """Raised when encoding fails."""
+
+
+class VideoCodec(str, Enum):
+    H264 = "h264"
+    H265 = "h265"
+    VP9 = "vp9"
+    AV1 = "av1"
+
+
+class EncoderBackend(str, Enum):
+    AUTO = "auto"
+    SOFTWARE = "software"
+    VIDEOTOOLBOX = "videotoolbox"
+    NVENC = "nvenc"
+    VAAPI = "vaapi"
+
+
+class EncoderMode(str, Enum):
+    """Which video encoder backend to use."""
+
+    AUTO = "auto"
+    PYAV = "pyav"
+    FFMPEG_CLI = "ffmpeg-cli"
+
+
+@dataclass(frozen=True, slots=True)
+class EncoderConfig:
+    """Configuration for a video encoder."""
+
+    width: int
+    height: int
+    codec_name: str
+
+
+@dataclass(frozen=True, slots=True)
+class DecompressedFrame:
+    """Result of decompressing a single video frame."""
+
+    data: bytes
+    """JPEG bytes (when ``is_jpeg=True``) or raw RGB24 bytes."""
+    width: int
+    height: int
+    is_jpeg: bool
+
+
+DEFAULT_FPS: float = 30.0
+DEFAULT_GOP_SIZE: int = 30
+
+# Mapping from short codec name to software (CPU) encoder name.
+SOFTWARE_CODEC_MAP: dict[str, str] = {
+    "h264": "libx264",
+    "h265": "libx265",
+    "vp9": "libvpx-vp9",
+    "av1": "libaom-av1",
+}
+
+# Mapping from short codec name to hardware encoder names per backend.
+HARDWARE_CODEC_MAP: dict[str, dict[str, str]] = {
+    "h264": {
+        "videotoolbox": "h264_videotoolbox",
+        "nvenc": "h264_nvenc",
+        "vaapi": "h264_vaapi",
+    },
+    "h265": {
+        "videotoolbox": "hevc_videotoolbox",
+        "nvenc": "hevc_nvenc",
+        "vaapi": "hevc_vaapi",
+    },
+}
+
+
+def get_software_encoder(codec: str) -> str:
+    """Return the software encoder name for *codec*.
+
+    Raises:
+        ValueError: If *codec* is not in SOFTWARE_CODEC_MAP.
+    """
+    sw = SOFTWARE_CODEC_MAP.get(codec)
+    if not sw:
+        raise ValueError(f"Unsupported codec '{codec}'. Supported: {', '.join(SOFTWARE_CODEC_MAP)}")
+    return sw
+
+
+def build_encoder_options(
+    codec_name: str,
+    quality: int,
+    width: int,
+    height: int,
+    *,
+    preset: str | None = None,
+) -> tuple[dict[str, str], int | None]:
+    """Build codec-specific encoder options from user-facing quality (CRF).
+
+    Returns (options_dict, bit_rate_or_none).
+    """
+    if codec_name == "libx264":
+        return {
+            "crf": str(quality),
+            "preset": preset or "superfast",
+            "tune": "zerolatency",
+        }, None
+    if codec_name == "libx265":
+        return {
+            "crf": str(quality),
+            "preset": preset or "superfast",
+        }, None
+    if codec_name in {"h264_videotoolbox", "hevc_videotoolbox"}:
+        pixel_scale = (width * height) / (1920 * 1080)
+        bit_rate = int(5_000_000 * (2 ** ((28 - quality) / 6)) * pixel_scale)
+        return {}, bit_rate
+    if codec_name in {"h264_nvenc", "hevc_nvenc"}:
+        return {"rc": "vbr", "cq": str(quality)}, None
+    if codec_name in {"h264_vaapi", "hevc_vaapi"}:
+        return {"qp": str(quality)}, None
+    if codec_name == "libvpx-vp9":
+        opts: dict[str, str] = {"cpu-used": "4", "crf": str(quality)}
+        if preset:
+            opts["deadline"] = preset
+        return opts, None
+    if codec_name == "libaom-av1":
+        opts = {"cpu-used": "6", "crf": str(quality)}
+        if preset:
+            opts["usage"] = preset
+        return opts, None
+    return {}, None
+
+
+def get_encoder_options(codec: VideoCodec, encoder_name: str) -> dict[str, str]:
+    """Return encoder preset options for bitrate-mode encoding (e.g. file output)."""
+    options: dict[str, str] = {}
+    if "nvenc" in encoder_name:
+        options["preset"] = "p4"
+    elif (
+        codec in (VideoCodec.H264, VideoCodec.H265) and "libx264" in encoder_name
+    ) or "libx265" in encoder_name:
+        options["preset"] = "medium"
+    return options
+
+
+# ---------------------------------------------------------------------------
+# Dimension helpers
+# ---------------------------------------------------------------------------
+
+
+def calculate_downscale_dimensions(width: int, height: int, max_dimension: int) -> tuple[int, int]:
+    """Downscale dimensions to fit within max_dimension, preserving aspect ratio.
+
+    Ensures both dimensions are even (required for yuv420p).
+    """
+
+    def ensure_even(value: int) -> int:
+        return value if value % 2 == 0 else max(value - 1, 2)
+
+    if width <= max_dimension and height <= max_dimension:
+        return ensure_even(width), ensure_even(height)
+
+    aspect_ratio = width / height
+    if width > height:
+        new_width = max_dimension
+        new_height = int(new_width / aspect_ratio)
+    else:
+        new_height = max_dimension
+        new_width = int(new_height * aspect_ratio)
+
+    return ensure_even(new_width), ensure_even(new_height)
+
+
+def raw_image_to_array(message: RawImageMessage) -> np.ndarray:
+    """Convert a ROS Image message to an RGB numpy array."""
+    import numpy as np  # noqa: PLC0415
+
+    if not message.data:
+        raise VideoEncoderError("Image has no data")
+
+    width = message.width
+    height = message.height
+    encoding = str(message.encoding).lower()
+    data = bytes(message.data)
+
+    if encoding in {"rgb", "rgb8"}:
+        array = np.frombuffer(data, dtype=np.uint8).reshape(height, width, 3)
+        return array.copy()
+    if encoding in {"bgr", "bgr8"}:
+        array = np.frombuffer(data, dtype=np.uint8).reshape(height, width, 3)
+        return array[..., ::-1].copy()
+    if encoding in {"mono", "mono8", "8uc1"}:
+        mono_array = np.frombuffer(data, dtype=np.uint8).reshape(height, width)
+        return np.repeat(mono_array[:, :, None], 3, axis=2)
+
+    raise VideoEncoderError(f"Unsupported image encoding: {message.encoding}")
+
+
+# ---------------------------------------------------------------------------
+# Encoder resolution
+# ---------------------------------------------------------------------------
+
+# Platform → hardware backend probe order.
+_HW_PROBE_ORDER: dict[str, list[str]] = {
+    "Darwin": ["videotoolbox"],
+    "Linux": ["nvenc", "vaapi"],
+}
+
+
+def resolve_encoder(
+    codec: str,
+    *,
+    test_fn: Callable[[str], bool],
+    use_hardware: bool = True,
+) -> str:
+    """Pick the best available encoder for *codec*.
+
+    Probes hardware encoders first (when *use_hardware* is True), then
+    falls back to the software encoder. Uses *test_fn* to check whether
+    a given encoder name is available on the system.
+
+    Raises:
+        ValueError: If no encoder is found for *codec*.
+    """
+    if use_hardware:
+        hw = HARDWARE_CODEC_MAP.get(codec)
+        if hw:
+            for backend_name in _HW_PROBE_ORDER.get(platform.system(), []):
+                encoder = hw.get(backend_name)
+                if encoder and test_fn(encoder):
+                    return encoder
+    return get_software_encoder(codec)
+
+
+def resolve_encoder_for_backend(
+    codec: str,
+    backend: str,
+    *,
+    test_fn: Callable[[str], bool],
+) -> str:
+    """Pick the encoder for *codec* using the specified *backend*.
+
+    *backend* must be one of the ``EncoderBackend`` values: ``"auto"``,
+    ``"software"``, ``"videotoolbox"``, ``"nvenc"``, or ``"vaapi"``.
+
+    Raises:
+        VideoEncoderError: If the encoder is unavailable or the backend is unknown.
+    """
+    if backend == "auto":
+        try:
+            return resolve_encoder(codec, test_fn=test_fn)
+        except ValueError as exc:
+            raise VideoEncoderError(str(exc)) from exc
+
+    if backend == "software":
+        try:
+            return get_software_encoder(codec)
+        except ValueError as exc:
+            raise VideoEncoderError(str(exc)) from exc
+
+    # Explicit hardware backend.
+    hw = HARDWARE_CODEC_MAP.get(codec, {})
+    encoder = hw.get(backend)
+    if not encoder:
+        raise VideoEncoderError(f"Hardware encoder '{backend}' not available for codec: {codec}")
+    if not test_fn(encoder):
+        raise VideoEncoderError(
+            f"Hardware encoder '{encoder}' not available on this system. Try --encoder software."
+        )
+    return encoder

mcap_codec_support/video/compression.py

@@ -0,0 +1,269 @@
+"""Video compression, image transcoding, and decompressor selection helpers."""
+
+from __future__ import annotations
+
+from collections import deque
+from typing import TYPE_CHECKING, Any
+
+from mcap_codec_support._schemas import normalize_schema_name
+from mcap_codec_support.video.common import (
+    DEFAULT_FPS,
+    DEFAULT_GOP_SIZE,
+    EncoderMode,
+    VideoEncoderError,
+    calculate_downscale_dimensions,
+    raw_image_to_array,
+)
+from mcap_codec_support.video.schemas import COMPRESSED_SCHEMAS
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator
+    from concurrent.futures import Future, ThreadPoolExecutor
+
+    from small_mcap import DecodedMessage
+
+    from mcap_codec_support._protocols import VideoCompressionBackend, VideoDecompressorProtocol
+
+
+class _PyAVCompressionBackend:
+    label = "pyav"
+    prefetch_supported = True
+
+    def test_encoder(self, encoder_name: str) -> bool:
+        from mcap_codec_support.video.pyav import test_encoder  # noqa: PLC0415
+
+        return test_encoder(encoder_name)
+
+    def resolve_encoder(self, codec: str) -> str:
+        from mcap_codec_support.video.pyav import resolve_encoder  # noqa: PLC0415
+
+        return resolve_encoder(codec)
+
+    def decode_compressed(self, data: bytes) -> tuple[Any, int, int]:
+        from mcap_codec_support.video.pyav import decode_compressed_frame  # noqa: PLC0415
+
+        frame = decode_compressed_frame(data)
+        return frame, frame.width, frame.height
+
+    def decode_image(self, msg: DecodedMessage, schema_name: str) -> tuple[Any, int, int]:
+        if schema_name in COMPRESSED_SCHEMAS:
+            return self.decode_compressed(bytes(msg.decoded_message.data))
+
+        import av  # noqa: PLC0415
+
+        rgb_array = raw_image_to_array(msg.decoded_message)
+        frame = av.VideoFrame.from_ndarray(rgb_array, format="rgb24")
+        return frame, frame.width, frame.height
+
+    def create_encoder(
+        self,
+        width: int,
+        height: int,
+        codec_name: str,
+        quality: int,
+        *,
+        input_pix_fmt: str | None = None,  # noqa: ARG002
+        scale: tuple[int, int] | None = None,  # noqa: ARG002
+    ) -> Any:
+        from mcap_codec_support.video.pyav import VideoEncoder  # noqa: PLC0415
+
+        return VideoEncoder(
+            width=width,
+            height=height,
+            codec_name=codec_name,
+            quality=quality,
+            target_fps=DEFAULT_FPS,
+            gop_size=DEFAULT_GOP_SIZE,
+        )
+
+    def get_pix_fmt(self, topic: str) -> str | None:
+        del topic
+        return None
+
+
+class _FfmpegCliCompressionBackend:
+    label = "ffmpeg-cli"
+    prefetch_supported = False
+
+    def __init__(self) -> None:
+        self._topic_pix_fmt: dict[str, str | None] = {}
+
+    def get_pix_fmt(self, topic: str) -> str | None:
+        return self._topic_pix_fmt.get(topic)
+
+    def test_encoder(self, encoder_name: str) -> bool:
+        from mcap_codec_support.video.ffmpeg import check_encoder_cli  # noqa: PLC0415
+
+        return check_encoder_cli(encoder_name)
+
+    def resolve_encoder(self, codec: str) -> str:
+        from mcap_codec_support.video.ffmpeg import resolve_encoder  # noqa: PLC0415
+
+        return resolve_encoder(codec)
+
+    def decode_compressed(self, data: bytes) -> tuple[Any, int, int]:
+        from mcap_codec_support.video.ffmpeg import probe_image_dimensions  # noqa: PLC0415
+
+        width, height = probe_image_dimensions(data)
+        return data, width, height
+
+    def decode_image(self, msg: DecodedMessage, schema_name: str) -> tuple[Any, int, int]:
+        data = bytes(msg.decoded_message.data)
+        topic = msg.channel.topic
+
+        if schema_name in COMPRESSED_SCHEMAS:
+            self._topic_pix_fmt[topic] = None
+            frame, width, height = self.decode_compressed(data)
+            return frame, width, height
+
+        from mcap_codec_support.video.ffmpeg import ROS_ENCODING_TO_PIX_FMT  # noqa: PLC0415
+
+        encoding = str(msg.decoded_message.encoding).lower()
+        pix_fmt = ROS_ENCODING_TO_PIX_FMT.get(encoding)
+        if not pix_fmt:
+            raise VideoEncoderError(f"Unsupported image encoding: {msg.decoded_message.encoding}")
+        self._topic_pix_fmt[topic] = pix_fmt
+        return data, msg.decoded_message.width, msg.decoded_message.height
+
+    def create_encoder(
+        self,
+        width: int,
+        height: int,
+        codec_name: str,
+        quality: int,
+        *,
+        input_pix_fmt: str | None = None,
+        scale: tuple[int, int] | None = None,
+    ) -> Any:
+        from mcap_codec_support.video.ffmpeg import FFmpegVideoEncoder  # noqa: PLC0415
+
+        return FFmpegVideoEncoder(
+            width=width,
+            height=height,
+            codec_name=codec_name,
+            quality=quality,
+            target_fps=DEFAULT_FPS,
+            gop_size=DEFAULT_GOP_SIZE,
+            input_pix_fmt=input_pix_fmt,
+            scale=scale,
+        )
+
+
+def create_video_compression_backend(
+    mode: EncoderMode, codec: str, *, do_video: bool
+) -> VideoCompressionBackend:
+    """Select the roscompress video backend."""
+    if mode is EncoderMode.FFMPEG_CLI:
+        return _FfmpegCliCompressionBackend()
+
+    pyav_backend = _PyAVCompressionBackend()
+    if mode is EncoderMode.AUTO and do_video:
+        try:
+            pyav_backend.resolve_encoder(codec)
+        except (ImportError, ValueError):
+            return _FfmpegCliCompressionBackend()
+    return pyav_backend
+
+
+def prefetch_image_decodes(
+    messages: Iterable[DecodedMessage],
+    backend: VideoCompressionBackend,
+    pool: ThreadPoolExecutor,
+    prefetch: int = 8,
+) -> Iterator[tuple[DecodedMessage, Future[Any] | None]]:
+    """Wrap message iterator to decode compressed images in background threads."""
+    buffer: deque[tuple[DecodedMessage, Future[Any] | None]] = deque()
+
+    for msg in messages:
+        schema_name = normalize_schema_name(msg.schema.name) if msg.schema else ""
+        if schema_name in COMPRESSED_SCHEMAS:
+            data = bytes(msg.decoded_message.data)
+            future: Future[Any] | None = pool.submit(backend.decode_compressed, data)
+        else:
+            future = None
+        buffer.append((msg, future))
+
+        if len(buffer) > prefetch:
+            yield buffer.popleft()
+
+    while buffer:
+        yield buffer.popleft()
+
+
+def _encode_rgb_array_to_jpeg(rgb_array: Any, quality: int) -> bytes:
+    try:
+        import imagecodecs  # noqa: PLC0415
+    except ImportError as exc:
+        raise VideoEncoderError(
+            "imagecodecs is required for JPEG image encoding. "
+            "Install with: uv add 'mcap-codec-support[image]'"
+        ) from exc
+    return bytes(imagecodecs.jpeg_encode(rgb_array, level=quality))
+
+
+def _resize_rgb_array(rgb_array: Any, width: int, height: int) -> Any:
+    import av  # noqa: PLC0415
+
+    frame = av.VideoFrame.from_ndarray(rgb_array, format="rgb24")
+    return frame.reformat(width=width, height=height, format="rgb24").to_ndarray(format="rgb24")
+
+
+def encode_raw_image_to_jpeg(
+    decoded_message: Any, *, jpeg_quality: int, scale: int | None
+) -> tuple[bytes, int, int]:
+    """Encode a raw ROS Image message to JPEG using imagecodecs for final encode."""
+    rgb_array = raw_image_to_array(decoded_message)
+    src_h, src_w = rgb_array.shape[:2]
+    if scale is not None:
+        target_w, target_h = calculate_downscale_dimensions(src_w, src_h, scale)
+    else:
+        target_w, target_h = src_w, src_h
+
+    target_w -= target_w % 2
+    target_h -= target_h % 2
+    if target_w < 2 or target_h < 2:
+        raise VideoEncoderError(f"Source frame too small ({target_w}x{target_h}) for JPEG encoding")
+
+    if target_w != src_w or target_h != src_h:
+        rgb_array = _resize_rgb_array(rgb_array, target_w, target_h)
+
+    return _encode_rgb_array_to_jpeg(rgb_array, jpeg_quality), target_w, target_h
+
+
+def decode_compressed_image_to_rgb_array(data: bytes) -> Any:
+    """Decode JPEG/PNG compressed image bytes to an RGB numpy array."""
+    from mcap_codec_support.video.pyav import decode_compressed_frame  # noqa: PLC0415
+
+    return decode_compressed_frame(data).to_ndarray(format="rgb24")
+
+
+def create_video_decompressor(
+    video_format: str = "compressed",
+    jpeg_quality: int = 90,
+    *,
+    mode: EncoderMode = EncoderMode.AUTO,
+) -> VideoDecompressorProtocol:
+    """Create a video decompressor using the requested backend."""
+    if mode == EncoderMode.PYAV:
+        from mcap_codec_support.video.pyav import PyAVVideoDecompressor  # noqa: PLC0415
+
+        return PyAVVideoDecompressor(video_format=video_format, jpeg_quality=jpeg_quality)
+
+    if mode == EncoderMode.FFMPEG_CLI:
+        from mcap_codec_support.video.ffmpeg import FFmpegVideoDecompressor  # noqa: PLC0415
+
+        return FFmpegVideoDecompressor(video_format=video_format, jpeg_quality=jpeg_quality)
+
+    try:
+        from mcap_codec_support.video.pyav import PyAVVideoDecompressor  # noqa: PLC0415
+
+        return PyAVVideoDecompressor(video_format=video_format, jpeg_quality=jpeg_quality)
+    except ImportError:
+        from mcap_codec_support.video.ffmpeg import (  # noqa: PLC0415
+            FFmpegVideoDecompressor,
+            find_ffmpeg,
+        )
+
+        if find_ffmpeg():
+            return FFmpegVideoDecompressor(video_format=video_format, jpeg_quality=jpeg_quality)
+        raise

mcap_codec_support/video/factories.py

@@ -0,0 +1,123 @@
+"""Decoder factory for decompressing CompressedVideo topics.
+
+Provides ``VideoDecompressFactory`` for use with ``read_message_decoded``.
+Uses ``VideoDecompressorProtocol`` — no direct ``av`` or ``subprocess`` imports.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from mcap_codec_support.video.common import EncoderMode
+from mcap_codec_support.video.schemas import COMPRESSED_VIDEO_SCHEMA
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from small_mcap import Channel, Schema
+
+    from mcap_codec_support._messages import CompressedImageDict, Header, ImageDict
+    from mcap_codec_support._protocols import VideoDecompressorProtocol
+    from mcap_codec_support.video.common import DecompressedFrame
+
+
+# ---------------------------------------------------------------------------
+# VideoDecompressFactory
+# ---------------------------------------------------------------------------
+
+
+class VideoDecompressFactory:
+    """Channel-aware decoder factory: CompressedVideo → CompressedImage or Image.
+
+    Creates a separate ``VideoDecompressorProtocol`` per channel for proper
+    P-frame handling. No direct ``av`` or ``subprocess`` imports.
+    """
+
+    channel_aware = True
+
+    def __init__(
+        self,
+        *,
+        video_format: Literal["compressed", "raw"] = "compressed",
+        jpeg_quality: int = 90,
+        backend: EncoderMode = EncoderMode.AUTO,
+    ) -> None:
+        self._video_format = video_format
+        self._jpeg_quality = jpeg_quality
+        self._backend = backend
+
+        from mcap_ros2_support_fast.decoder import DecoderFactory  # noqa: PLC0415
+
+        self._cdr_factory = DecoderFactory()
+        self._decompressors: dict[int, VideoDecompressorProtocol] = {}
+
+    def flush_all(self) -> list[DecompressedFrame]:
+        """Flush all decompressors and return remaining frames."""
+        return [frame for _, frame in self.flush_all_by_channel()]
+
+    def flush_all_by_channel(self) -> list[tuple[int, DecompressedFrame]]:
+        """Flush all decompressors and keep channel ownership for each frame."""
+        frames: list[tuple[int, DecompressedFrame]] = []
+        for channel_id, decompressor in self._decompressors.items():
+            frames.extend((channel_id, frame) for frame in decompressor.flush())
+        return frames
+
+    def _get_decompressor(self, channel_id: int) -> VideoDecompressorProtocol:
+        if channel_id not in self._decompressors:
+            from mcap_codec_support.video.compression import (  # noqa: PLC0415
+                create_video_decompressor,
+            )
+
+            self._decompressors[channel_id] = create_video_decompressor(
+                video_format=self._video_format,
+                jpeg_quality=self._jpeg_quality,
+                mode=self._backend,
+            )
+        return self._decompressors[channel_id]
+
+    def decoder_for(
+        self,
+        message_encoding: str,
+        schema: Schema | None,
+        channel: Channel,
+    ) -> Callable[[bytes | memoryview], CompressedImageDict | ImageDict | None] | None:
+        if schema is None or schema.name != COMPRESSED_VIDEO_SCHEMA:
+            return None
+
+        cdr_decoder = self._cdr_factory.decoder_for(message_encoding, schema)
+        if cdr_decoder is None:
+            return None
+
+        decompressor = self._get_decompressor(channel.id)
+
+        def _decode(data: bytes | memoryview) -> CompressedImageDict | ImageDict | None:
+            decoded = cdr_decoder(data)
+            codec = decoded.format
+            video_data = decoded.data
+            if isinstance(video_data, memoryview):
+                video_data = bytes(video_data)
+
+            frame: DecompressedFrame | None = decompressor.decompress(video_data, codec)
+            if frame is None:
+                return None
+
+            timestamp = decoded.timestamp
+            header: Header = {
+                "stamp": {"sec": timestamp.sec, "nanosec": timestamp.nanosec},
+                "frame_id": decoded.frame_id,
+            }
+
+            if frame.is_jpeg:
+                return {"header": header, "format": "jpeg", "data": frame.data}
+
+            return {
+                "header": header,
+                "height": frame.height,
+                "width": frame.width,
+                "encoding": "rgb8",
+                "is_bigendian": 0,
+                "step": frame.width * 3,
+                "data": frame.data,
+            }
+
+        return _decode