code/tt_meteor/tests/test_host_preprocess.py

# SPDX-License-Identifier: Apache-2.0
"""Host pre-processing (no device; numpy + PIL): the PandaSet 019 frame 40 sample's graph feed (staged, not shipped:
``tests/paths.py``) equals the research golden's byte for byte through every input path (scene reader, Python-API request, absent-camera donor rule), and the INTER_AREA
resize / K scaling / extrinsics inversion follow METEOR's runtimes.

    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_meteor/tests/test_host_preprocess.py
"""
from __future__ import annotations


import numpy as np
import pytest

from tt_meteor import io as tio
from tt_meteor.api import load_sample
from tt_meteor.host.inputs import prepare_request
from tt_meteor.host.preprocess import (assemble_frame, invert_extrinsics, load_scene_frame, normalize_images,
                                       prepare_camera, scene_dir_of)
from tt_meteor.reference import config as C
from tt_meteor.tests import paths as P
from tt_meteor.ttaw.image_area import meteor_resize

SAMPLE_ROOT = P.PANDASET_SCENE
SAMPLE_JSON = P.PANDASET_SAMPLE


@pytest.fixture
def GOLD():
    return P.pandaset_golden("inputs")


def test_scene_frame_matches_the_golden_feed(GOLD):
    frame = load_scene_frame(scene_dir_of(SAMPLE_ROOT), 0)
    assert frame.sha256() == GOLD["input_sha256"]
    assert frame.imgs.shape == (1, 8, 3, 432, 768) and frame.imgs.dtype == np.uint8
    assert frame.K.dtype == frame.T_cam_ego.dtype == frame.v0.dtype == np.float32
    assert frame.present.tolist() == [True] * 7 + [False] and not frame.imgs[0, 7].any()
    assert frame.v0[0] == pytest.approx(10.8792, abs=1e-4) and frame.pose is not None
    assert len(frame.meta["wp"]) == 6 and frame.meta["valid"] == 1.0


def test_api_request_matches_the_golden_feed(GOLD):
    kw = load_sample(SAMPLE_JSON)
    frame = prepare_request(kw["images"], kw["calibration"], kw["ego_speed"], kw["stream"])
    assert frame.sha256() == GOLD["input_sha256"]
    assert frame.meta["donors"] == {"CAM_BACK_NARROW": "CAM_BACK_WIDE"}    # the all-zero image = absent
    scene = load_scene_frame(scene_dir_of(SAMPLE_ROOT), 0)
    assert np.allclose(frame.pose, scene.pose, atol=1e-6)                  # T_world_from_ego (x, y, yaw) round trip
    # the narrow camera omitted: zero image + the donor's K and pose (the trained 7-camera rig)
    cams = [c for c in kw["images"] if c.name != "CAM_BACK_NARROW"]
    assert prepare_request(cams, kw["calibration"], kw["ego_speed"]).sha256() == GOLD["input_sha256"]
    # a zero narrow image with another calibration still takes the donor's (absent convention)
    other = [tio.CameraImage(c.name, c.image, intrinsics=c.intrinsics * 2 if c.name == "CAM_BACK_NARROW" else
                             c.intrinsics, T_ref_from_camera=c.T_ref_from_camera) for c in kw["images"]]
    assert prepare_request(other, kw["calibration"], kw["ego_speed"]).sha256() == GOLD["input_sha256"]


def test_resize_and_intrinsics_scaling():
    rng = np.random.default_rng(1)
    big = rng.integers(0, 256, (1080, 1920, 3), dtype=np.uint8)
    K = np.array([[1000.0, 0.0, 960.0], [0.0, 1000.0, 540.0], [0.0, 0.0, 1.0]])
    chw, k, src = prepare_camera(big, K)
    assert src == (1080, 1920) and chw.shape == (3, 432, 768)
    np.testing.assert_array_equal(chw.transpose(1, 2, 0), meteor_resize(big))
    np.testing.assert_allclose(k[:2], K[:2] * 0.4)
    cv2 = pytest.importorskip("cv2")
    np.testing.assert_array_equal(chw.transpose(1, 2, 0), cv2.resize(big, (768, 432), interpolation=cv2.INTER_AREA))
    # anisotropic (METEOR's 2880x1860 training frames): fx != fy afterwards
    _, k2, _ = prepare_camera(rng.integers(0, 256, (1860, 2880, 3), dtype=np.uint8), K)
    assert k2[0, 0] == pytest.approx(1000 * 768 / 2880) and k2[1, 1] == pytest.approx(1000 * 432 / 1860)


def test_assemble_frame_and_donors():
    rng = np.random.default_rng(2)
    K = np.array([[400.0, 0.0, 384.0], [0.0, 400.0, 216.0], [0.0, 0.0, 1.0]])
    T = {c: np.eye(4) for c in C.CAMERAS}
    for i, c in enumerate(C.CAMERAS):
        T[c] = T[c].copy()
        T[c][0, 3] = float(i)
    imgs = {c: rng.integers(1, 256, (432, 768, 3), dtype=np.uint8) for c in C.CAMERAS[:6]}
    f = assemble_frame(imgs, {c: K for c in C.CAMERAS}, T, 5.0)
    assert f.present.tolist() == [True] * 6 + [False, False]
    np.testing.assert_array_equal(f.K[0, 6], f.K[0, 0])                        # FRONT_NARROW <- FRONT_WIDE
    np.testing.assert_array_equal(f.T_cam_ego[0, 7], f.T_cam_ego[0, 3])        # BACK_NARROW <- BACK_WIDE
    assert not f.imgs[0, 6:].any() and f.meta["donors"] == {"CAM_FRONT_NARROW": "CAM_FRONT_WIDE",
                                                             "CAM_BACK_NARROW": "CAM_BACK_WIDE"}
    with pytest.raises(ValueError, match="no present donor"):
        assemble_frame({c: v for c, v in imgs.items() if c != "CAM_FRONT_LEFT"}, {c: K for c in C.CAMERAS}, T, 5.0)


@pytest.mark.parametrize("bad", ["small", "missing", "no_speed", "nan_speed", "no_calib"])
def test_request_errors(bad):
    kw = load_sample(P.SHIPPED_SAMPLE)
    images, calib, speed = list(kw["images"]), kw["calibration"], kw["ego_speed"]
    if bad == "small":
        images[0] = tio.CameraImage(images[0].name, np.zeros((400, 700, 3), np.uint8), images[0].intrinsics,
                                    images[0].T_ref_from_camera)
    elif bad == "missing":
        images = images[1:]
    elif bad == "no_speed":
        speed = None
    elif bad == "nan_speed":
        speed = float("nan")
    elif bad == "no_calib":
        images = [{"camera": c.name, "image": c.image} for c in images]
        calib = {}
    with pytest.raises(tio.InputError):
        prepare_request(images, calib, speed)


def test_invert_extrinsics_modes():
    a = np.radians(30.0)
    T = np.array([[np.cos(a), -np.sin(a), 0, 1.5], [np.sin(a), np.cos(a), 0, -0.2], [0, 0, 1, 1.9], [0, 0, 0, 1]])
    i32, i64 = invert_extrinsics(T, "float32"), invert_extrinsics(T, "float64")
    assert i32.dtype == i64.dtype == np.float32
    np.testing.assert_allclose(i32 @ T.astype(np.float32), np.eye(4), atol=1e-6)
    np.testing.assert_allclose(i32, i64, atol=1e-6)
    with pytest.raises(ValueError):
        invert_extrinsics(T, "float16")


def test_normalize_images():
    x = np.random.default_rng(3).integers(0, 256, (1, 8, 3, 432, 768), dtype=np.uint8)
    np.testing.assert_array_equal(normalize_images(x), x.reshape(8, 3, 432, 768).astype(np.float32) / np.float32(255))
    present = [True] * 7 + [False]
    y = normalize_images(x, "imagenet", present)
    mean = np.asarray(C.IMAGENET_MEAN, np.float32).reshape(3, 1, 1)
    std = np.asarray(C.IMAGENET_STD, np.float32).reshape(3, 1, 1)
    np.testing.assert_allclose(y[0], (x[0, 0] / np.float32(255) - mean) / std, rtol=1e-6)
    assert not y[7].any()                                                       # absent: zero AFTER normalising
    with pytest.raises(ValueError):
        normalize_images(x, "caffe")