Add support for yolox models to onnx detector (#17773)

frigate/util/model.py

@@ -187,7 +187,12 @@ def __post_process_multipart_yolo(
 
 
 def __post_process_nms_yolo(predictions: np.ndarray, width, height) -> np.ndarray:
-    predictions = np.squeeze(predictions).T
+    predictions = np.squeeze(predictions)
+
+    # transpose the output so it has order (inferences, class_ids)
+    if predictions.shape[0] < predictions.shape[1]:
+        predictions = predictions.T
+
     scores = np.max(predictions[:, 4:], axis=1)
     predictions = predictions[scores > 0.4, :]
     scores = scores[scores > 0.4]
@@ -225,6 +230,53 @@ def post_process_yolo(output: list[np.ndarray], width: int, height: int) -> np.n
         return __post_process_nms_yolo(output[0], width, height)
 
 
+def post_process_yolox(
+    predictions: np.ndarray,
+    width: int,
+    height: int,
+    grids: np.ndarray,
+    expanded_strides: np.ndarray,
+) -> np.ndarray:
+    predictions[..., :2] = (predictions[..., :2] + grids) * expanded_strides
+    predictions[..., 2:4] = np.exp(predictions[..., 2:4]) * expanded_strides
+
+    # process organized predictions
+    predictions = predictions[0]
+    boxes = predictions[:, :4]
+    scores = predictions[:, 4:5] * predictions[:, 5:]
+
+    boxes_xyxy = np.ones_like(boxes)
+    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2
+    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
+    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2
+    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2
+
+    cls_inds = scores.argmax(1)
+    scores = scores[np.arange(len(cls_inds)), cls_inds]
+
+    indices = cv2.dnn.NMSBoxes(
+        boxes_xyxy, 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_xyxy[indices], scores[indices], cls_inds[indices])
+    ):
+        if i == 20:
+            break
+
+        detections[i] = [
+            class_id,
+            confidence,
+            bbox[1] / height,
+            bbox[0] / width,
+            bbox[3] / height,
+            bbox[2] / width,
+        ]
+
+    return detections
+
+
 ### ONNX Utilities