mcap-codec-support 0.3.0__tar.gz → 0.6.0__tar.gz

{mcap_codec_support-0.3.0 → mcap_codec_support-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: mcap-codec-support
-Version: 0.3.0
+Version: 0.6.0
 Summary: Reusable MCAP encoder and decoder factories for robotics codecs
 Keywords: mcap,robotics,factories,video,pointcloud
 Author: Marko Bausch

{mcap_codec_support-0.3.0 → mcap_codec_support-0.6.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mcap-codec-support"
-version = "0.3.0"
+version = "0.6.0"
 description = "Reusable MCAP encoder and decoder factories for robotics codecs"
 readme = "README.md"
 requires-python = ">=3.10"

{mcap_codec_support-0.3.0 → mcap_codec_support-0.6.0}/src/mcap_codec_support/_protocols.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Protocol
+from typing import TYPE_CHECKING, Protocol, TypeAlias, TypeVar
 
 if TYPE_CHECKING:
     from av import VideoFrame
@@ -10,6 +10,11 @@ if TYPE_CHECKING:
 
     from mcap_codec_support.video.common import DecompressedFrame, EncoderConfig
 
+# The frame representation differs by backend: PyAV works on ``av.VideoFrame``,
+# the ffmpeg-CLI backend on raw ``bytes``. Parameterizing keeps decode→encode
+# paired per backend instead of pretending they share one frame type.
+FrameT = TypeVar("FrameT")
+
 
 class RawImageMessage(Protocol):
     """Structural shape of a ROS ``sensor_msgs/Image`` message."""
@@ -17,6 +22,7 @@ class RawImageMessage(Protocol):
     width: int
     height: int
     encoding: str
+    step: int
     data: bytes
 
 
@@ -27,12 +33,12 @@ class CompressedImageMsg(Protocol):
     def data(self) -> bytes | bytearray | memoryview: ...
 
 
-class VideoEncoderProtocol(Protocol):
-    """Structural interface shared by VideoEncoder and FFmpegVideoEncoder."""
+class VideoEncoderProtocol(Protocol[FrameT]):
+    """Encoder interface; ``FrameT`` is the per-backend frame type."""
 
     config: EncoderConfig
 
-    def encode(self, frame: VideoFrame) -> bytes | None: ...
+    def encode(self, frame: FrameT) -> bytes | None: ...
 
     def flush_packets(self) -> list[bytes]: ...
 
@@ -49,8 +55,13 @@ class VideoDecompressorProtocol(Protocol):
         ...
 
 
-class VideoCompressionBackend(Protocol):
-    """Backend used by roscompress for CompressedVideo output."""
+class VideoCompressionBackend(Protocol[FrameT]):
+    """Backend used by roscompress for CompressedVideo output.
+
+    ``FrameT`` ties ``decode_*`` output to the frame type ``create_encoder``'s
+    encoder consumes, so a backend can't decode to one frame type and encode
+    another.
+    """
 
     label: str
     prefetch_supported: bool
@@ -59,11 +70,9 @@ class VideoCompressionBackend(Protocol):
 
     def resolve_encoder(self, codec: str) -> str: ...
 
-    def decode_compressed(self, data: bytes) -> tuple[VideoFrame, int, int]: ...
+    def decode_compressed(self, data: bytes) -> tuple[FrameT, int, int]: ...
 
-    def decode_image(
-        self, msg: DecodedMessage, schema_name: str
-    ) -> tuple[VideoFrame, int, int]: ...
+    def decode_image(self, msg: DecodedMessage, schema_name: str) -> tuple[FrameT, int, int]: ...
 
     def create_encoder(
         self,
@@ -74,6 +83,11 @@ class VideoCompressionBackend(Protocol):
         *,
         input_pix_fmt: str | None = None,
         scale: tuple[int, int] | None = None,
-    ) -> VideoEncoderProtocol: ...
+    ) -> VideoEncoderProtocol[FrameT]: ...
 
     def get_pix_fmt(self, topic: str) -> str | None: ...
+
+
+# A backend chosen at runtime is either the PyAV (VideoFrame) or ffmpeg-CLI
+# (bytes) flavor; this union is the honest type at that dynamic boundary.
+AnyVideoBackend: TypeAlias = "VideoCompressionBackend[VideoFrame] | VideoCompressionBackend[bytes]"

{mcap_codec_support-0.3.0 → mcap_codec_support-0.6.0}/src/mcap_codec_support/pointcloud/compression.py

@@ -97,10 +97,12 @@ class CloudiniPointCloudCompressor:
         info = _build_encoding_info(
             msg, self._encoding_opt, self._compression_opt, self._resolution
         )
-        if self._cached_info != info:
+        encoder = self._cached_encoder
+        if self._cached_info != info or encoder is None:
             self._cached_info = info
-            self._cached_encoder = self._PointcloudEncoder(info)
-        return self._cached_encoder.encode(bytes(msg.data))  # type: ignore[union-attr]
+            encoder = self._PointcloudEncoder(info)
+            self._cached_encoder = encoder
+        return encoder.encode(bytes(msg.data))
 
 
 def _compute_position_quantization(

{mcap_codec_support-0.3.0 → mcap_codec_support-0.6.0}/src/mcap_codec_support/pointcloud/factories.py

@@ -226,7 +226,9 @@ def _decode_draco_payload(payload: bytes, header: Header) -> Pointcloud2Dict:
     dtype = np.dtype([(name, values.dtype) for name, values in columns])
     point_data = np.empty(point_count, dtype=dtype)
     for name, values in columns:
-        point_data[name] = values
+        # Structured-array field assignment; numpy's stubs type ``point_data`` as a
+        # plain float64 array, so they don't model string-key (field) assignment.
+        point_data[name] = values  # ty: ignore[invalid-assignment]
 
     fields: list[PointFieldDict] = [
         {"name": field.name, "offset": field.offset, "datatype": field.datatype, "count": 1}

{mcap_codec_support-0.3.0 → mcap_codec_support-0.6.0}/src/mcap_codec_support/video/__init__.py

@@ -1,6 +1,6 @@
 """Video MCAP factories, schema constants, and backend helpers."""
 
-from mcap_codec_support._protocols import VideoCompressionBackend
+from mcap_codec_support._protocols import AnyVideoBackend, VideoCompressionBackend
 from mcap_codec_support.video.common import (
     EncoderBackend,
     EncoderConfig,
@@ -33,6 +33,7 @@ __all__ = [
     "IMAGE",
     "IMAGE_SCHEMAS",
     "RAW_SCHEMAS",
+    "AnyVideoBackend",
     "EncoderBackend",
     "EncoderConfig",
     "EncoderMode",

{mcap_codec_support-0.3.0 → mcap_codec_support-0.6.0}/src/mcap_codec_support/video/compression.py

@@ -6,6 +6,8 @@ import io
 from collections import deque
 from typing import TYPE_CHECKING, Any
 
+import numpy as np
+
 from mcap_codec_support._schemas import normalize_schema_name
 from mcap_codec_support.video.common import (
     DEFAULT_FPS,
@@ -29,13 +31,17 @@ if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator
     from concurrent.futures import Future, ThreadPoolExecutor
 
+    import numpy.typing as npt
+    from av import VideoFrame
     from small_mcap import DecodedMessage
 
     from mcap_codec_support._protocols import (
+        AnyVideoBackend,
         RawImageMessage,
-        VideoCompressionBackend,
         VideoDecompressorProtocol,
     )
+    from mcap_codec_support.video.ffmpeg import FFmpegVideoEncoder
+    from mcap_codec_support.video.pyav import VideoEncoder
 
 
 class _PyAVCompressionBackend:
@@ -52,13 +58,13 @@ class _PyAVCompressionBackend:
 
         return resolve_encoder(codec)
 
-    def decode_compressed(self, data: bytes) -> tuple[Any, int, int]:
+    def decode_compressed(self, data: bytes) -> tuple[VideoFrame, 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]:
+    def decode_image(self, msg: DecodedMessage, schema_name: str) -> tuple[VideoFrame, int, int]:
         if schema_name in COMPRESSED_SCHEMAS:
             return self.decode_compressed(bytes(msg.decoded_message.data))
 
@@ -75,9 +81,12 @@ class _PyAVCompressionBackend:
         codec_name: str,
         quality: int,
         *,
-        input_pix_fmt: str | None = None,  # noqa: ARG002
-        scale: tuple[int, int] | None = None,  # noqa: ARG002
-    ) -> Any:
+        input_pix_fmt: str | None = None,
+        scale: tuple[int, int] | None = None,
+    ) -> VideoEncoder:
+        # PyAV reformats input frames per-frame inside VideoEncoder.encode, so
+        # the protocol's pix-fmt / scale knobs are FFmpeg-CLI-only.
+        del input_pix_fmt, scale
         from mcap_codec_support.video.pyav import VideoEncoder  # noqa: PLC0415
 
         return VideoEncoder(
@@ -114,13 +123,13 @@ class _FfmpegCliCompressionBackend:
 
         return resolve_encoder(codec)
 
-    def decode_compressed(self, data: bytes) -> tuple[Any, int, int]:
+    def decode_compressed(self, data: bytes) -> tuple[bytes, 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]:
+    def decode_image(self, msg: DecodedMessage, schema_name: str) -> tuple[bytes, int, int]:
         data = bytes(msg.decoded_message.data)
         topic = msg.channel.topic
 
@@ -147,7 +156,7 @@ class _FfmpegCliCompressionBackend:
         *,
         input_pix_fmt: str | None = None,
         scale: tuple[int, int] | None = None,
-    ) -> Any:
+    ) -> FFmpegVideoEncoder:
         from mcap_codec_support.video.ffmpeg import FFmpegVideoEncoder  # noqa: PLC0415
 
         return FFmpegVideoEncoder(
@@ -164,7 +173,7 @@ class _FfmpegCliCompressionBackend:
 
 def create_video_compression_backend(
     mode: EncoderMode, codec: str, *, do_video: bool
-) -> VideoCompressionBackend:
+) -> AnyVideoBackend:
     """Select the roscompress video backend."""
     if mode is EncoderMode.FFMPEG_CLI:
         return _FfmpegCliCompressionBackend()
@@ -180,7 +189,7 @@ def create_video_compression_backend(
 
 def prefetch_image_decodes(
     messages: Iterable[DecodedMessage],
-    backend: VideoCompressionBackend,
+    backend: AnyVideoBackend,
     pool: ThreadPoolExecutor,
     prefetch: int = 8,
 ) -> Iterator[tuple[DecodedMessage, Future[Any] | None]]:
@@ -227,11 +236,12 @@ def encode_raw_image_to_jpeg(
     return buf.getvalue(), 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."""
+def decode_compressed_image_to_rgb_array(data: bytes) -> npt.NDArray[np.uint8]:
+    """Decode JPEG/PNG compressed image bytes to an RGB (uint8) numpy array."""
     from mcap_codec_support.video.pyav import decode_compressed_frame  # noqa: PLC0415
 
-    return decode_compressed_frame(data).to_ndarray(format="rgb24")
+    rgb = decode_compressed_frame(data).to_ndarray(format="rgb24")
+    return np.asarray(rgb, dtype=np.uint8)
 
 
 def create_video_decompressor(