Spaces:
Running on Zero
Running on Zero
File size: 3,147 Bytes
087643a | 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 | # Thin wrapper around Ultralytics so nothing else in the app imports it
# directly - lets me stub this out in tests without a GPU or real weights.
from __future__ import annotations
import os
import time
from pathlib import Path
from app.constants import DEFAULT_QUERY_BUDGET, ID_TO_CLASS
from app.schemas import BBox, Detection
class Detector:
"""Loads RT-DETR on first use, converts raw output to Detection objects."""
def __init__(self, weights_path: str | None = None):
self._weights_path = weights_path or os.environ.get("MODEL_PATH", "./weights/best.pt")
self._model = None
self.query_budget = DEFAULT_QUERY_BUDGET
@property
def is_loaded(self) -> bool:
return self._model is not None
def load(self) -> None:
"""Lazy load - keeps a bad MODEL_PATH from crashing the whole app
at startup. Checks the file exists first so the error is readable
instead of a torch.load stack trace.
If the weights aren't on disk but MODEL_URL is set, downloads them
first - this is what lets a fresh HF Space or Docker container just
work without manually copying a 66MB file into every host."""
if self.is_loaded:
return
path = Path(self._weights_path)
if not path.exists():
model_url = os.environ.get("MODEL_URL")
if model_url:
path.parent.mkdir(parents=True, exist_ok=True)
import urllib.request
urllib.request.urlretrieve(model_url, path)
else:
raise FileNotFoundError(
f"Model weights not found at '{self._weights_path}'. "
"Set MODEL_PATH to an existing file, or set MODEL_URL to "
"download from - see the README for the weights link."
)
from ultralytics import RTDETR
self._model = RTDETR(self._weights_path)
def predict(self, image, conf: float = 0.25) -> tuple[list[Detection], float]:
"""Runs inference, returns (detections, inference_ms). Timing is
measured here, not in the route handler, so it's not conflated with
request/network overhead.
conf=0.25 is intentionally permissive - the guardrail decides what
counts as confident enough, not this default."""
if not self.is_loaded:
self.load()
started = time.perf_counter()
results = self._model.predict(image, conf=conf, verbose=False)[0]
elapsed_ms = (time.perf_counter() - started) * 1000
detections = [
Detection(
class_name=ID_TO_CLASS.get(int(class_id), f"unknown_{int(class_id)}"),
class_id=int(class_id),
confidence=float(confidence),
bbox=BBox(x1=float(x1), y1=float(y1), x2=float(x2), y2=float(y2)),
)
for class_id, confidence, (x1, y1, x2, y2) in zip(
results.boxes.cls.tolist(),
results.boxes.conf.tolist(),
results.boxes.xyxy.tolist(),
)
]
return detections, elapsed_ms
|