Initial implementation of D-FINE model via ONNX (#16772)

* initial implementation of D-FINE model

* revert docker-compose

* add docs for D-FINE

* remove weird auto-format issue

frigate/detectors/detector_config.py

@@ -37,6 +37,7 @@ class ModelTypeEnum(str, Enum):
     yolox = "yolox"
     yolov9 = "yolov9"
     yolonas = "yolonas"
+    dfine = "dfine"
 
 
 class ModelConfig(BaseModel):

frigate/detectors/plugins/onnx.py

@@ -9,7 +9,11 @@ from frigate.detectors.detector_config import (
     BaseDetectorConfig,
     ModelTypeEnum,
 )
-from frigate.util.model import get_ort_providers, post_process_yolov9
+from frigate.util.model import (
+    get_ort_providers,
+    post_process_dfine,
+    post_process_yolov9,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -41,6 +45,7 @@ class ONNXDetector(DetectionApi):
         providers, options = get_ort_providers(
             detector_config.device == "CPU", detector_config.device
         )
+
         self.model = ort.InferenceSession(
             path, providers=providers, provider_options=options
         )
@@ -55,6 +60,16 @@ class ONNXDetector(DetectionApi):
         logger.info(f"ONNX: {path} loaded")
 
     def detect_raw(self, tensor_input: np.ndarray):
+        if self.onnx_model_type == ModelTypeEnum.dfine:
+            tensor_output = self.model.run(
+                None,
+                {
+                    "images": tensor_input,
+                    "orig_target_sizes": np.array([[self.h, self.w]], dtype=np.int64),
+                },
+            )
+            return post_process_dfine(tensor_output, self.w, self.h)
+
         model_input_name = self.model.get_inputs()[0].name
         tensor_output = self.model.run(None, {model_input_name: tensor_input})
 

frigate/util/model.py

@@ -9,7 +9,34 @@ import onnxruntime as ort
 
 logger = logging.getLogger(__name__)
 
+
 ### Post Processing
+def post_process_dfine(tensor_output: np.ndarray, width, height) -> np.ndarray:
+    class_ids = tensor_output[0][tensor_output[2] > 0.4]
+    boxes = tensor_output[1][tensor_output[2] > 0.4]
+    scores = tensor_output[2][tensor_output[2] > 0.4]
+
+    input_shape = np.array([height, width, height, width])
+    boxes = np.divide(boxes, input_shape, dtype=np.float32)
+    indices = cv2.dnn.NMSBoxes(boxes, scores, score_threshold=0.4, nms_threshold=0.4)
+    detections = np.zeros((20, 6), np.float32)
+
+    for i, (bbox, confidence, class_id) in enumerate(
+        zip(boxes[indices], scores[indices], class_ids[indices])
+    ):
+        if i == 20:
+            break
+
+        detections[i] = [
+            class_id,
+            confidence,
+            bbox[1],
+            bbox[0],
+            bbox[3],
+            bbox[2],
+        ]
+
+    return detections
 
 
 def post_process_yolov9(predictions: np.ndarray, width, height) -> np.ndarray: