code/tt_transfusion/tests/_data.py

# SPDX-License-Identifier: Apache-2.0
"""Where the tests find data that does not ship in code/: the AutowareFoundation/lidar_transfusion weights, the
research scripts and goldens of the workspace, the git-ignored PandaSet sample staging copy, and the bundle's own small
goldens. Missing data skips the test that needs it: the container image and a Hub clone have neither the workspace
nor, offline, the weights. The one exception is a sample (``sample_path``): inside the workspace a missing sample
FAILS (``missing``), so that the frozen sample gates never silently stop running on it.

- weights: ``$TRANSFUSION_WEIGHTS_DIR`` (holds ``transfusion.onnx`` and the yaml files), else the workspace download
  ``assets/transfusion/hf_lidar_transfusion`` (sha256-verified, S:transfusion:40-48), else the local HF cache snapshot
  at the pinned revision (``local_files_only``: tests never download);
- research: ``$TRANSFUSION_RESEARCH_DIR``, else ``research/transfusion`` of the workspace (scripts, sample frames,
  ORT goldens, public_data);
- large goldens: ``$TRANSFUSION_GOLDEN_DIR``, else ``research/transfusion/goldens`` (written by
  ``code/scripts/make_goldens.py``; never shipped);
- samples (``sample_path`` / ``find_sample``): the shipped ``samples/`` (``test_pcd.npz``, Apache-2.0), then the
  git-ignored ``staging_samples_pandaset/`` of the bundle (``pandaset_019_f40.npz`` and its stored reference:
  PandaSet-derived, not shipped until the user approves PandaSet samples, research/DATASETS.md N5;
  ``code/scripts/make_samples.py`` and ``make_goldens.py`` write it there).
"""
from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Optional

import numpy as np
import pytest

PKG_DIR = Path(__file__).resolve().parents[1]
SAMPLES = PKG_DIR / "samples"
GOLDENS = Path(__file__).resolve().parent / "goldens"
WORKSPACE = Path("/home/ubuntu/experiments/tt-models")
# The PandaSet sample waits in the bundle's git-ignored staging directory until it may ship (DATASETS.md N5); it stays
# the frame of the frozen ``sample.*`` gates of test_e2e_device.py. staging_samples_pandaset/README.md: swap-in steps.
PANDASET_STAGING = PKG_DIR.parents[1] / "staging_samples_pandaset"
SAMPLE_PANDASET_NAME = "pandaset_019_f40.npz"
SAMPLE_TESTPCD = SAMPLES / "test_pcd.npz"      # the shipped sample: autoware_universe test.pcd (Apache-2.0)
FIELDS5 = ("x", "y", "z", "intensity", "time_lag")      # the research (M, 5) arrays
# The development frames of test_e2e_device.py (and of every precision choice) are the 16 public frames of
# research/transfusion/public_data, test.pcd and the sample-rosbag frames. The HELD-OUT frames (fix round 1,
# VERIFY_PORT.md M2) are 34 other 2-sweep frames of the same public sequences, converted by the research converter
# (research/transfusion/public_data/scripts/tf_public_convert.py) into research/transfusion/heldout/inputs, never
# used to choose anything: PandaSet 019 / 005 / 029 / 090 x 4, nuScenes scene-0103 key-frames 5 / 15 / 25 / 35 and
# scene-0916 5 / 15 / 20 / 30 / 35, Argoverse 2 (3 logs) sweeps 30 / 120 / 150.
HELDOUT = (tuple(f"ps{s}_f{f:02d}" for s, fs in (("019", (5, 25, 55, 75)), ("005", (15, 25, 55, 65)),
                                                  ("029", (15, 25, 55, 65)), ("090", (5, 25, 55, 75))) for f in fs)
           + tuple(f"ns0103_k{k:02d}" for k in (5, 15, 25, 35)) + tuple(f"ns0916_k{k:02d}" for k in (5, 15, 20, 30, 35))
           + tuple(f"av2_{log}_s{s:03d}" for log in ("280269f9", "51bbdd4d", "e13c06cb") for s in (30, 120, 150)))


def _first_dir(*candidates: Optional[Path]) -> Optional[Path]:
    for c in candidates:
        if c is not None and c.is_dir():
            return c
    return None


def in_workspace() -> bool:
    """True inside the porting workspace (the research tree exists): there, missing sample data is an error."""
    return (WORKSPACE / "research" / "transfusion" / "SPEC.md").is_file()


def missing(message: str):
    """Missing sample data: FAIL inside the workspace (a gate must never run on less data than it was frozen on), skip
    elsewhere (a Hub clone or the container image has only the shipped sample)."""
    if in_workspace():
        pytest.fail(f"{message} [inside the workspace missing sample data is an error, not a skip]")
    pytest.skip(message)


def find_sample(name: str) -> Optional[Path]:
    """A sample file by name (input or stored reference): the shipped ``samples/``, then the PandaSet staging copy."""
    for d in (SAMPLES, PANDASET_STAGING):
        if (d / name).is_file():
            return d / name
    return None


def sample_path(name: str) -> Path:
    """:func:`find_sample`, or :func:`missing` (fail in the workspace, skip elsewhere)."""
    p = find_sample(name)
    if p is None:
        missing(f"sample {name} found neither in {SAMPLES} nor in {PANDASET_STAGING} (code/scripts/make_samples.py; "
                "the stored references: code/scripts/make_goldens.py)")
    return p


def pandaset_sample() -> Path:
    """The PandaSet 019 f40 sample (213,007 points, 2 sweeps, 38,158 pillars): ``samples/`` once it ships, until
    then the staging copy."""
    return sample_path(SAMPLE_PANDASET_NAME)


def find_weights_dir() -> Optional[Path]:
    env = os.environ.get("TRANSFUSION_WEIGHTS_DIR")
    found = _first_dir(Path(env) if env else None, WORKSPACE / "assets" / "transfusion" / "hf_lidar_transfusion")
    if found is not None and (found / "transfusion.onnx").is_file():
        return found
    try:
        from huggingface_hub import snapshot_download

        from tt_transfusion.api import TransFusion

        snap = Path(snapshot_download(TransFusion.DEFAULT_REPO, revision=TransFusion.DEFAULT_REVISION,
                                      allow_patterns=list(TransFusion.ALLOW_PATTERNS), local_files_only=True))
        if (snap / "transfusion.onnx").is_file():
            return snap
    except Exception:  # noqa: BLE001 -- not cached / no huggingface_hub: the caller skips
        pass
    return None


def weights_dir() -> Path:
    d = find_weights_dir()
    if d is None:
        pytest.skip("lidar_transfusion weights not found (set TRANSFUSION_WEIGHTS_DIR)")
    return d


def research_dir() -> Path:
    env = os.environ.get("TRANSFUSION_RESEARCH_DIR")
    d = _first_dir(Path(env) if env else None, WORKSPACE / "research" / "transfusion")
    if d is None or not (d / "scripts" / "tf_ref_common.py").is_file():
        pytest.skip("research/transfusion (the research reference scripts) not found")
    return d


def research_module(name: str):
    """Import a research script by name (``tf_ref_common``, ``verify_ort_independent``, ...); their directory is
    workspace code, not a download."""
    scripts = str(research_dir() / "scripts")
    if scripts not in sys.path:
        sys.path.insert(0, scripts)
    import importlib

    return importlib.import_module(name)


def research_cloud(name: str):
    """A research input as a ``PointCloud`` with fields x, y, z, intensity, time_lag: ``frame_0150`` (the Autoware
    sample-rosbag frames, license unstated: research only) or a public-data id (``ps019_f40``, ``ns0103_k20``, ...)."""
    from tt_transfusion import io as tio

    r = research_dir()
    path = (r / "sample_frames" / f"{name}_sweeps2.npy" if name.startswith("frame_")
            else r / "public_data" / "inputs" / f"{name}.npy")
    if not path.is_file():
        pytest.skip(f"research input {path} not found")
    return tio.PointCloud(np.load(path, allow_pickle=False), FIELDS5)


def golden_dir() -> Path:
    env = os.environ.get("TRANSFUSION_GOLDEN_DIR")
    d = _first_dir(Path(env) if env else None, WORKSPACE / "research" / "transfusion" / "goldens")
    if d is None:
        pytest.skip("large goldens (research/transfusion/goldens) not found")
    return d