parking_solution/apps/vision/tests/test_vehicle.py

"""Vehicle stage (phase A) — pure post-processing on synthetic tensors, and the
recognizer composition over the stub with a fake detector. No weights needed."""

from __future__ import annotations

import os

import numpy as np
from fastapi.testclient import TestClient

from vision_service.schemas import BBox, VehicleResult
from vision_service.vehicle import (
    COCO_VEHICLE_CLASSES,
    Detection,
    decode,
    letterbox,
    nms,
    pick_vehicle,
    vehicles_from_output,
)

SIZE = 64  # tiny "model" input: grids 8x8 + 4x4 + 2x2 = 84 rows
ROWS = (SIZE // 8) ** 2 + (SIZE // 16) ** 2 + (SIZE // 32) ** 2


def raw_output(hits: list[tuple[int, int, int, float, float, float, float]]) -> np.ndarray:
    """Build a YOLOX-style raw tensor [ROWS, 85] with the given (row, coco_class, _, obj,
    cls_score, log_w, log_h) hits; everything else is background."""
    raw = np.zeros((ROWS, 85), dtype=np.float32)
    raw[:, 2:4] = -10.0  # exp → ~0 size for background rows
    for row, cls, _, obj, score, lw, lh in hits:
        raw[row, 0:2] = 0.5  # centre of its grid cell
        raw[row, 2] = lw
        raw[row, 3] = lh
        raw[row, 4] = obj
        raw[row, 5 + cls] = score
    return raw


def test_decode_maps_grid_offsets_and_log_sizes_to_pixels() -> None:
    raw = raw_output([(0, 2, 0, 1.0, 1.0, np.log(2.0), np.log(3.0))])
    dec = decode(raw, SIZE)
    # Row 0 = stride-8 grid cell (0,0): centre (0.5+0)*8 = 4, size exp(log 2)*8 = 16 / 24.
    assert dec[0, :4].tolist() == [4.0, 4.0, 16.0, 24.0]
    # Last row = stride-32 cell (1,1): centre (0.5+1)*32 = 48.
    raw2 = raw_output([(ROWS - 1, 7, 0, 1.0, 1.0, 0.0, 0.0)])
    dec2 = decode(raw2, SIZE)
    assert dec2[ROWS - 1, :4].tolist() == [48.0, 48.0, 32.0, 32.0]


def test_vehicles_only_above_floor_mapped_to_vocabulary_and_scaled_back() -> None:
    raw = raw_output(
        [
            (0, 2, 0, 0.9, 0.9, np.log(2.0), np.log(2.0)),  # car, score .81
            (1, 0, 0, 0.99, 0.99, np.log(2.0), np.log(2.0)),  # person → ignored
            (2, 7, 0, 0.5, 0.5, np.log(2.0), np.log(2.0)),  # truck, score .25 → below floor
        ]
    )
    found = vehicles_from_output(raw, SIZE, scale=0.5, min_confidence=0.4)
    assert [d.body_type for d in found] == ["car"]
    assert round(found[0].confidence, 2) == 0.81
    # Box 16px wide in the letterboxed input → 32px in the original (scale 0.5).
    assert round(found[0].x2 - found[0].x1) == 32
    assert set(COCO_VEHICLE_CLASSES.values()) == {"car", "motorcycle", "bus", "truck"}


def test_nms_keeps_the_best_of_overlapping_boxes() -> None:
    boxes = np.array([[0, 0, 10, 10], [1, 1, 11, 11], [50, 50, 60, 60]], dtype=np.float32)
    scores = np.array([0.5, 0.9, 0.7], dtype=np.float32)
    assert sorted(nms(boxes, scores, 0.45)) == [1, 2]


def test_pick_prefers_the_box_holding_the_plate_else_the_largest() -> None:
    near = Detection("car", 0.9, 0, 0, 100, 100)
    far = Detection("truck", 0.8, 200, 200, 400, 400)  # larger
    inside = Detection("car", 0.7, 10, 10, 60, 60)  # tighter box also holding the plate
    assert pick_vehicle([near, far], None) is far
    assert pick_vehicle([near, far], BBox(x1=20, y1=20, x2=30, y2=30)) is near
    assert pick_vehicle([near, far, inside], BBox(x1=20, y1=20, x2=30, y2=30)) is inside
    assert pick_vehicle([near, far], BBox(x1=900, y1=900, x2=910, y2=910)) is far  # plate outside every box
    assert pick_vehicle([], None) is None


def test_letterbox_keeps_aspect_and_pads_with_114() -> None:
    frame = np.zeros((30, 60, 3), dtype=np.uint8)
    tensor, scale = letterbox(frame, 64)
    assert tensor.shape == (1, 3, 64, 64) and tensor.dtype == np.float32
    assert abs(scale - 64 / 60) < 1e-9
    assert tensor[0, 0, 63, 63] == 114.0  # padding
    assert tensor[0, 0, 0, 0] == 0.0  # image


class FakeDetector:
    model_version = "fake-vehicle"

    def __init__(self, result: VehicleResult | None) -> None:
        self.result = result
        self.calls: list[BBox | None] = []

    def detect(self, image_bytes: bytes, plate: BBox | None) -> VehicleResult | None:
        self.calls.append(plate)
        return self.result


def test_composition_fills_vehicle_over_the_stub_and_survives_a_failing_stage() -> None:
    from vision_service.recognizer import StubRecognizer, WithVehicle
    from vision_service.settings import Settings

    det = FakeDetector(VehicleResult(body_type="truck", confidence=0.77))
    rec = WithVehicle(StubRecognizer(Settings()), det)
    res = rec.analyze(b"jpeg-bytes")
    assert res.plate is None
    assert res.vehicle == VehicleResult(body_type="truck", confidence=0.77)
    assert res.model_version == "stub-0+fake-vehicle"
    assert det.calls == [None]

    class Boom:
        model_version = "boom"

        def detect(self, image_bytes: bytes, plate: BBox | None) -> VehicleResult | None:
            raise RuntimeError("no model")

    rec2 = WithVehicle(StubRecognizer(Settings()), Boom())
    res2 = rec2.analyze(b"jpeg-bytes")
    assert res2.vehicle is None
    assert rec2.ready is True
    assert "vehicle: RuntimeError: no model" in (rec2.error or "")


def test_app_reports_a_missing_model_file_and_keeps_serving() -> None:
    from vision_service.app import app

    os.environ["VISION_VEHICLE_MODEL_PATH"] = "/nonexistent/yolox.onnx"
    try:
        with TestClient(app) as client:
            health = client.get("/health").json()
            assert health["ready"] is True  # the plate stage (stub) is fine
            assert "vehicle:" in (health["detail"] or "")
            res = client.post("/analyze", content=b"x", headers={"content-type": "application/octet-stream"})
            assert res.status_code == 200
            assert res.json()["vehicle"] is None
    finally:
        os.environ.pop("VISION_VEHICLE_MODEL_PATH", None)