scripts/libraries: Vectorize NMS using numpy.

The new NMS algorithm with numpy can now handle a large number of
bounding box candidates without choking. Note that FOMO had to be
rewritten to be compatible with the new NMS algorithm.
This commit is contained in:
Kwabena W. Agyeman 2025-06-14 23:39:31 -07:00
parent 9958d8ee62
commit fafa8271ec
3 changed files with 206 additions and 213 deletions

View File

@ -9,9 +9,8 @@
import sensor import sensor
import time import time
import ml import ml
from ml.utils import NMS from ml.postprocessing import fomo_postprocess
import math import math
import image
sensor.reset() # Reset and initialize the sensor. sensor.reset() # Reset and initialize the sensor.
sensor.set_pixformat(sensor.RGB565) # Set pixel format to RGB565 (or GRAYSCALE) sensor.set_pixformat(sensor.RGB565) # Set pixel format to RGB565 (or GRAYSCALE)
@ -19,9 +18,6 @@ sensor.set_framesize(sensor.QVGA) # Set frame size to QVGA (320x240)
sensor.set_windowing((240, 240)) # Set 240x240 window. sensor.set_windowing((240, 240)) # Set 240x240 window.
sensor.skip_frames(time=2000) # Let the camera adjust. sensor.skip_frames(time=2000) # Let the camera adjust.
min_confidence = 0.4
threshold_list = [(math.ceil(min_confidence * 255), 255)]
# Load built-in FOMO face detection model # Load built-in FOMO face detection model
model = ml.Model("/rom/fomo_face_detection.tflite") model = ml.Model("/rom/fomo_face_detection.tflite")
print(model) print(model)
@ -40,31 +36,7 @@ colors = [ # Add more colors if you are detecting more than 7 types of classes
(255, 255, 255), (255, 255, 255),
] ]
fomo = fomo_postprocess(threshold=0.4)
# FOMO outputs an image per class where each pixel in the image is the centroid of the trained
# object. So, we will get those output images and then run find_blobs() on them to extract the
# centroids. We will also run get_stats() on the detected blobs to determine their score.
# The Non-Max-Supression (NMS) object then filters out overlapping detections and maps their
# position in the output image back to the original input image. The function then returns a
# list per class which each contain a list of (rect, score) tuples representing the detected
# objects.
def fomo_post_process(model, inputs, outputs):
n, oh, ow, oc = model.output_shape[0]
nms = NMS(ow, oh, inputs[0].roi)
for i in range(oc):
img = image.Image(outputs[0][0, :, :, i] * 255)
blobs = img.find_blobs(
threshold_list, x_stride=1, area_threshold=1, pixels_threshold=1
)
for b in blobs:
rect = b.rect()
x, y, w, h = rect
score = (
img.get_statistics(thresholds=threshold_list, roi=rect).l_mean() / 255.0
)
nms.add_bounding_box(x, y, x + w, y + h, score, i)
return nms.get_bounding_boxes()
clock = time.clock() clock = time.clock()
while True: while True:
@ -72,7 +44,7 @@ while True:
img = sensor.snapshot() img = sensor.snapshot()
for i, detection_list in enumerate(model.predict([img], callback=fomo_post_process)): for i, detection_list in enumerate(model.predict([img], callback=fomo)):
if i == 0: if i == 0:
continue # background class continue # background class
if len(detection_list) == 0: if len(detection_list) == 0:

View File

@ -26,14 +26,30 @@
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import math import ml.utils
import image
from ml.utils import NMS
from micropython import const from micropython import const
from ulab import numpy as np from ulab import numpy as np
_NO_DETECTION = const(()) _NO_DETECTION = const(())
_FOMO_CLASSES = const(1)
_YOLO_V2_TX = const(0)
_YOLO_V2_TY = const(1)
_YOLO_V2_TW = const(2)
_YOLO_V2_TH = const(3)
_YOLO_V2_SCORE = const(4)
_YOLO_V2_CLASSES = const(5)
_YOLO_V5_CX = const(0)
_YOLO_V5_CY = const(1)
_YOLO_V5_CW = const(2)
_YOLO_V5_CH = const(3)
_YOLO_V5_SCORE = const(4)
_YOLO_V5_CLASSES = const(5)
_YOLO_V8_CX = const(0)
_YOLO_V8_CY = const(1)
_YOLO_V8_CW = const(2)
_YOLO_V8_CH = const(3)
_YOLO_V8_CLASSES = const(4)
def dequantize(value, dtype, zero_point, scale): def dequantize(value, dtype, zero_point, scale):
@ -42,48 +58,74 @@ def dequantize(value, dtype, zero_point, scale):
return (value - zero_point) * scale return (value - zero_point) * scale
# FOMO generates an image per class, where each pixel represents the centroid def sigmoid(x):
# of the trained object. These images are processed with `find_blobs()` to return 1.0 / (1.0 + np.exp(-x))
# extract centroids, and `get_stats()` is used to get their scores. Overlapping
# detections are then filtered with NMS and positions are mapped back to the
# original image, and a list of (rect, score) tuples is returned for each class, def mod(a, b):
# representing detected objects. return a - (b * (a // b))
def softmax(x):
e_x = np.exp(x - np.max(x, axis=1, keepdims=True))
return e_x / np.sum(e_x, axis=1, keepdims=True)
class fomo_postprocess: class fomo_postprocess:
def __init__(self, threshold=0.4): def __init__(self, threshold=0.4, w_scale=1.414214, h_scale=1.414214,
self.threshold_list = [(math.ceil(threshold * 255), 255)] nms_threshold=0.1, nms_sigma=0.001,
scale_aspect=ml.utils.NMS_SCALE_ASPECT_KEEP):
self.threshold = threshold
self.w_scale = w_scale
self.h_scale = h_scale
self.nms_threshold = nms_threshold
self.nms_sigma = nms_sigma
self.scale_aspect = scale_aspect
def __call__(self, model, inputs, outputs): def __call__(self, model, inputs, outputs):
n, oh, ow, oc = model.output_shape[0] ob, oh, ow, oc = model.output_shape[0]
s = model.output_scale[0] s = model.output_scale[0]
zp = model.output_zero_point[0] zp = model.output_zero_point[0]
dt = model.output_dtype[0] dt = model.output_dtype[0]
nms = NMS(ow, oh, inputs[0].roi)
for i in range(oc): # Reshape the output to a 2D array
img = image.Image(dequantize(outputs[0][0, :, :, i], dt, zp, s) * 255) row_outputs = outputs[0].reshape((oh * ow, oc))
blobs = img.find_blobs(
self.threshold_list, x_stride=1, area_threshold=1, pixels_threshold=1 # Threshold all the scores
) score_indices = np.max(dequantize(row_outputs[:, _FOMO_CLASSES:], dt, zp, s), axis=1)
for b in blobs: score_indices = np.nonzero(score_indices > self.threshold)[0]
rect = b.rect() if not len(score_indices):
x, y, w, h = rect return _NO_DETECTION
score = (
img.get_statistics(thresholds=self.threshold_list, roi=rect).l_mean() / 255.0 # Get the bounding boxes that have a valid score
) bb = dequantize(np.take(row_outputs, score_indices, axis=0), dt, zp, s)
nms.add_bounding_box(x, y, x + w, y + h, score, i)
return nms.get_bounding_boxes() # Extract rows and columns
bb_rows = score_indices // ow
bb_cols = mod(score_indices, ow)
# Get the score information
bb_scores = np.max(bb[:, _FOMO_CLASSES:], axis=1)
# Get the class information
bb_classes = np.argmax(bb[:, _FOMO_CLASSES:], axis=1) + _FOMO_CLASSES
# Compute the bounding box information
x_center = (bb_cols + 0.5) / ow
y_center = (bb_rows + 0.5) / oh
w_rel = np.full(len(bb_cols), self.w_scale / ow)
h_rel = np.full(len(bb_rows), self.h_scale / oh)
return ml.utils.box_nms(x_center, y_center, w_rel, h_rel, bb_scores,
bb_classes, model.input_shape[0][1:3], inputs[0].roi,
self.nms_threshold, self.nms_sigma, self.scale_aspect)
# This is a lightweight version of the tiny yolo v2 object detection algorithm. # This is a lightweight version of the tiny yolo v2 object detection algorithm.
# It was optimized to work well on embedded devices with limited computational resources. # It was optimized to work well on embedded devices with limited computational resources.
class yolo_v2_postprocess: class yolo_v2_postprocess:
_YOLO_V2_TX = const(0) def __init__(self, threshold=0.6, anchors=None, nms_threshold=0.1, nms_sigma=0.1,
_YOLO_V2_TY = const(1) scale_aspect=ml.utils.NMS_SCALE_ASPECT_KEEP):
_YOLO_V2_TW = const(2)
_YOLO_V2_TH = const(3)
_YOLO_V2_SCORE = const(4)
_YOLO_V2_CLASSES = const(5)
def __init__(self, threshold=0.6, anchors=None, nms_threshold=0.1, nms_sigma=0.1):
self.threshold = threshold self.threshold = threshold
self.anchors = anchors self.anchors = anchors
if self.anchors is None: if self.anchors is None:
@ -95,6 +137,7 @@ class yolo_v2_postprocess:
self.anchors_len = len(self.anchors) self.anchors_len = len(self.anchors)
self.nms_threshold = nms_threshold self.nms_threshold = nms_threshold
self.nms_sigma = nms_sigma self.nms_sigma = nms_sigma
self.scale_aspect = scale_aspect
def __call__(self, model, inputs, outputs): def __call__(self, model, inputs, outputs):
ob, oh, ow, oc = model.output_shape[0] ob, oh, ow, oc = model.output_shape[0]
@ -103,16 +146,6 @@ class yolo_v2_postprocess:
dt = model.output_dtype[0] dt = model.output_dtype[0]
class_count = (oc // self.anchors_len) - _YOLO_V2_CLASSES class_count = (oc // self.anchors_len) - _YOLO_V2_CLASSES
def sigmoid(x):
return 1.0 / (1.0 + np.exp(-x))
def mod(a, b):
return a - (b * (a // b))
def softmax(x):
e_x = np.exp(x - np.max(x, axis=1, keepdims=True))
return e_x / np.sum(e_x, axis=1, keepdims=True)
# Reshape the output to a 2D array # Reshape the output to a 2D array
row_outputs = outputs[0].reshape((oh * ow * self.anchors_len, row_outputs = outputs[0].reshape((oh * ow * self.anchors_len,
_YOLO_V2_CLASSES + class_count)) _YOLO_V2_CLASSES + class_count))
@ -146,48 +179,32 @@ class yolo_v2_postprocess:
w_rel = (bb_a_array[:, 0] * np.exp(bb[:, _YOLO_V2_TW])) / ow w_rel = (bb_a_array[:, 0] * np.exp(bb[:, _YOLO_V2_TW])) / ow
h_rel = (bb_a_array[:, 1] * np.exp(bb[:, _YOLO_V2_TH])) / oh h_rel = (bb_a_array[:, 1] * np.exp(bb[:, _YOLO_V2_TH])) / oh
# Scale the bounding boxes to have enough integer precision for NMS return ml.utils.box_nms(x_center, y_center, w_rel, h_rel, bb_scores, bb_classes,
ib, ih, iw, ic = model.input_shape[0] model.input_shape[0][1:3], inputs[0].roi,
x_center = x_center * iw self.nms_threshold, self.nms_sigma, self.scale_aspect)
y_center = y_center * ih
w_rel = w_rel * iw
h_rel = h_rel * ih
nms = NMS(iw, ih, inputs[0].roi)
for i in range(bb.shape[0]):
nms.add_bounding_box(x_center[i] - (w_rel[i] / 2),
y_center[i] - (h_rel[i] / 2),
x_center[i] + (w_rel[i] / 2),
y_center[i] + (h_rel[i] / 2),
bb_scores[i], bb_classes[i])
return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)
# This is a lightweight version of the YOLO (You Only Look Once) object detection algorithm. # This is a lightweight version of the YOLO (You Only Look Once) object detection algorithm.
# It is designed to work well on embedded devices with limited computational resources. # It is designed to work well on embedded devices with limited computational resources.
class yolo_lc_postprocess(yolo_v2_postprocess): class yolo_lc_postprocess(yolo_v2_postprocess):
def __init__(self, threshold=0.6, anchors=None, nms_threshold=0.1, nms_sigma=0.1): def __init__(self, threshold=0.6, anchors=None, nms_threshold=0.1, nms_sigma=0.1,
scale_aspect=ml.utils.NMS_SCALE_ASPECT_KEEP):
if anchors is None: if anchors is None:
anchors = np.array([[0.076023, 0.258508], anchors = np.array([[0.076023, 0.258508],
[0.163031, 0.413531], [0.163031, 0.413531],
[0.234769, 0.702585], [0.234769, 0.702585],
[0.427054, 0.715892], [0.427054, 0.715892],
[0.748154, 0.857092]]) [0.748154, 0.857092]])
super().__init__(threshold, anchors, nms_threshold, nms_sigma) super().__init__(threshold, anchors, nms_threshold, nms_sigma, scale_aspect)
class yolo_v5_postprocess: class yolo_v5_postprocess:
_YOLO_V5_CX = const(0) def __init__(self, threshold=0.6, nms_threshold=0.1, nms_sigma=0.1,
_YOLO_V5_CY = const(1) scale_aspect=ml.utils.NMS_SCALE_ASPECT_KEEP):
_YOLO_V5_CW = const(2)
_YOLO_V5_CH = const(3)
_YOLO_V5_SCORE = const(4)
_YOLO_V5_CLASSES = const(5)
def __init__(self, threshold=0.6, nms_threshold=0.1, nms_sigma=0.1):
self.threshold = threshold self.threshold = threshold
self.nms_threshold = nms_threshold self.nms_threshold = nms_threshold
self.nms_sigma = nms_sigma self.nms_sigma = nms_sigma
self.scale_aspect = scale_aspect
def __call__(self, model, inputs, outputs): def __call__(self, model, inputs, outputs):
oh, ow, oc = model.output_shape[0] oh, ow, oc = model.output_shape[0]
@ -217,34 +234,21 @@ class yolo_v5_postprocess:
# Compute the bounding box information # Compute the bounding box information
x_center = bb[:, _YOLO_V5_CX] x_center = bb[:, _YOLO_V5_CX]
y_center = bb[:, _YOLO_V5_CY] y_center = bb[:, _YOLO_V5_CY]
w_rel = bb[:, _YOLO_V5_CW] * 0.5 w_rel = bb[:, _YOLO_V5_CW]
h_rel = bb[:, _YOLO_V5_CH] * 0.5 h_rel = bb[:, _YOLO_V5_CH]
# Scale the bounding boxes to have enough integer precision for NMS return ml.utils.box_nms(x_center, y_center, w_rel, h_rel, bb_scores, bb_classes,
ib, ih, iw, ic = model.input_shape[0] model.input_shape[0][1:3], inputs[0].roi,
xmin = (x_center - w_rel) * iw self.nms_threshold, self.nms_sigma, self.scale_aspect)
ymin = (y_center - h_rel) * ih
xmax = (x_center + w_rel) * iw
ymax = (y_center + h_rel) * ih
nms = NMS(iw, ih, inputs[0].roi)
for i in range(bb.shape[0]):
nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i],
bb_scores[i], bb_classes[i])
return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)
class yolo_v8_postprocess: class yolo_v8_postprocess:
_YOLO_V8_CX = const(0) def __init__(self, threshold=0.6, nms_threshold=0.1, nms_sigma=0.1,
_YOLO_V8_CY = const(1) scale_aspect=ml.utils.NMS_SCALE_ASPECT_KEEP):
_YOLO_V8_CW = const(2)
_YOLO_V8_CH = const(3)
_YOLO_V8_CLASSES = const(4)
def __init__(self, threshold=0.6, nms_threshold=0.1, nms_sigma=0.1):
self.threshold = threshold self.threshold = threshold
self.nms_threshold = nms_threshold self.nms_threshold = nms_threshold
self.nms_sigma = nms_sigma self.nms_sigma = nms_sigma
self.scale_aspect = scale_aspect
def __call__(self, model, inputs, outputs): def __call__(self, model, inputs, outputs):
oh, ow, oc = model.output_shape[0] oh, ow, oc = model.output_shape[0]
@ -274,18 +278,9 @@ class yolo_v8_postprocess:
# Compute the bounding box information # Compute the bounding box information
x_center = bb[_YOLO_V8_CX, :] x_center = bb[_YOLO_V8_CX, :]
y_center = bb[_YOLO_V8_CY, :] y_center = bb[_YOLO_V8_CY, :]
w_rel = bb[_YOLO_V8_CW, :] * 0.5 w_rel = bb[_YOLO_V8_CW, :]
h_rel = bb[_YOLO_V8_CH, :] * 0.5 h_rel = bb[_YOLO_V8_CH, :]
# Scale the bounding boxes to have enough integer precision for NMS return ml.utils.box_nms(x_center, y_center, w_rel, h_rel, bb_scores, bb_classes,
ib, ih, iw, ic = model.input_shape[0] model.input_shape[0][1:3], inputs[0].roi,
xmin = (x_center - w_rel) * iw self.nms_threshold, self.nms_sigma, self.scale_aspect)
ymin = (y_center - h_rel) * ih
xmax = (x_center + w_rel) * iw
ymax = (y_center + h_rel) * ih
nms = NMS(iw, ih, inputs[0].roi)
for i in range(bb.shape[1]):
nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i],
bb_scores[i], bb_classes[i])
return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)

View File

@ -26,98 +26,124 @@
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import math from micropython import const
from ulab import numpy as np
class NMS: NMS_SCALE_ASPECT_KEEP = const(0)
def __init__( NMS_SCALE_ASPECT_EXPAND = const(1)
self, NMS_SCALE_ASPECT_IGNORE = const(2)
window_w,
window_h,
roi, def box_nms(
x_center, # shape: (N,)
y_center, # shape: (N,)
w_rel, # shape: (N,)
h_rel, # shape: (N,)
scores, # shape: (N,)
classes, # shape: (N,)
input_shape, # (input_h, input_w)
image_roi, # (image_x, image_y, image_w, image_h)
threshold=0.1, # IoU threshold for NMS
sigma=0.1, # Sigma for NMS
scale_aspect=NMS_SCALE_ASPECT_KEEP # Scale aspect ratio
): ):
self.window_w = window_w N = len(classes)
self.window_h = window_h assert all(N == len(arr) for arr in [x_center, y_center, w_rel, h_rel, scores]), \
self.roi = roi f"Inconsistent lengths: classes={N}, x_center={len(x_center)}, " \
if roi[2] < 1 or roi[3] < 1: f"y_center={len(y_center)}, w_rel={len(w_rel)}, h_rel={len(h_rel)}, " \
raise ValueError("Invalid ROI dimensions!") f"scores={len(scores)}"
self.boxes = []
def add_bounding_box(self, xmin, ymin, xmax, ymax, score, label_index): input_h, input_w = input_shape
if score >= 0.0 and score <= 1.0: roi_x, roi_y, roi_w, roi_h = image_roi
xmin = max(0.0, min(xmin, self.window_w))
ymin = max(0.0, min(ymin, self.window_h))
xmax = max(0.0, min(xmax, self.window_w))
ymax = max(0.0, min(ymax, self.window_h))
w = int(xmax - xmin)
h = int(ymax - ymin)
if w > 0 and h > 0:
self.boxes.append([int(xmin), int(ymin), w, h, score, label_index])
def get_bounding_boxes(self, threshold=0.1, sigma=0.1): sigma_inv = (-1.0 / sigma) if (sigma > 0.0) else 0.0
sorted_boxes = sorted(self.boxes, key=lambda x: x[4], reverse=True)
sigma_scale = (-1.0 / sigma) if (sigma > 0.0) else 0.0
def iou(box1, box2): x_scale = roi_w / float(input_w)
x1 = max(box1[0], box2[0]) y_scale = roi_h / float(input_h)
y1 = max(box1[1], box2[1])
x2 = min(box1[0] + box1[2], box2[0] + box2[2])
y2 = min(box1[1] + box1[3], box2[1] + box2[3])
w = max(0, x2 - x1)
h = max(0, y2 - y1)
intersection = w * h
union = (box1[2] * box1[3]) + (box2[2] * box2[3]) - intersection
return float(intersection) / float(union)
# Perform Non Max Supression. if scale_aspect == NMS_SCALE_ASPECT_KEEP:
max_index = 0
output_boxes = []
max_label_index = 0
while len(sorted_boxes):
box = sorted_boxes.pop(max_index)
output_boxes.append(box)
max_label_index = max(max_label_index, box[5])
# Compare and supress the remaining boxes in the list against the max.
for i in range(len(sorted_boxes)):
v = iou(box, sorted_boxes[i])
sorted_boxes[i][4] = sorted_boxes[i][4] * math.exp(sigma_scale * v * v)
if sorted_boxes[i][4] < threshold:
sorted_boxes[i][4] = 0.0
# Filter out supressed boxes and find the next largest.
sorted_boxes = list(filter(lambda x: x[4] > 0.0, sorted_boxes))
if len(sorted_boxes):
max_index = max(enumerate(sorted_boxes), key=lambda x: x[1][4])[0]
# Map the output boxes back to the input image.
x_scale = self.roi[2] / float(self.window_w)
y_scale = self.roi[3] / float(self.window_h)
scale = min(x_scale, y_scale) scale = min(x_scale, y_scale)
x_offset = ((self.roi[2] - (self.window_w * scale)) / 2) + self.roi[0] x_scale = scale
y_offset = ((self.roi[3] - (self.window_h * scale)) / 2) + self.roi[1] y_scale = scale
elif scale_aspect == NMS_SCALE_ASPECT_EXPAND:
scale = max(x_scale, y_scale)
x_scale = scale
y_scale = scale
elif scale_aspect != NMS_SCALE_ASPECT_IGNORE:
raise ValueError("Invalid scale_aspect value!")
for i in range(len(output_boxes)): x_offset = ((roi_w - (input_w * x_scale)) * 0.5) + roi_x
output_boxes[i][0] = int((output_boxes[i][0] * scale) + x_offset) y_offset = ((roi_h - (input_h * y_scale)) * 0.5) + roi_y
output_boxes[i][1] = int((output_boxes[i][1] * scale) + y_offset)
output_boxes[i][2] = int(output_boxes[i][2] * scale)
output_boxes[i][3] = int(output_boxes[i][3] * scale)
# Create a list per class with (rect, score) tuples. # Convert boxes to (x1, y1, x2, y2) format.
w_rel_2 = w_rel * 0.5
h_rel_2 = h_rel * 0.5
x1 = (x_center - w_rel_2) * input_w
y1 = (y_center - h_rel_2) * input_h
x2 = (x_center + w_rel_2) * input_w
y2 = (y_center + h_rel_2) * input_h
areas = (x2 - x1) * (y2 - y1)
boxes = np.array([x1, y1, x2, y2, areas]).T
output_list = [[] for i in range(max_label_index + 1)] # Allocate output list for each class.
output = [[] for _ in range(int(np.max(classes)) + 1)]
for i in range(len(output_boxes)): # Filter out invalid boxes.
output_list[output_boxes[i][5]].append( valid_indices = np.nonzero((areas > 0.0) & (scores > threshold) & (scores <= 1.0))[0]
(output_boxes[i][0:4], output_boxes[i][4]) if not len(valid_indices):
) return output
return output_list # Sort boxes by scores in descending order.
valid_scores = np.take(scores, valid_indices, axis=0)
sorted_valid_score_indices = np.argsort(valid_scores, axis=0)[::-1]
sorted_valid_indices = np.take(valid_indices, sorted_valid_score_indices, axis=0)
while True:
# Grab the box with the highest score.
i = sorted_valid_indices[0]
x1i, y1i, x2i, y2i, area = boxes[i]
# Project and store the box.
px = round((x1i * x_scale) + x_offset)
py = round((y1i * y_scale) + y_offset)
pw = round((x2i - x1i) * x_scale)
ph = round((y2i - y1i) * y_scale)
output[classes[i]].append(((px, py, pw, ph), scores[i]))
# Stop if there's only one box left.
if len(sorted_valid_indices) == 1:
break
# Get the rest of the boxes.
sorted_valid_indices = sorted_valid_indices[1:]
boxes = np.take(boxes, sorted_valid_indices, axis=0)
scores = np.take(scores, sorted_valid_indices, axis=0)
classes = np.take(classes, sorted_valid_indices, axis=0)
# Compute IoU of the max box with the rest.
xx1 = np.maximum(x1i, boxes[:, 0])
yy1 = np.maximum(y1i, boxes[:, 1])
xx2 = np.minimum(x2i, boxes[:, 2])
yy2 = np.minimum(y2i, boxes[:, 3])
iw = np.maximum(0.0, xx2 - xx1)
ih = np.maximum(0.0, yy2 - yy1)
intersection = iw * ih
union = area + boxes[:, 4] - intersection
iou = intersection / (union + 1e-6)
scores *= np.exp((iou ** 2.0) * sigma_inv)
# Filter out boxes with low scores.
valid_indices = np.nonzero(scores > threshold)[0]
if not len(valid_indices):
break
# Sort boxes by scores in descending order.
valid_scores = np.take(scores, valid_indices, axis=0)
sorted_valid_score_indices = np.argsort(valid_scores, axis=0)[::-1]
sorted_valid_indices = np.take(valid_indices, sorted_valid_score_indices, axis=0)
return output
def draw_predictions( def draw_predictions(