From e3342cd2dd7f6da494fccd6e5d2f3ece0b6f31e9 Mon Sep 17 00:00:00 2001
From: MarcA711 <>
Date: Wed, 1 Nov 2023 19:28:10 +0000
Subject: [PATCH] format rknn.py

---
 frigate/detectors/plugins/rknn.py | 49 +++++++++++++++++++++----------
 1 file changed, 34 insertions(+), 15 deletions(-)

diff --git a/frigate/detectors/plugins/rknn.py b/frigate/detectors/plugins/rknn.py
index aa64bf45b..9a2de2daf 100644
--- a/frigate/detectors/plugins/rknn.py
+++ b/frigate/detectors/plugins/rknn.py
@@ -13,10 +13,16 @@ logger = logging.getLogger(__name__)
 
 DETECTOR_KEY = "rknn"
 
+
 class RknnDetectorConfig(BaseDetectorConfig):
     type: Literal[DETECTOR_KEY]
-    score_thresh: float = Field(default=0.5, ge=0, le= 1, title="Minimal confidence for detection.")
-    nms_thresh: float = Field(default=0.45, ge=0, le= 1, title="IoU threshold for non-maximum suppression.")
+    score_thresh: float = Field(
+        default=0.5, ge=0, le=1, title="Minimal confidence for detection."
+    )
+    nms_thresh: float = Field(
+        default=0.45, ge=0, le=1, title="IoU threshold for non-maximum suppression."
+    )
+
 
 class Rknn(DetectionApi):
     type_key = DETECTOR_KEY
@@ -30,11 +36,12 @@ class Rknn(DetectionApi):
         self.model_path = config.model.path or "/models/yolov8n-320x320.rknn"
 
         from rknnlite.api import RKNNLite
+
         self.rknn = RKNNLite(verbose=False)
         if self.rknn.load_rknn(self.model_path) != 0:
-            logger.error('Error initializing rknn model.')
+            logger.error("Error initializing rknn model.")
         if self.rknn.init_runtime() != 0:
-            logger.error('Error initializing rknn runtime.')
+            logger.error("Error initializing rknn runtime.")
 
     def __del__(self):
         self.rknn.release()
@@ -50,22 +57,30 @@ class Rknn(DetectionApi):
         detections: array with shape (20, 6) with 20 rows of (class, confidence, y_min, x_min, y_max, x_max)
         """
 
-        results = np.transpose(results[0,:,:,0]) # array shape (2100, 84)
-        classes = np.argmax(results[:, 4:], axis=1) # array shape (2100,); index of class with max confidence of each row
-        scores = np.max(results[:, 4:], axis=1) # array shape (2100,); max confidence of each row
+        results = np.transpose(results[0, :, :, 0])  # array shape (2100, 84)
+        classes = np.argmax(
+            results[:, 4:], axis=1
+        )  # array shape (2100,); index of class with max confidence of each row
+        scores = np.max(
+            results[:, 4:], axis=1
+        )  # array shape (2100,); max confidence of each row
 
         # array shape (2100, 4); bounding box of each row
         boxes = np.transpose(
-            np.vstack((
-                results[:,0] - 0.5 * results[:,2],
-                results[:,1] - 0.5 * results[:,3],
-                results[:,2],
-                results[:,3])
+            np.vstack(
+                (
+                    results[:, 0] - 0.5 * results[:, 2],
+                    results[:, 1] - 0.5 * results[:, 3],
+                    results[:, 2],
+                    results[:, 3],
+                )
             )
         )
 
         # indices of rows with confidence > SCORE_THRESH with Non-maximum Suppression (NMS)
-        result_boxes = cv2.dnn.NMSBoxes(boxes, scores, self.score_thresh, self.nms_thresh, 0.5)
+        result_boxes = cv2.dnn.NMSBoxes(
+            boxes, scores, self.score_thresh, self.nms_thresh, 0.5
+        )
 
         detections = np.zeros((20, 6), np.float32)
 
@@ -80,7 +95,7 @@ class Rknn(DetectionApi):
                 (boxes[index][1]) / self.height,
                 (boxes[index][0]) / self.width,
                 (boxes[index][1] + boxes[index][3]) / self.height,
-                (boxes[index][0] + boxes[index][2]) / self.width
+                (boxes[index][0] + boxes[index][2]) / self.width,
             ]
 
         return detections
@@ -90,5 +105,9 @@ class Rknn(DetectionApi):
         return self.rknn.inference(inputs=tensor_input)
 
     def detect_raw(self, tensor_input):
-        output = self.inference([tensor_input,])
+        output = self.inference(
+            [
+                tensor_input,
+            ]
+        )
         return self.postprocess(output[0])