From b561f00ff9c07e99b5c745a0b394e099a0cc2547 Mon Sep 17 00:00:00 2001
From: Anil Ozyalcin <aeozyalcin@gmail.com>
Date: Thu, 16 Feb 2023 20:30:35 -0800
Subject: [PATCH] Initial commit that adds YOLOv5 and YOLOv8 support for
 OpenVINO detector

---
 frigate/detectors/detector_config.py  |  2 +
 frigate/detectors/plugins/openvino.py | 66 +++++++++++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/frigate/detectors/detector_config.py b/frigate/detectors/detector_config.py
index e321af087..25dede107 100644
--- a/frigate/detectors/detector_config.py
+++ b/frigate/detectors/detector_config.py
@@ -26,6 +26,8 @@ class InputTensorEnum(str, Enum):
 class ModelTypeEnum(str, Enum):
     ssd = "ssd"
     yolox = "yolox"
+    yolov5 = "yolov5"
+    yolov8 = "yolov8"
 
 
 class ModelConfig(BaseModel):
diff --git a/frigate/detectors/plugins/openvino.py b/frigate/detectors/plugins/openvino.py
index e04ac1352..9848232f7 100644
--- a/frigate/detectors/plugins/openvino.py
+++ b/frigate/detectors/plugins/openvino.py
@@ -133,3 +133,69 @@ class OvDetector(DetectionApi):
                 else:
                     break
             return detections
+        elif self.ov_model_type == ModelTypeEnum.yolov8:
+            infer_request = self.interpreter.create_infer_request()
+            infer_request.infer([tensor_input])
+            out_tensor = infer_request.get_output_tensor()
+            results = out_tensor.data[0]
+            output_data = np.transpose(results)
+            scores = np.max(output_data[:, 4:], axis=1)
+            if len(scores) == 0:
+                return np.zeros((20, 6), np.float32)
+
+            scores = np.expand_dims(scores, axis=1)
+            dets = np.concatenate((output_data, scores), axis=1) # add scores to the last column
+            dets = dets[dets[:,-1] > 0.5,:] # filter out lines with scores below threshold
+
+            ordered = dets[dets[:, -1].argsort()[::-1]][:20] # limit to top 20 scores, descending order
+            detections = np.zeros((20, 6), np.float32)
+            i = 0
+
+            for object_detected in ordered:
+                if i < 20:
+                    detections[i] = [
+                        np.argmax(object_detected[4:-1])  ,  # Label ID
+                        object_detected[-1],  # Confidence
+                        (object_detected[1] - (object_detected[3] / 2))
+                        / self.h,  # y_min
+                        (object_detected[0] - (object_detected[2] / 2))
+                        / self.w,  # x_min
+                        (object_detected[1] + (object_detected[3] / 2))
+                        / self.h,  # y_max
+                        (object_detected[0] + (object_detected[2] / 2))
+                        / self.w,  # x_max
+                    ]
+                    i += 1
+                else:
+                    break              
+            return detections
+        elif self.ov_model_type == ModelTypeEnum.yolov5:
+            infer_request = self.interpreter.create_infer_request()
+            infer_request.infer([tensor_input])
+            out_tensor = infer_request.get_output_tensor()
+            output_data = out_tensor.data[0]
+            conf_mask = (output_data[:, 4] >= 0.5).squeeze()
+            output_data = output_data[conf_mask]
+            ordered = output_data[output_data[:, 4].argsort()[::-1]][:20] # limit to top 20 scores, descending order
+
+            detections = np.zeros((20, 6), np.float32)
+            i = 0
+
+            for object_detected in ordered:
+                if i < 20:
+                    detections[i] = [
+                        np.argmax(object_detected[5:])  ,  # Label ID
+                        object_detected[4],  # Confidence
+                        (object_detected[1] - (object_detected[3] / 2))
+                        / self.h,  # y_min
+                        (object_detected[0] - (object_detected[2] / 2))
+                        / self.w,  # x_min
+                        (object_detected[1] + (object_detected[3] / 2))
+                        / self.h,  # y_max
+                        (object_detected[0] + (object_detected[2] / 2))
+                        / self.w,  # x_max
+                    ]
+                    i += 1
+                else:
+                    break
+            return detections
\ No newline at end of file