From 4522715a60f5df564520ade123f2d2b10f5503aa Mon Sep 17 00:00:00 2001 From: Jacob Lambert Date: Sat, 10 Oct 2026 22:58:54 +0900 Subject: [PATCH] feat(common): load camera frames from video or image files by their own format --- pyproject.toml | 1 + t4_devkit/common/io.py | 62 ++++++++++++++++++++++++++++++++++++++++- tests/common/test_io.py | 25 ++++++++++++++++- 3 files changed, 86 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b04a8855..3de1b2af 100644 --- a/pyproject.toml +++ b/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] diff --git a/t4_devkit/common/io.py b/t4_devkit/common/io.py index fc9e5c04..4c74c784 100644 --- a/t4_devkit/common/io.py +++ b/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 diff --git a/tests/common/test_io.py b/tests/common/test_io.py index 6aba84a0..75bf2af2 100644 --- a/tests/common/test_io.py +++ b/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)