From 8951647f3370dab4c8cf9e55727003300bb72dda Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Fri, 18 Apr 2025 07:51:57 -0600
Subject: [PATCH] Use existing variable

---
 frigate/detectors/plugins/onnx.py | 24 +++++++------
 frigate/detectors/plugins/rknn.py | 60 +++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 10 deletions(-)

diff --git a/frigate/detectors/plugins/onnx.py b/frigate/detectors/plugins/onnx.py
index 6c6a39e783..8affc624dc 100644
--- a/frigate/detectors/plugins/onnx.py
+++ b/frigate/detectors/plugins/onnx.py
@@ -52,8 +52,6 @@ class ONNXDetector(DetectionApi):
             path, providers=providers, provider_options=options
         )
 
-        self.h = detector_config.model.height
-        self.w = detector_config.model.width
         self.onnx_model_type = detector_config.model.model_type
         self.onnx_model_px = detector_config.model.input_pixel_format
         self.onnx_model_shape = detector_config.model.input_tensor
@@ -70,10 +68,12 @@ class ONNXDetector(DetectionApi):
                 None,
                 {
                     "images": tensor_input,
-                    "orig_target_sizes": np.array([[self.h, self.w]], dtype=np.int64),
+                    "orig_target_sizes": np.array(
+                        [[self.height, self.width]], dtype=np.int64
+                    ),
                 },
             )
-            return post_process_dfine(tensor_output, self.w, self.h)
+            return post_process_dfine(tensor_output, self.width, self.height)
 
         model_input_name = self.model.get_inputs()[0].name
         tensor_output = self.model.run(None, {model_input_name: tensor_input})
@@ -95,17 +95,21 @@ class ONNXDetector(DetectionApi):
                 detections[i] = [
                     class_id,
                     confidence,
-                    y_min / self.h,
-                    x_min / self.w,
-                    y_max / self.h,
-                    x_max / self.w,
+                    y_min / self.height,
+                    x_min / self.width,
+                    y_max / self.height,
+                    x_max / self.width,
                 ]
             return detections
         elif self.onnx_model_type == ModelTypeEnum.yologeneric:
-            return post_process_yolo(tensor_output, self.w, self.h)
+            return post_process_yolo(tensor_output, self.width, self.height)
         elif self.onnx_model_type == ModelTypeEnum.yolox:
             return post_process_yolox(
-                tensor_output[0], self.w, self.h, self.grids, self.expanded_strides
+                tensor_output[0],
+                self.width,
+                self.height,
+                self.grids,
+                self.expanded_strides,
             )
         else:
             raise Exception(
diff --git a/frigate/detectors/plugins/rknn.py b/frigate/detectors/plugins/rknn.py
index 407c93917a..3805d9fba3 100644
--- a/frigate/detectors/plugins/rknn.py
+++ b/frigate/detectors/plugins/rknn.py
@@ -4,6 +4,7 @@ import re
 import urllib.request
 from typing import Literal
 
+import numpy as np
 from pydantic import Field
 
 from frigate.const import MODEL_CACHE_DIR
@@ -60,6 +61,9 @@ class Rknn(DetectionApi):
                 "Error initializing rknn runtime. Do you run docker in privileged mode?"
             )
 
+        if self.detector_config.model.model_type == ModelTypeEnum.yolox:
+            self.calculate_grids_strides()
+
     def __del__(self):
         self.rknn.release()
 
@@ -150,6 +154,62 @@ class Rknn(DetectionApi):
                 'Make sure to set the model input_tensor to "nhwc" in your config.'
             )
 
+    def post_process_yolonas(self, output: list[np.ndarray]):
+        """
+        @param output: output of inference
+        expected shape: [np.array(1, N, 4), np.array(1, N, 80)]
+        where N depends on the input size e.g. N=2100 for 320x320 images
+
+        @return: best results: np.array(20, 6) where each row is
+        in this order (class_id, score, y1/height, x1/width, y2/height, x2/width)
+        """
+
+        N = output[0].shape[1]
+
+        boxes = output[0].reshape(N, 4)
+        scores = output[1].reshape(N, 80)
+
+        class_ids = np.argmax(scores, axis=1)
+        scores = scores[np.arange(N), class_ids]
+
+        args_best = np.argwhere(scores > self.thresh)[:, 0]
+
+        num_matches = len(args_best)
+        if num_matches == 0:
+            return np.zeros((20, 6), np.float32)
+        elif num_matches > 20:
+            args_best20 = np.argpartition(scores[args_best], -20)[-20:]
+            args_best = args_best[args_best20]
+
+        boxes = boxes[args_best]
+        class_ids = class_ids[args_best]
+        scores = scores[args_best]
+
+        boxes = np.transpose(
+            np.vstack(
+                (
+                    boxes[:, 1] / self.height,
+                    boxes[:, 0] / self.width,
+                    boxes[:, 3] / self.height,
+                    boxes[:, 2] / self.width,
+                )
+            )
+        )
+
+        results = np.hstack(
+            (class_ids[..., np.newaxis], scores[..., np.newaxis], boxes)
+        )
+
+        return np.resize(results, (20, 6))
+
+    def post_process(self, output):
+        if self.detector_config.model.model_type == ModelTypeEnum.yolonas:
+            return self.post_process_yolonas(output)
+        else:
+            raise ValueError(
+                f'Model type "{self.detector_config.model.model_type}" is currently not supported.'
+            )
+
     def detect_raw(self, tensor_input):
         output = self.rknn.inference(
             [