test_image_dataset.py

"""
test_image_dataset.py — unit tests for image_dataset.py (pure CPU, no sim/render).

Synthetic HDF5 encodes frame identity directly in pixel values so windows,
padding, alignment, and cross-episode isolation are all checkable exactly.

Run: python test_image_dataset.py
"""

from __future__ import annotations

import tempfile
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader

from datasets import WindowConfig, sample_sequence
from image_dataset import (
    ImageSequenceDataset,
    read_padded_image_window,
    flat_to_local,
    worker_init_fn,
)

CAMS = ("camA", "camB")
H = W = 84


def _pixval(e: int, t: int, c: int) -> int:
    """Unique-ish identity value for (episode e, timestep t, camera index c)."""
    return (10 * e + t + 100 * c) % 256


def make_synthetic_image_hdf5(path, ep_lens=(20, 30), cams=CAMS):
    import h5py
    with h5py.File(path, "w") as f:
        data = f.create_group("data")
        names = []
        for e, L in enumerate(ep_lens):
            name = f"demo_{e}"
            names.append(name)
            g = data.create_group(name)
            g.attrs["num_samples"] = L
            obs = g.create_group("obs")
            for c, cam in enumerate(cams):
                arr = np.empty((L, H, W, 3), dtype=np.uint8)
                for t in range(L):
                    arr[t] = _pixval(e, t, c)
                obs.create_dataset(f"{cam}_image", data=arr, chunks=(1, H, W, 3))
            # proprio: col0 = episode, col1 = timestep, rest arange-ish (9 dims total)
            prop_pos = np.stack([np.full(L, e), np.arange(L), np.arange(L) * 0.1], axis=1)
            obs.create_dataset("robot0_eef_pos", data=prop_pos.astype(np.float32))
            obs.create_dataset("robot0_eef_quat",
                               data=np.tile([0, 0, 0, 1], (L, 1)).astype(np.float32))
            obs.create_dataset("robot0_gripper_qpos",
                               data=np.zeros((L, 2), dtype=np.float32))
            g.create_dataset("actions", data=np.arange(L * 7).reshape(L, 7).astype(np.float32))
            g.create_dataset("states", data=np.zeros((L, 5), dtype=np.float32))
        mask = f.create_group("mask")
        mask.create_dataset("train", data=np.array(names, dtype="S"))
    return names


def _decode(img_chw_float):
    """(3,h,w) float [0,1] constant frame -> integer identity value."""
    return int(round(float(img_chw_float[0, 0, 0].item()) * 255.0))


def test_shapes_dtypes_ranges(path):
    cfg = WindowConfig(pred_horizon=16, obs_horizon=2, action_horizon=8)
    ds = ImageSequenceDataset(path, CAMS, cfg=cfg, crop=76, train=True, filter_key="train")
    s = ds[len(ds) // 2]
    for cam in CAMS:
        assert s["images"][cam].shape == (2, 3, 76, 76), s["images"][cam].shape
        assert s["images"][cam].dtype == torch.float32
        assert 0.0 <= float(s["images"][cam].min()) and float(s["images"][cam].max()) <= 1.0
    assert s["proprio"].shape == (2, 9), s["proprio"].shape
    assert s["action"].shape == (16, 7), s["action"].shape
    assert float(s["proprio"].abs().max()) <= 1.0 + 1e-6
    assert float(s["action"].abs().max()) <= 1.0 + 1e-6
    ds.close()
    print("  ok: shapes/dtypes/ranges")


def test_identity_alignment(path):
    cfg = WindowConfig(16, 2, 8)
    ds = ImageSequenceDataset(path, CAMS, cfg=cfg, train=False, filter_key="train")
    # find a mid-episode window (no head padding): ss == 0 and bs well inside ep 1
    for i in range(len(ds)):
        bs, be, ss, se = (int(v) for v in ds.indices[i])
        ep, local = flat_to_local(bs, ds.episode_ends)
        if ss == 0 and ep == 1 and local >= 3:
            break
    s = ds[i]
    for k in range(cfg.obs_horizon):
        gidx = bs + k
        e = int(round(ds.proprio_raw[gidx, 0]))   # episode from proprio col0
        t = int(round(ds.proprio_raw[gidx, 1]))   # timestep from proprio col1
        for c, cam in enumerate(CAMS):
            got = _decode(s["images"][cam][k])
            assert got == _pixval(e, t, c), (cam, k, got, _pixval(e, t, c))
    ds.close()
    print("  ok: image (e,t) matches proprio (e,t) for both cameras")


def test_start_of_episode_padding(path):
    cfg = WindowConfig(16, 2, 8)
    ds = ImageSequenceDataset(path, CAMS, cfg=cfg, train=False, filter_key="train")
    # first window of episode 0 has pad_before -> ss > 0, frame 0 repeated
    i0 = 0
    bs, be, ss, se = (int(v) for v in ds.indices[i0])
    assert ss > 0, "expected head padding on first window"
    s = ds[i0]
    for cam_i, cam in enumerate(CAMS):
        v0 = _decode(s["images"][cam][0])
        v1 = _decode(s["images"][cam][1])
        assert v0 == _pixval(0, 0, cam_i), (v0, _pixval(0, 0, cam_i))  # padded == frame 0
        assert v1 == _pixval(0, 0, cam_i), (v1, _pixval(0, 0, cam_i))  # real frame 0
    ds.close()
    print("  ok: start-of-episode padding repeats frame 0")


def test_no_cross_episode_bleed(path):
    cfg = WindowConfig(16, 2, 8)
    ds = ImageSequenceDataset(path, CAMS, cfg=cfg, train=False, filter_key="train")
    for i in range(len(ds)):
        bs, be, ss, se = (int(v) for v in ds.indices[i])
        ep, _ = flat_to_local(bs, ds.episode_ends)
        # obs frames must all belong to this window's episode (proprio identity)
        for k in range(cfg.obs_horizon):
            gidx = min(bs + k, be - 1)
            e = int(round(ds.proprio_raw[gidx, 0]))
            assert e == ep, (i, k, e, ep)
    ds.close()
    print("  ok: no cross-episode bleed in obs windows")


def test_crop_modes(path):
    cfg = WindowConfig(16, 2, 8)
    # eval: deterministic center crop
    ds_eval = ImageSequenceDataset(path, CAMS, cfg=cfg, train=False, filter_key="train")
    a = ds_eval[5]["images"]["camA"]
    b = ds_eval[5]["images"]["camB"]  # different read
    a2 = ds_eval[5]["images"]["camA"]
    assert torch.equal(a, a2), "eval center crop must be deterministic"
    # within one sample, both To frames share crop offset (constant-value frames -> trivially equal;
    # verify shape is the crop size)
    assert a.shape == (2, 3, 76, 76)
    ds_eval.close()

    # train: offsets vary across samples (statistically); constant across To within a sample.
    # Constant-valued frames can't reveal spatial offset, so instead check the raw offset path:
    ds_train = ImageSequenceDataset(path, CAMS, cfg=cfg, train=True, seed=0, filter_key="train")
    from image_dataset import random_crop_params
    rng = np.random.default_rng(0)
    offs = {random_crop_params(rng, H, W, 76, 76) for _ in range(50)}
    assert len(offs) > 1, "random crop should produce varied offsets"
    ds_train.close()
    print("  ok: eval center crop deterministic; train crop offsets vary")


def test_worker_safety(path):
    cfg = WindowConfig(16, 2, 8)

    def collect(num_workers):
        ds = ImageSequenceDataset(path, CAMS, cfg=cfg, train=False, filter_key="train")
        dl = DataLoader(ds, batch_size=4, shuffle=False, num_workers=num_workers,
                        worker_init_fn=worker_init_fn if num_workers > 0 else None,
                        persistent_workers=False)
        out = []
        for batch in dl:
            out.append(batch["images"]["camA"].clone())
        ds.close()
        return torch.cat(out, dim=0)

    a = collect(0)
    b = collect(2)
    assert a.shape == b.shape
    assert torch.equal(a, b), (a - b).abs().max().item()
    print("  ok: num_workers=2 matches num_workers=0 (center crop)")


def test_read_padded_equiv_sample_sequence():
    rng = np.random.default_rng(1)
    from datasets import create_sample_indices
    L, n = 12, 5
    data = rng.integers(0, 255, size=(L, 2, 2, 3)).astype(np.uint8)  # single episode
    ends = np.array([L], dtype=np.int64)
    idxs = create_sample_indices(ends, n, pad_before=n - 1, pad_after=n - 1)
    for bs, be, ss, se in idxs:
        want = sample_sequence(data, n, int(bs), int(be), int(ss), int(se))
        got = read_padded_image_window(data, int(bs), int(be), int(ss), n)
        assert np.array_equal(want, got), (bs, be, ss, se)
    print(f"  ok: read_padded_image_window == sample_sequence over {len(idxs)} windows")


if __name__ == "__main__":
    with tempfile.TemporaryDirectory() as d:
        path = str(Path(d) / "synthetic_image.hdf5")
        make_synthetic_image_hdf5(path)
        tests = [
            ("shapes_dtypes_ranges", lambda: test_shapes_dtypes_ranges(path)),
            ("identity_alignment", lambda: test_identity_alignment(path)),
            ("start_of_episode_padding", lambda: test_start_of_episode_padding(path)),
            ("no_cross_episode_bleed", lambda: test_no_cross_episode_bleed(path)),
            ("crop_modes", lambda: test_crop_modes(path)),
            ("worker_safety", lambda: test_worker_safety(path)),
            ("read_padded_equiv_sample_sequence", test_read_padded_equiv_sample_sequence),
        ]
        for name, fn in tests:
            print(f"[{name}]")
            fn()
    print("\nALL IMAGE DATASET TESTS PASSED")