Create ObjectDetector.py

.ai/models/ObjectDetector.py

@@ -0,0 +1,42 @@
+import cv2
+import numpy as np
+
+class ObjectDetector:
+    def __init__(self, model_path):
+        self.model_path = model_path
+        self.net = cv2.dnn.readNetFromDarknet(self.model_path, self.model_path)
+
+    def detect_objects(self, image):
+        blob = cv2.dnn.blobFromImage(image, 1/255, (416, 416), [0,0,0], 1, crop=False)
+        self.net.setInput(blob)
+        outs = self.net.forward(self.getOutputsNames(self.net))
+        class_ids = []
+        confidences = []
+        boxes = []
+        for out in outs:
+            for detection in out:
+                scores = detection[5:]
+                class_id = np.argmax(scores)
+                confidence = scores[class_id]
+                if confidence > 0.5 and class_id == 0:
+                    # Filter out weak detections
+                    center_x = int(detection[0] * image.shape[1])
+                    center_y = int(detection[1] * image.shape[0])
+                    w = int(detection[2] * image.shape[1])
+                    h = int(detection[3] * image.shape[0])
+                    x = center_x - w / 2
+                    y = center_y - h / 2
+                    class_ids.append(class_id)
+                    confidences.append(float(confidence))
+                    boxes.append([x, y, w, h])
+        indices = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
+        for i in indices:
+            i = i[0]
+            box = boxes[i]
+            x, y, w, h = box[0], box[1], box[2], box[3]
+            cv2.rectangle(image, (int(x), int(y)), (int(x+w), int(y+h)), (255, 0, 0), 2)
+        return image
+
+    def getOutputsNames(self, net):
+        layersNames = net.getLayerNames()
+        return [layersNames[i[0] - 1] for i in net.getUnconnectedOutLayers()]