54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
# Copyright © 2020 Arm Ltd and Contributors. All rights reserved.
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
"""
|
|
Contains functions specific to decoding and processing inference results for SSD Mobilenet V1 models.
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def ssd_processing(output: np.ndarray, confidence_threshold=0.60):
|
|
"""
|
|
Gets class, bounding box positions and confidence from the four outputs of the SSD model.
|
|
|
|
Args:
|
|
output: Vector of outputs from network.
|
|
confidence_threshold: Selects only strong detections above this value.
|
|
|
|
Returns:
|
|
A list of detected objects in the form [class, [box positions], confidence]
|
|
"""
|
|
if len(output) != 4:
|
|
raise RuntimeError('Number of outputs from SSD model does not equal 4')
|
|
|
|
position, classification, confidence, num_detections = [index[0] for index in output]
|
|
|
|
detections = []
|
|
for i in range(int(num_detections)):
|
|
if confidence[i] > confidence_threshold:
|
|
class_idx = classification[i]
|
|
box = position[i, :4]
|
|
# Reorder positions in format [x_min, y_min, x_max, y_max]
|
|
box[0], box[1], box[2], box[3] = box[1], box[0], box[3], box[2]
|
|
confidence_value = confidence[i]
|
|
detections.append((class_idx, box, confidence_value))
|
|
return detections
|
|
|
|
|
|
def ssd_resize_factor(video: cv2.VideoCapture):
|
|
"""
|
|
Gets a multiplier to scale the bounding box positions to
|
|
their correct position in the frame.
|
|
|
|
Args:
|
|
video: Video capture object, contains information about data source.
|
|
|
|
Returns:
|
|
Resizing factor to scale box coordinates to output frame size.
|
|
"""
|
|
frame_height = video.get(cv2.CAP_PROP_FRAME_HEIGHT)
|
|
frame_width = video.get(cv2.CAP_PROP_FRAME_WIDTH)
|
|
return max(frame_height, frame_width)
|