Added files

detector.py

@@ -0,0 +1,113 @@
+import numpy as np
+import cv2
+from pathlib import Path
+
+
+class Detector():
+    def __init__(self, weights="yolov4-tiny.weights", config_file = "tiny-yolov4.cfg",
+                 classes_file=Path(__file__).parent / 'cfg' / 'coco.names',  conf_thresh=0.25, nms_thresh=0.4,
+                 netsize=[416, 416]):
+
+        """ Detector class that implements yolo detector
+              Parameters
+              ----------
+              weights : string
+                  path to the weights file, eg. yolov4-tiny.weights
+              config_file : string
+                  path to the configs file, eg. tiny-yolov4.cfg
+              classes_file : string
+                  path to the classes file, eg. coco.names
+              conf_thresh : float
+                  confidence threshold for the detection
+              nms_thresh : float
+                  non maxima supression threshold for the detections
+              netsize : tuple
+                  netwrork's config size eg. [416,416]
+              """
+        self.weights = weights
+        self.config = config_file
+        self.classes_file = classes_file
+        # initialize the network
+        self.net = cv2.dnn.readNet(self.weights, self.config)
+        # self.net  = cv2.dnn.readNetFromDarknet(self.config, self.weights)
+        # set the backend
+        self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
+        self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
+
+        ln = self.net.getLayerNames()
+        # an error checking due to OPENCV versions changes
+        try:
+            self.vers = 0
+            self.layers = [ln[i[0] - 1] for i in self.net.getUnconnectedOutLayers()]
+        except:
+            self.vers = 1
+            self.layers = [ln[i - 1] for i in self.net.getUnconnectedOutLayers()]
+
+        self.netsize = netsize
+        self.conf_thresh = conf_thresh
+        self.nms_thresh = nms_thresh
+        with open(self.classes_file, 'r') as f:
+            classes = [line.strip() for line in f.readlines()]
+        self.labels = np.array(classes, dtype=str)
+        self.model = cv2.dnn_DetectionModel(self.net)
+        self.model.setInputParams(size=(self.netsize[0], self.netsize[1]), scale=1 / 256)
+        self.COLORS = np.random.randint(0, 255, size=(len(self.labels), 3))
+
+    def detect(self, img_or):
+
+        # copy the original image
+        img = img_or.copy()
+        # do some pre-processing (eg. resizing)
+        image = self._preprocess(img)
+        try:
+            # detect the image
+            classes, scores, boxes = self.model.detect(image, self.conf_thresh,
+                                                       self.nms_thresh)
+        except Exception as e:
+            print("[ERROR ]" + e)
+        if len(boxes)==0:
+           return img, [], [], []
+
+        # de-normalize the results
+        fixboxes = self.de_normalize(boxes,self.netsize, self.imgsize)
+        boxes = np.array(fixboxes,dtype=np.int32)
+        # draw the detections on the image
+        image_det = self.postprocess(classes, scores, boxes, img)
+
+        return image_det, classes, scores, boxes
+
+    def _preprocess(self, frame):
+        self.imgsize = [frame.shape[1], frame.shape[0]]
+        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        return cv2.resize(img, (self.netsize[0], self.netsize[1]), interpolation=cv2.INTER_LINEAR)
+
+    def de_normalize(self, boxes,netsize,imgsize):
+        finBoxes = []
+        for i in range(len(boxes)):
+            box = boxes[i]
+            x,y,w,h= int((box[0] / netsize[0]) * imgsize[0]),int((box[1] / netsize[1]) * imgsize[1]),\
+                     int((box[2] / netsize[0]) * imgsize[0]), int((box[3] / netsize[1]) * imgsize[1])
+            finBoxes.append((x,y,w,h))
+        return finBoxes
+
+    # drawing bounding boxes and their label on the detections
+    def postprocess(self, classes, scores, boxes, img):
+        for i in range(len(boxes)):
+            box = boxes[i]
+            if scores[i] > self.conf_thresh:
+                x,y,w,h = box
+                if self.vers == 1:
+                    class_id = classes[i]
+                    score = scores[i]
+                    color = [int(c) for c in self.COLORS[class_id]]
+                else:
+                    class_id = classes[i][0]
+                    score = scores[i][0]
+                    color = self.COLORS[class_id]
+                    color = (int(color[0]), int(color[1]), int(color[2]))
+                cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
+                text = "{}: {:.4f}".format(self.labels[class_id], score)
+                cv2.putText(img, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX,
+                            0.5, color, 2)
+        return img
+

packages.txt

@@ -0,0 +1,2 @@
+python3-opencv
+ffmpeg

requirements.txt

@@ -0,0 +1,5 @@
+opencv-python
+ffmpeg-python
+numpy
+pandas
+matplotlib