tests/test_pointworld_read_clip.py

"""Tests for :meth:`fpgm.data.pointworld.FlowsReader.read_clip`.

Synthetic ``*_flows.h5`` written to ``tmp_path`` with h5py -- no network, no GPU.
Covers the two datasets added on top of the pre-existing reader
(``initial_depth``, ``scene_normals``): that they round-trip with the right
dtype when present, and that :class:`~fpgm.types.SceneFlowClip` still builds
fine (fields simply ``None``) when a file predates them, since real files in
the wild were downloaded before these fields existed on the dataclass.
"""

from __future__ import annotations

import io
from pathlib import Path

import h5py
import numpy as np
import pytest
from PIL import Image

from fpgm.data.pointworld import FlowsReader
from fpgm.types import DataError

_T = 11
_N = 17
_H, _W = 180, 320


def _jpeg_bytes(width: int = 8, height: int = 6) -> bytes:
    img = Image.fromarray(np.zeros((height, width, 3), dtype=np.uint8))
    buf = io.BytesIO()
    img.save(buf, format="JPEG")
    return buf.getvalue()


def _write_flows_h5(
    path: Path,
    *,
    with_initial_depth: bool,
    with_scene_normals: bool,
    clip_key: str = "0:11",
    camera_serial: str = "22008760",
) -> None:
    rng = np.random.default_rng(0)
    with h5py.File(path, "w") as f:
        clip_group = f.create_group(clip_key)
        cam_group = clip_group.create_group(f"camera_{camera_serial}_ext")

        cam_group.create_dataset(
            "scene_flows", data=rng.standard_normal((_T, _N, 3)).astype(np.float16)
        )
        cam_group.create_dataset(
            "scene_visibility", data=rng.integers(0, 2, size=(_T, _N)).astype(bool)
        )
        cam_group.create_dataset(
            "scene_depth_valid_mask", data=rng.integers(0, 2, size=(_T, _N)).astype(bool)
        )
        cam_group.create_dataset(
            "scene_colors", data=rng.integers(0, 256, size=(_T, _N, 3)).astype(np.uint8)
        )
        cam_group.create_dataset("intrinsic", data=np.eye(3, dtype=np.float32))
        cam_group.create_dataset("extrinsic", data=np.eye(4, dtype=np.float32))

        jpeg_dtype = h5py.special_dtype(vlen=np.dtype("uint8"))
        rgb_ds = cam_group.create_dataset("initial_rgb", shape=(1,), dtype=jpeg_dtype)
        rgb_ds[0] = np.frombuffer(_jpeg_bytes(), dtype=np.uint8)

        if with_initial_depth:
            depth_mm = rng.integers(0, 4000, size=(_H, _W)).astype(np.uint16)
            cam_group.create_dataset("initial_depth", data=depth_mm)

        if with_scene_normals:
            normals = rng.integers(-127, 128, size=(_T, _N, 3)).astype(np.int8)
            cam_group.create_dataset("scene_normals", data=normals)


class TestReadClipNewFields:
    def test_initial_depth_and_scene_normals_round_trip(self, tmp_path: Path) -> None:
        path = tmp_path / "ep_flows.h5"
        _write_flows_h5(path, with_initial_depth=True, with_scene_normals=True)

        with FlowsReader(path) as reader:
            clip = reader.read_clip("0:11", "22008760")

        assert clip.initial_depth is not None
        assert clip.initial_depth.shape == (_H, _W)
        assert clip.initial_depth.dtype == np.uint16

        assert clip.scene_normals is not None
        assert clip.scene_normals.shape == (_T, _N, 3)
        assert clip.scene_normals.dtype == np.int8

        # Round-trip against the raw h5 contents (not just shape/dtype).
        with h5py.File(path, "r") as f:
            expected_depth = np.asarray(f["0:11"]["camera_22008760_ext"]["initial_depth"])
            expected_normals = np.asarray(f["0:11"]["camera_22008760_ext"]["scene_normals"])
        assert np.array_equal(clip.initial_depth, expected_depth)
        assert np.array_equal(clip.scene_normals, expected_normals)

    def test_fields_are_none_when_datasets_absent(self, tmp_path: Path) -> None:
        """Backward compatibility: files written before these fields existed."""
        path = tmp_path / "ep_flows.h5"
        _write_flows_h5(path, with_initial_depth=False, with_scene_normals=False)

        with FlowsReader(path) as reader:
            clip = reader.read_clip("0:11", "22008760")

        assert clip.initial_depth is None
        assert clip.scene_normals is None
        # Everything else on the pre-existing contract still works unaffected.
        assert clip.scene_flows.shape == (_T, _N, 3)
        assert clip.initial_rgb is not None

    def test_partial_presence_is_independent_per_field(self, tmp_path: Path) -> None:
        path = tmp_path / "ep_flows.h5"
        _write_flows_h5(path, with_initial_depth=True, with_scene_normals=False)

        with FlowsReader(path) as reader:
            clip = reader.read_clip("0:11", "22008760")

        assert clip.initial_depth is not None
        assert clip.scene_normals is None

    def test_read_clip_requires_context_manager(self, tmp_path: Path) -> None:
        path = tmp_path / "ep_flows.h5"
        _write_flows_h5(path, with_initial_depth=True, with_scene_normals=True)

        reader = FlowsReader(path)
        with pytest.raises(DataError):
            reader.read_clip("0:11", "22008760")