From 622a45ef12a455a0f1fc4b0a5719a5c4c368b193 Mon Sep 17 00:00:00 2001
From: Indrek Mandre <indrek@mare.ee>
Date: Sat, 27 Jan 2024 22:35:32 +0200
Subject: [PATCH] detectors/rocm: separated preprocessing functions into
 yolo_utils.py

---
 frigate/detectors/plugins/rocm.py | 27 +++++--------------------
 frigate/detectors/yolo_utils.py   | 33 +++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 22 deletions(-)
 create mode 100644 frigate/detectors/yolo_utils.py

diff --git a/frigate/detectors/plugins/rocm.py b/frigate/detectors/plugins/rocm.py
index ad7d1f6b0..5a306cefd 100644
--- a/frigate/detectors/plugins/rocm.py
+++ b/frigate/detectors/plugins/rocm.py
@@ -12,29 +12,12 @@ import cv2
 from frigate.detectors.detection_api import DetectionApi
 from frigate.detectors.detector_config import BaseDetectorConfig
 
+import frigate.detectors.yolo_utils as yolo_utils
+
 logger = logging.getLogger(__name__)
 
 DETECTOR_KEY = "rocm"
 
-# XXX several detectors run yolov8, this should probably be common code in some utils module
-def postprocess_yolov8(model_input_shape, tensor_output, box_count = 20):
-    model_box_count = tensor_output.shape[2]
-    model_class_count = tensor_output.shape[1] - 4
-    probs = tensor_output[0, 4:, :]
-    all_ids = np.argmax(probs, axis=0)
-    all_confidences = np.take(probs.T, model_class_count*np.arange(0, model_box_count) + all_ids)
-    all_boxes = tensor_output[0, 0:4, :].T
-    mask = (all_confidences > 0.30)
-    class_ids = all_ids[mask]
-    confidences = all_confidences[mask]
-    cx, cy, w, h = all_boxes[mask].T
-    scale_y, scale_x = 1 / model_input_shape[2], 1 / model_input_shape[3]
-    detections = np.stack((class_ids, confidences, scale_y * (cy - h / 2), scale_x * (cx - w / 2), scale_y * (cy + h / 2), scale_x * (cx + w / 2)), axis=1)
-    if detections.shape[0] > box_count:
-        detections = detections[np.argpartition(detections[:,1], -box_count)[-box_count:]]
-    detections.resize((box_count, 6))
-    return detections
-
 class ROCmDetectorConfig(BaseDetectorConfig):
     type: Literal[DETECTOR_KEY]
 
@@ -47,7 +30,7 @@ class ROCmDetector(DetectionApi):
             import migraphx
 
             logger.info(f"AMD/ROCm: loaded migraphx module")
-        except ValueError:
+        except ModuleNotFoundError:
             logger.error(
                 "AMD/ROCm: module loading failed, missing ROCm environment?"
             )
@@ -88,12 +71,12 @@ class ROCmDetector(DetectionApi):
         model_input_name = self.model.get_parameter_names()[0];
         model_input_shape = tuple(self.model.get_parameter_shapes()[model_input_name].lens());
 
-        tensor_input = cv2.dnn.blobFromImage(tensor_input[0], 1.0 / 255, (model_input_shape[3], model_input_shape[2]), None, swapRB=False)
+        tensor_input = yolo_utils.yolov8_preprocess(tensor_input, model_input_shape)
 
         detector_result = self.model.run({model_input_name: tensor_input})[0]
 
         addr = ctypes.cast(detector_result.data_ptr(), ctypes.POINTER(ctypes.c_float))
         tensor_output = np.ctypeslib.as_array(addr, shape=detector_result.get_shape().lens())
 
-        return postprocess_yolov8(model_input_shape, tensor_output)
+        return yolo_utils.yolov8_postprocess(model_input_shape, tensor_output)
 
diff --git a/frigate/detectors/yolo_utils.py b/frigate/detectors/yolo_utils.py
new file mode 100644
index 000000000..38083e705
--- /dev/null
+++ b/frigate/detectors/yolo_utils.py
@@ -0,0 +1,33 @@
+import logging
+
+import numpy as np
+import cv2
+
+logger = logging.getLogger(__name__)
+
+def yolov8_preprocess(tensor_input, model_input_shape):
+  # tensor_input must be nhwc
+  assert tensor_input.shape[3] == 3
+  if tuple(tensor_input.shape[1:3]) != tuple(model_input_shape[2:4]):
+    logger.warn(f"yolov8_preprocess: tensor_input.shape {tensor_input.shape} and model_input_shape {model_input_shape} do not match!")
+  # cv2.dnn.blobFromImage is faster than numpying it
+  return cv2.dnn.blobFromImage(tensor_input[0], 1.0 / 255, (model_input_shape[3], model_input_shape[2]), None, swapRB=False)
+
+def yolov8_postprocess(model_input_shape, tensor_output, box_count = 20):
+    model_box_count = tensor_output.shape[2]
+    model_class_count = tensor_output.shape[1] - 4
+    probs = tensor_output[0, 4:, :]
+    all_ids = np.argmax(probs, axis=0)
+    all_confidences = np.take(probs.T, model_class_count*np.arange(0, model_box_count) + all_ids)
+    all_boxes = tensor_output[0, 0:4, :].T
+    mask = (all_confidences > 0.30)
+    class_ids = all_ids[mask]
+    confidences = all_confidences[mask]
+    cx, cy, w, h = all_boxes[mask].T
+    scale_y, scale_x = 1 / model_input_shape[2], 1 / model_input_shape[3]
+    detections = np.stack((class_ids, confidences, scale_y * (cy - h / 2), scale_x * (cx - w / 2), scale_y * (cy + h / 2), scale_x * (cx + w / 2)), axis=1)
+    if detections.shape[0] > box_count:
+        detections = detections[np.argpartition(detections[:,1], -box_count)[-box_count:]]
+    detections.resize((box_count, 6))
+    return detections
+