code/tt_ptv3/tests/_data.py

# SPDX-License-Identifier: Apache-2.0
"""Where the tests find data that does not ship in ``code/``: the AutowareFoundation/ptv3 weights, the research
scripts and goldens of the workspace, the git-ignored PandaSet sample staging copy, and the 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`` / ``small_golden``): inside the workspace a missing sample
FAILS (``missing``), so that the frozen gates never silently stop running on it.

- weights: ``$PTV3_WEIGHTS_DIR``, else the workspace download ``assets/ptv3/hf_ptv3`` (sha256-verified, SPEC 1.4),
  else the local HF cache snapshot at the pinned revision (``local_files_only``: tests never download);
- research: ``$PTV3_RESEARCH_DIR``, else ``research/ptv3`` (scripts, public_data inputs + goldens);
- T4 sample frames: ``$PTV3_T4_DIR``, else ``assets/ptv3/t4_lidar_detection_sample/lidar_detection_sample_dataset``;
- large goldens (per-module taps): ``$PTV3_GOLDEN_DIR``, else ``research/ptv3/goldens``
  (written by ``code/scripts/make_goldens.py``; never shipped);
- samples (``sample_path`` / ``find_sample``): the shipped ``samples/`` (``test_pcd.npz``, Autoware's ``test.pcd``,
  Apache-2.0), then the git-ignored ``staging_samples_pandaset/`` of the bundle (``pandaset_019_f40.npz`` and its
  stored references: PandaSet-derived, not shipped until the user approves PandaSet samples, research/DATASETS.md
  N5; ``code/scripts/make_samples.py`` writes it there);
- small goldens (``small_golden``: counts, flags, host-table digests, boxes; ``make_goldens.py``): ``tests/goldens``
  (shipped: ``test_pcd`` only), then ``staging_samples_pandaset/goldens`` (``ps019_f40``), then
  ``research/ptv3/goldens/summaries`` (T4 and nuScenes frames: never shipped).
"""
from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Dict, 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 "sample" frame of the frozen e2e / PCC gates. staging_samples_pandaset/README.md: swap-in steps.
PANDASET_STAGING = PKG_DIR.parents[1] / "staging_samples_pandaset"
SAMPLE_PANDASET_NAME = "pandaset_019_f40.npz"
SHIPPED_SAMPLE = SAMPLES / "test_pcd.npz"          # the shipped sample: autoware_universe test.pcd (Apache-2.0)
SMALL_GOLDEN_DIRS = (GOLDENS, PANDASET_STAGING / "goldens", WORKSPACE / "research" / "ptv3" / "goldens" / "summaries")
# T4 sample-dataset frames (SPEC 7; licence UNVERIFIED: never shipped) and the public golden frames
# (research/ptv3/public_data: PandaSet CC BY 4.0, nuScenes CC BY-NC-SA 4.0 -- never shipped)
T4_FRAMES = {"t4_scene10_00795": "scene_10/data/LIDAR_CONCAT/00795.pcd.bin",
             "t4_scene6_00860": "scene_6/data/LIDAR_CONCAT/00860.pcd.bin",
             "t4_scene7_00270": "scene_7/data/LIDAR_CONCAT/00270.pcd.bin",
             "t4_scene7_00720": "scene_7/data/LIDAR_CONCAT/00720.pcd.bin"}
PUBLIC_FRAMES = ("ps019_f20", "ps019_f40", "ps019_f60", "ps005_f20", "ps005_f40", "ps005_f60", "ps029_f20",
                 "ps029_f40", "ps029_f60", "ps090_f20", "ps090_f40", "ps090_f60", "ns0103_k00", "ns0103_k20",
                 "ns0916_k00", "ns0916_k20")
# frames with per-module tap goldens (make_goldens.py --taps): the shipped sample, a T4 frame, a sparse nuScenes one
TAP_FRAMES = ("ps019_f40", "t4_scene10_00795", "ns0103_k20")


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" / "ptv3" / "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)")
    return p


def pandaset_sample() -> Path:
    """The PandaSet 019 f40 sample (106,501 points, 73,992 voxels): ``samples/`` once it ships, until then the
    staging copy. It is the "sample" frame of the frozen device gates."""
    return sample_path(SAMPLE_PANDASET_NAME)


def small_golden(stem: str) -> dict:
    """The small golden ``<frame>.<mode>.json`` of ``make_goldens.py`` (:data:`SMALL_GOLDEN_DIRS`, first hit)."""
    import json
    for d in SMALL_GOLDEN_DIRS:
        if (d / f"{stem}.json").is_file():
            return json.loads((d / f"{stem}.json").read_text())
    missing(f"small golden {stem}.json not found in {[str(d) for d in SMALL_GOLDEN_DIRS]} (make_goldens.py)")


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

        from tt_ptv3.api import PTv3

        snap = Path(snapshot_download(PTv3.DEFAULT_REPO, revision=PTv3.DEFAULT_REVISION,
                                      allow_patterns=list(PTv3.ALLOW_PATTERNS), local_files_only=True))
        if (snap / "ptv3_encoder.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("AutowareFoundation/ptv3 weights not found (set PTV3_WEIGHTS_DIR)")
    return d


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


def research_modules():
    """Import the research reference scripts (``preprocess_ref``, ``postprocess_ref``, ``ptv3_torch_ref``,
    ``ptv3_public_common``) from the workspace."""
    r = research_dir()
    for p in (r / "scripts", r / "public_data" / "scripts"):
        if str(p) not in sys.path:
            sys.path.insert(0, str(p))
    import importlib
    return {m: importlib.import_module(m) for m in ("preprocess_ref", "postprocess_ref", "ptv3_public_common")}


def golden_dir() -> Path:
    env = os.environ.get("PTV3_GOLDEN_DIR")
    d = _first_dir(Path(env) if env else None, WORKSPACE / "research" / "ptv3" / "goldens")
    if d is None:
        pytest.skip("research/ptv3/goldens not found (run code/scripts/make_goldens.py)")
    return d


def t4_dir() -> Optional[Path]:
    env = os.environ.get("PTV3_T4_DIR")
    return _first_dir(Path(env) if env else None,
                      WORKSPACE / "assets" / "ptv3" / "t4_lidar_detection_sample" / "lidar_detection_sample_dataset")


def load_frame(name: str) -> np.ndarray:
    """Points ``(N, >= 4)`` (x, y, z, intensity 0-255) of a named frame: ``sample``, a T4 frame, a public frame."""
    if name == "sample":                            # the PandaSet sample (staged until it may ship)
        with np.load(pandaset_sample()) as z:
            return np.asarray(z["points"], np.float32)
    if name == "test_pcd":                          # the shipped sample (structured array: x, y, z, intensity)
        from tt_ptv3 import io as tio
        return np.asarray(tio.load_points(str(SHIPPED_SAMPLE)).points, np.float32)
    if name in T4_FRAMES:
        d = t4_dir()
        if d is None or not (d / T4_FRAMES[name]).is_file():
            pytest.skip(f"T4 frame {name} not found")
        return np.fromfile(d / T4_FRAMES[name], dtype=np.float32).reshape(-1, 5)
    p = research_dir() / "public_data" / "inputs" / f"{name}.pcd5.npy"
    if not p.is_file():
        pytest.skip(f"public frame {name} not found")
    return np.load(p)


def public_golden(name: str, mode: str = "exact") -> Dict[str, np.ndarray]:
    stem = name if mode == "exact" else f"{name}.{mode}"
    p = research_dir() / "public_data" / "goldens" / f"{stem}.golden.npz"
    if not p.is_file():
        pytest.skip(f"public golden {stem} not found")
    with np.load(p) as z:
        return {k: z[k] for k in z.files}


def public_summary(name: str, mode: str = "exact") -> dict:
    import json
    stem = name if mode == "exact" else f"{name}.{mode}"
    p = research_dir() / "public_data" / "goldens" / f"{stem}.summary.json"
    if not p.is_file():
        pytest.skip(f"public golden summary {stem} not found")
    return json.loads(p.read_text())


def need(module: str) -> None:
    import importlib.util
    if importlib.util.find_spec(module) is None:
        pytest.skip(f"{module} not installed (research venv: tools/research-venv)")