Spaces:
Sleeping
Sleeping
File size: 4,947 Bytes
b2a84cb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
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
|