feat(common): load camera frames from video or image files by their own format

pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "returns>=0.26.0",
     "pyarrow<24.0.0", # NOTE: pyarrow==24.0.0 is only supported on macOS
     "pypcd4>=1.4.3,<2",
+    "av>=13.1",
 ]
 
 [dependency-groups]

t4_devkit/common/io.py

@@ -2,18 +2,21 @@
 
 import json
 import os
+import threading
 from functools import lru_cache
 from typing import Any, Sequence
 
+import av
 import numpy as np
 import pyarrow.parquet as pq
 import pypcd4
+from PIL import Image
 from typing_extensions import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from t4_devkit.typing import PathLike
 
-__all__ = ("load_json", "save_json", "load_pointcloud")
+__all__ = ("load_json", "save_json", "load_pointcloud", "load_image")
 
 
 def load_json(filename: PathLike) -> Any:
@@ -100,3 +103,60 @@ def _open_parquet(filepath: str, mtime: int, pid: int) -> pq.ParquetFile:
     so a rewritten file or a forked worker gets its own handle.
     """
     return pq.ParquetFile(filepath, memory_map=True)
+
+
+def load_image(filepath: PathLike, *, frame_index: int | None = None) -> np.ndarray:
+    """Load one camera frame as an RGB image.
+
+    With `frame_index`, the file is a video (one per channel) read with PyAV: container, codec,
+    pixel format, bit depth, color range and matrix, and timing all come from the file.
+    Without it, the file is one image per frame (`.jpg`, `.png`, ...) read with Pillow.
+
+    Args:
+        filepath (PathLike): Video file of a channel, or image file of one frame.
+        frame_index (int | None, optional): Frame to read from a video file.
+
+    Returns:
+        Image of shape (H, W, 3); uint8, or uint16 for video stored with more than 8 bits.
+    """
+    filepath = os.fspath(filepath)
+    if frame_index is None:
+        with Image.open(filepath) as image:
+            return np.asarray(image.convert("RGB"))
+
+    mtime = os.stat(filepath).st_mtime_ns
+    container = _open_video(filepath, mtime, os.getpid(), threading.get_ident())
+    stream = container.streams.video[0]
+    # Constant frame rate: frame i is shown at i / rate seconds after the first frame.
+    pts = (stream.start_time or 0) + round(frame_index / (stream.average_rate * stream.time_base))
+    # Seeks to the nearest keyframe at or before `pts` (the frame itself for intra-only video).
+    container.seek(pts, stream=stream)
+    for frame in container.decode(stream):
+        if frame.pts == pts:
+            # Decoded frames are YUV (e.g. yuv420p); convert to RGB at the stored bit depth so
+            # 8 bits per channel -> "rgb24" (uint8), more (e.g. 10-bit yuv420p10le) ->
+            # "rgb48le" (uint16). `components[0]` is the luma plane.
+            bit_depth = frame.format.components[0].bits
+            rgb_format = "rgb48le" if bit_depth > 8 else "rgb24"
+            return frame.reformat(
+                format=rgb_format,
+                src_colorspace=frame.colorspace,
+                src_color_range=frame.color_range,
+            ).to_ndarray()
+        if frame.pts > pts:
+            break
+    raise IndexError(f"Frame {frame_index} not found in {filepath}")
+
+
+@lru_cache(maxsize=32)
+def _open_video(filepath: str, mtime: int, pid: int, thread: int) -> av.container.InputContainer:
+    """Open a video file once per process, thread and file version.
+
+    A container keeps decoding state, so threads do not share one. `mtime`, `pid` and `thread`
+    are only part of the cache key.
+    """
+    container = av.open(filepath)
+    # One decoding thread: random access decodes one frame after each seek, where FFmpeg's
+    # default threading only adds latency (2880x1860 AV1: ~30 ms/frame default, ~20 ms with 1).
+    container.streams.video[0].codec_context.thread_count = 1
+    return container

tests/common/test_io.py

@@ -6,10 +6,12 @@
 import pyarrow.parquet as pq
 import pypcd4
 import pytest
+from PIL import Image
 
-from t4_devkit.common.io import load_json, load_pointcloud, save_json
+from t4_devkit.common.io import load_image, load_json, load_pointcloud, save_json
 
 SAMPLE_PARQUET = Path(__file__).parents[1] / "sample/t4dataset_v1/data/LIDAR_TOP/LIDAR_TOP.parquet"
+SAMPLE_VIDEO = Path(__file__).parents[1] / "sample/t4dataset_v1/data/CAM_FRONT/CAM_FRONT.mp4"
 
 
 def test_json_io() -> None:
@@ -73,3 +75,24 @@ def test_load_pointcloud_pcd(tmp_path: Path) -> None:
     frame = load_pointcloud(tmp_path / "frame.pcd")
     assert list(frame) == ["x", "y", "z", "intensity", "ring"]
     assert frame["ring"].dtype == np.uint16
+
+
+def test_load_image_video() -> None:
+    """Sample frames: red ramps 0..255 across the width, green is 128 * frame_index."""
+    for frame_index in (1, 0):
+        image = load_image(SAMPLE_VIDEO, frame_index=frame_index)
+        assert image.shape == (48, 64, 3)
+        assert image.dtype == np.uint8
+        assert abs(image[..., 1].mean() - 128 * frame_index) < 4
+        assert image[:, 0, 0].mean() < 8 and image[:, -1, 0].mean() > 240
+
+    with pytest.raises(IndexError):
+        load_image(SAMPLE_VIDEO, frame_index=2)
+
+
+def test_load_image_file(tmp_path: Path) -> None:
+    path = tmp_path / "frame.png"
+    pixels = np.random.default_rng(0).integers(0, 256, (4, 6, 3), dtype=np.uint8)
+    Image.fromarray(pixels).save(path)
+
+    assert np.array_equal(load_image(path), pixels)