Spaces:
Sleeping
Sleeping
File size: 1,835 Bytes
b30e7a3 f89fa0b 3c61b44 29e1c2b b30e7a3 f89fa0b b30e7a3 f89fa0b b30e7a3 3c61b44 b30e7a3 45eb65b b30e7a3 45eb65b b30e7a3 bfc6bae 45eb65b b30e7a3 | 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 | import os
from functools import lru_cache
from typing import Callable, Dict, Optional
from models.detectors.base import ObjectDetector
from models.detectors.detr import DetrDetector
from models.detectors.grounding_dino import GroundingDinoDetector
from models.detectors.yolov11 import Yolo11Detector
from models.detectors.yolov8_visdrone import YoloV8VisDroneDetector
DEFAULT_DETECTOR = "yolo11"
_REGISTRY: Dict[str, Callable[[], ObjectDetector]] = {
"yolo11": Yolo11Detector,
"detr_resnet50": DetrDetector,
"grounding_dino": GroundingDinoDetector,
"yolov8_visdrone": YoloV8VisDroneDetector,
}
def _create_detector(name: str, **kwargs) -> ObjectDetector:
try:
factory = _REGISTRY[name]
except KeyError as exc:
available = ", ".join(sorted(_REGISTRY))
raise ValueError(f"Unknown detector '{name}'. Available: {available}") from exc
return factory(**kwargs)
@lru_cache(maxsize=None)
def _get_cached_detector(name: str) -> ObjectDetector:
return _create_detector(name)
def load_detector(name: Optional[str] = None) -> ObjectDetector:
"""Return a cached detector instance selected via arg or OBJECT_DETECTOR env."""
detector_name = name or os.getenv("OBJECT_DETECTOR", DEFAULT_DETECTOR)
return _get_cached_detector(detector_name)
def prefetch_weights(name: str) -> None:
"""Pre-download model weights (call before parallel GPU init)."""
factory = _REGISTRY.get(name)
if factory and hasattr(factory, "ensure_weights"):
factory.ensure_weights()
def load_detector_on_device(name: str, device: str) -> ObjectDetector:
"""Create a new detector instance on the specified device (no caching)."""
return _create_detector(name, device=device)
# Backwards compatibility for existing callers.
def load_model():
return load_detector()
|