RT-DETR compatibility

perception_roi_server.py

@@ -31,7 +31,7 @@ import numpy as np
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, StreamingResponse, JSONResponse
-from ultralytics import YOLO
+from ultralytics import YOLO, RTDETR
 
 DEFAULT_WEIGHTS = os.environ.get("YOLO_WEIGHTS", "yolov8s.pt")
 WEIGHTS_DIR = os.environ.get("WEIGHTS_DIR", os.path.dirname(__file__))
@@ -57,7 +57,13 @@ def root():
     return {"status": "ok", "service": "roi-compression"}
 
 _model_lock = threading.Lock()
-_models: Dict[str, YOLO] = {}
+_models: Dict[str, Any] = {}
+
+def _infer_model_type(weights: str) -> str:
+    name = os.path.basename(str(weights or "")).lower()
+    if name.startswith("rtdetr"):
+        return "rtdetr"
+    return "yolo"
 
 def _resolve_weights_path(weights: str) -> (str, List[str]):
     if not weights:
@@ -93,8 +99,9 @@ def _resolve_weights_path(weights: str) -> (str, List[str]):
             return os.path.abspath(cand), checked
     return w, checked
 
-def get_model(weights: str) -> YOLO:
+def get_model(weights: str) -> Any:
     key, checked = _resolve_weights_path(weights or DEFAULT_WEIGHTS)
+    model_type = _infer_model_type(key)
     if str(key).endswith(".pt") and not os.path.exists(key):
         search_list = ", ".join(checked) if checked else "(no local paths searched)"
         raise RuntimeError(
@@ -102,9 +109,13 @@ def get_model(weights: str) -> YOLO:
             f"Set WEIGHTS_DIR or upload the weights to the app directory."
         )
     with _model_lock:
-        if key not in _models:
-            _models[key] = YOLO(key)
-        return _models[key]
+        cache_key = f"{model_type}:{key}"
+        if cache_key not in _models:
+            if model_type == "rtdetr":
+                _models[cache_key] = RTDETR(key)
+            else:
+                _models[cache_key] = YOLO(key)
+        return _models[cache_key]
 
 
 def _parse_queries(q: str) -> List[str]:
@@ -124,7 +135,7 @@ def _keep_det(label: str, queries: List[str]) -> bool:
 
 
 def _yolo_detect_frame(
-    model: YOLO,
+    model: Any,
     frame_bgr: np.ndarray,
     conf: float,
     queries: List[str],