detection.py

import cv2
import numpy as np

import settings
from track import Track


def get_output_layers(net) -> list:
    layer_names = net.getLayerNames()
    output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
    return output_layers


def detect(image, net) -> list:
    width, height = image.shape[1], image.shape[0]

    scale = 0.00392  # constant scale for normalizing the color format
    blob = cv2.dnn.blobFromImage(image, scale, settings.YOLO_RESIZE_SIZE, (0, 0, 0), True, crop=False)
    net.setInput(blob)

    outs = net.forward(get_output_layers(net))

    detected_tracks = []
    confidences = []

    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            # confidence = np.max(scores)
            if class_id not in settings.YOLO_CLASSES_IDENTIFIERS:
                continue
            confidence = scores[class_id]
            if confidence > settings.CONFIDENCE_THRESHOLD:
                center_x = int(detection[0] * width)
                center_y = int(detection[1] * height)
                w = int(detection[2] * width)
                h = int(detection[3] * height)
                confidences.append(float(confidence))
                detected_tracks.append(Track(center_x, center_y, w, h))

    indices = cv2.dnn.NMSBoxes([track.rect for track in detected_tracks], confidences,
                               settings.CONFIDENCE_THRESHOLD, settings.NMS_THRESHOLD)

    return [track for i, track in enumerate(detected_tracks) if i in indices]