80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
"""Object-detection tests with a fake YOLO predictor (no torch/weights)."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
|
||
from fastapi.testclient import TestClient
|
||
from PIL import Image
|
||
|
||
from extractors.app import app
|
||
from extractors.features import detect
|
||
|
||
|
||
class _Box:
|
||
def __init__(self, cls_id, conf, xyxy):
|
||
self.cls = [cls_id]
|
||
self.conf = [conf]
|
||
self.xyxy = [xyxy]
|
||
|
||
|
||
class _Result:
|
||
names = {0: "person", 2: "car"}
|
||
|
||
def __init__(self, boxes):
|
||
self.boxes = boxes
|
||
|
||
|
||
class FakeYOLO:
|
||
def __init__(self, boxes):
|
||
self._boxes = boxes
|
||
|
||
def predict(self, img, **kwargs):
|
||
return [_Result(self._boxes)]
|
||
|
||
|
||
def _png() -> bytes:
|
||
buf = io.BytesIO()
|
||
Image.new("RGB", (64, 64), (0, 0, 0)).save(buf, format="PNG")
|
||
return buf.getvalue()
|
||
|
||
|
||
def test_summarize_counts():
|
||
dets = [
|
||
{"label": "person", "confidence": 0.9, "box": [0, 0, 10, 10]},
|
||
{"label": "person", "confidence": 0.8, "box": [5, 5, 20, 20]},
|
||
{"label": "car", "confidence": 0.7, "box": [0, 0, 30, 30]},
|
||
]
|
||
res = detect._summarize(dets)
|
||
assert res.results["count"] == 3
|
||
assert res.results["by_label"] == {"person": 2, "car": 1}
|
||
assert "person×2" in res.evidence[0]
|
||
|
||
|
||
def test_analyze_parses_and_sorts():
|
||
fake = FakeYOLO([_Box(2, 0.6, [1, 1, 9, 9]), _Box(0, 0.95, [0, 0, 8, 8])])
|
||
res = detect.analyze(_png(), model=fake)
|
||
assert res.ok is True
|
||
assert res.results["count"] == 2
|
||
# sorted by confidence desc → person (0.95) first
|
||
assert res.results["detections"][0]["label"] == "person"
|
||
assert res.results["detections"][0]["confidence"] == 0.95
|
||
|
||
|
||
def test_analyze_empty_image():
|
||
res = detect.analyze(b"", model=FakeYOLO([]))
|
||
assert res.ok is False
|
||
|
||
|
||
def test_detect_endpoint_503_without_ml_extra():
|
||
client = TestClient(app)
|
||
r = client.post("/v1/detect", files={"file": ("x.png", _png(), "image/png")})
|
||
assert r.status_code == 503
|
||
|
||
|
||
def test_detect_endpoint_with_monkeypatched_model(monkeypatch):
|
||
monkeypatch.setattr(detect, "get_model", lambda: FakeYOLO([_Box(0, 0.9, [0, 0, 5, 5])]))
|
||
client = TestClient(app)
|
||
r = client.post("/v1/detect", files={"file": ("x.png", _png(), "image/png")})
|
||
assert r.status_code == 200
|
||
assert r.json()["results"]["by_label"] == {"person": 1}
|