tests/test_superpoint_consistency.py

import sys
from pathlib import Path

import cv2
import numpy as np
import pytest

# Ensure project root is available on sys.path when tests run directly.
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

try:
    import onnxruntime as ort  # type: ignore
except ImportError:  # pragma: no cover - dependency managed by test skip
    ort = None  # type: ignore

from backend.py.app.inference.dl_adapters.superpoint import (
    SuperPointAdapter,
    SuperPointTransformersAdapter,
)

try:
    import torch
except ImportError:  # pragma: no cover - dependency managed by test skips
    torch = None  # type: ignore


def _synthetic_corner_image(size: int = 256) -> np.ndarray:
    img = np.zeros((size, size, 3), dtype=np.uint8)
    cv2.rectangle(img, (size // 8, size // 8), (7 * size // 8, 7 * size // 8), (255, 255, 255), thickness=3)
    cv2.line(img, (size // 8, size // 8), (7 * size // 8, 7 * size // 8), (255, 255, 255), thickness=2)
    cv2.line(img, (size // 8, 7 * size // 8), (7 * size // 8, size // 8), (255, 255, 255), thickness=2)
    cv2.circle(img, (size // 2, size // 2), size // 4, (255, 255, 255), thickness=2)
    return img


def _normalized_heatmap(heat: np.ndarray) -> np.ndarray:
    heat_min = float(np.min(heat))
    heat_max = float(np.max(heat))
    eps = 1e-8
    return (heat - heat_min) / (heat_max - heat_min + eps)


@pytest.mark.skipif(ort is None, reason="onnxruntime is required for SuperPoint ONNX comparison")
@pytest.mark.xfail(
    reason="Current superpoint.onnx export diverges from the transformers reference implementation",
    strict=True,
)
def test_superpoint_onnx_matches_transformers_heatmap():
    model_path = ROOT / "models" / "superpoint.onnx"
    if not model_path.is_file():
        pytest.skip("superpoint.onnx model not available in ./models directory")

    try:
        hf_adapter = SuperPointTransformersAdapter(device="cpu")
    except ImportError as exc:  # pragma: no cover - dependency checked by skip
        pytest.skip(str(exc))
    if torch is None:  # pragma: no cover - dependency checked by skip
        pytest.skip("PyTorch is required for the transformers comparison test")

    sess = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"])
    onnx_adapter = SuperPointAdapter()

    image = _synthetic_corner_image()

    feed_onnx, ctx_onnx = onnx_adapter.preprocess(image, sess)
    outputs_onnx = sess.run(None, feed_onnx)
    semi_onnx, _ = onnx_adapter._pick_outputs(outputs_onnx)
    heat_onnx = onnx_adapter._semi_to_heat(semi_onnx)
    heat_onnx = cv2.resize(heat_onnx, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
    heat_onnx = _normalized_heatmap(heat_onnx)

    feed_hf, ctx_hf = hf_adapter.preprocess(image, None)
    outputs_hf = hf_adapter._forward(feed_hf[hf_adapter._PIXEL_VALUES_KEY])
    mask = outputs_hf.mask[0] if outputs_hf.mask is not None else torch.ones_like(outputs_hf.scores[0], dtype=torch.bool)
    mask = mask.bool()
    keypoints = outputs_hf.keypoints[0][mask]
    scores = outputs_hf.scores[0][mask]

    heat_hf = np.zeros_like(heat_onnx)
    keypoints_np = keypoints.detach().cpu().numpy()
    scores_np = scores.detach().cpu().numpy()
    H, W = image.shape[:2]
    for (x_rel, y_rel), score in zip(keypoints_np, scores_np):
        x = int(round(float(np.clip(x_rel * (W - 1), 0, W - 1))))
        y = int(round(float(np.clip(y_rel * (H - 1), 0, H - 1))))
        heat_hf[y, x] = max(heat_hf[y, x], float(score))
    heat_hf = _normalized_heatmap(heat_hf)

    correlation = np.corrcoef(heat_onnx.flatten(), heat_hf.flatten())[0, 1]
    mean_absolute_error = float(np.mean(np.abs(heat_onnx - heat_hf)))

    _, meta_onnx = onnx_adapter.postprocess(outputs_onnx, image, ctx_onnx, "Corners (SuperPoint)")
    _, meta_hf = hf_adapter.postprocess(outputs_hf, image, ctx_hf, "Corners (SuperPoint)")

    assert correlation > 0.9
    assert mean_absolute_error < 0.05
    assert meta_onnx["num_corners"] == pytest.approx(meta_hf["num_keypoints"], rel=0.1, abs=10)
    assert meta_onnx["heat_mean"] == pytest.approx(meta_hf["scores_mean"], rel=0.1, abs=1e-3)


@pytest.mark.skipif(torch is None, reason="PyTorch is required for the transformers adapter test")
def test_superpoint_transformers_adapter_infer_returns_overlay_and_meta():
    try:
        adapter = SuperPointTransformersAdapter(device="cpu")
    except ImportError as exc:  # pragma: no cover - dependency checked by skip
        pytest.skip(str(exc))

    image = _synthetic_corner_image()
    overlay, meta = adapter.infer(image, detector="Corners (SuperPoint)")

    assert overlay.shape == image.shape
    assert overlay.dtype == np.uint8
    assert meta["adapter"] == "superpoint_transformers"
    assert meta["backend"] == "transformers"
    assert isinstance(meta["num_keypoints"], int)
    assert meta["descriptors_shape"] is None or meta["descriptors_shape"][1] == 256