scripts/libraries: Move nms to python.

This commit is contained in:
Kwabena W. Agyeman 2024-07-07 15:25:55 -07:00
parent 068c7c11f3
commit 3f8491cb0e
5 changed files with 104 additions and 145 deletions

View File

@ -9,6 +9,7 @@
import sensor
import time
import ml
from ml.nms import NMS
import math
import image
@ -47,7 +48,7 @@ colors = [ # Add more colors if you are detecting more than 7 types of classes
# objects.
def fomo_post_process(model, inputs, outputs):
n, oh, ow, oc = model.output_shape[0]
nms = ml.NMS(ow, oh, inputs[0].roi)
nms = NMS(ow, oh, inputs[0].roi)
for i in range(oc):
img = image.Image(outputs[0], shape=(oh, ow, 1), strides=(i, oc), scale=(255, 0))
blobs = img.find_blobs(

View File

@ -7,5 +7,4 @@
#
# This is an extension package to the ml C user-module.
from uml import NMS # noqa
from .model import * # noqa

View File

@ -0,0 +1,102 @@
# This file is part of the OpenMV project.
#
# Copyright (c) 2024 Ibrahim Abdelkader <iabdalkader@openmv.io>
# Copyright (c) 2024 Kwabena W. Agyeman <kwagyeman@openmv.io>
#
# This work is licensed under the MIT license, see the file LICENSE for details.
import math
class NMS:
def __init__(
self,
window_w,
window_h,
roi,
):
self.window_w = window_w
self.window_h = window_h
self.roi = roi
if roi[2] < 1 or roi[3] < 1:
raise ValueError("Invalid ROI dimensions!")
self.boxes = []
@micropython.native
def add_bounding_box(self, xmin, ymin, xmax, ymax, score, label_index):
if score >= 0.0 and score <= 1.0:
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])
@micropython.native
def get_bounding_boxes(self, threshold=0.1, sigma=0.1):
sorted_boxes = sorted(self.boxes, key=lambda x: x[4], reverse=True)
sigma_scale = (-1.0 / sigma) if (sigma > 0.0) else 0.0
@micropython.native
def iou(box1, box2):
x1 = max(box1[0], box2[0])
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.
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)
x_offset = ((self.roi[2] - (self.window_w * scale)) / 2) + self.roi[0]
y_offset = ((self.roi[3] - (self.window_h * scale)) / 2) + self.roi[1]
for i in range(len(output_boxes)):
output_boxes[i][0] = int((output_boxes[i][0] * scale) + x_offset)
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.
output_list = [[] for i in range(max_label_index + 1)]
for i in range(len(output_boxes)):
output_list[output_boxes[i][5]].append(
(output_boxes[i][0:4], output_boxes[i][4])
)
return output_list

View File

@ -327,7 +327,6 @@ extern const mp_obj_type_t py_ml_nms_type;
static const mp_rom_map_elem_t py_ml_globals_dict_table[] = {
{ MP_ROM_QSTR(MP_QSTR___name__), MP_OBJ_NEW_QSTR(MP_QSTR_ml) },
{ MP_ROM_QSTR(MP_QSTR_Model), MP_ROM_PTR(&py_ml_model_type) },
{ MP_ROM_QSTR(MP_QSTR_NMS), MP_ROM_PTR(&py_ml_nms_type) },
};
static MP_DEFINE_CONST_DICT(py_ml_globals_dict, py_ml_globals_dict_table);

View File

@ -1,142 +0,0 @@
/*
* This file is part of the OpenMV project.
*
* Copyright (c) 2013-2024 Ibrahim Abdelkader <iabdalkader@openmv.io>
* Copyright (c) 2013-2024 Kwabena W. Agyeman <kwagyeman@openmv.io>
*
* This work is licensed under the MIT license, see the file LICENSE for details.
*
* Python NMS class.
*/
#include "imlib_config.h"
#ifdef IMLIB_ENABLE_TFLM
#include "py/runtime.h"
#include "py_helper.h"
// TF NMS Object.
typedef struct py_ml_nms_obj {
mp_obj_base_t base;
int window_w;
int window_h;
rectangle_t roi;
list_t bounding_boxes;
} py_ml_nms_obj_t;
const mp_obj_type_t py_ml_nms_type;
// The use of mp_arg_parse_all() is deliberately avoided here to ensure this method remains fast.
static mp_obj_t py_ml_nms_add_bounding_box(uint n_args, const mp_obj_t *pos_args) {
enum { ARG_self, ARG_xmin, ARG_ymin, ARG_xmax, ARG_ymax, ARG_score, ARG_label_index };
py_ml_nms_obj_t *self_in = MP_OBJ_TO_PTR(pos_args[ARG_self]);
bounding_box_lnk_data_t lnk_data;
lnk_data.score = mp_obj_get_float(pos_args[ARG_score]);
if ((lnk_data.score >= 0.0f) && (lnk_data.score <= 1.0f)) {
float xmin = IM_CLAMP(mp_obj_get_float(pos_args[ARG_xmin]), 0.0f, ((float) self_in->window_w));
float ymin = IM_CLAMP(mp_obj_get_float(pos_args[ARG_ymin]), 0.0f, ((float) self_in->window_h));
float xmax = IM_CLAMP(mp_obj_get_float(pos_args[ARG_xmax]), 0.0f, ((float) self_in->window_w));
float ymax = IM_CLAMP(mp_obj_get_float(pos_args[ARG_ymax]), 0.0f, ((float) self_in->window_h));
lnk_data.rect.w = fast_floorf(xmax - xmin);
lnk_data.rect.h = fast_floorf(ymax - ymin);
if ((lnk_data.rect.w > 0) && (lnk_data.rect.h > 0)) {
lnk_data.rect.x = fast_floorf(xmin);
lnk_data.rect.y = fast_floorf(ymin);
lnk_data.label_index = mp_obj_get_int(pos_args[ARG_label_index]);
rectangle_nms_add_bounding_box(&self_in->bounding_boxes, &lnk_data);
}
}
return mp_const_none;
}
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(py_ml_nms_add_bounding_box_obj, 7, 7, py_ml_nms_add_bounding_box);
static mp_obj_t py_ml_nms_get_bounding_boxes(uint n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
enum { ARG_threshold, ARG_sigma };
static const mp_arg_t allowed_args[] = {
{ MP_QSTR_threshold, MP_ARG_OBJ | MP_ARG_KW_ONLY, {.u_rom_obj = MP_ROM_NONE } },
{ MP_QSTR_sigma, MP_ARG_OBJ | MP_ARG_KW_ONLY, {.u_rom_obj = MP_ROM_NONE } },
};
mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
mp_arg_parse_all(n_args - 1, pos_args + 1, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
py_ml_nms_obj_t *self_in = MP_OBJ_TO_PTR(pos_args[0]);
float threshold = py_helper_arg_to_float(args[ARG_threshold].u_obj, 0.1f);
float sigma = py_helper_arg_to_float(args[ARG_sigma].u_obj, 0.1f);
int max_label = rectangle_nms_get_bounding_boxes(&self_in->bounding_boxes, threshold, sigma);
rectangle_map_bounding_boxes(&self_in->bounding_boxes, self_in->window_w, self_in->window_h, &self_in->roi);
// Create a list per class label.
mp_obj_list_t *list = MP_OBJ_TO_PTR(mp_obj_new_list(max_label + 1, NULL));
for (size_t i = 0; i <= max_label; i++) {
list->items[i] = mp_obj_new_list(0, NULL);
}
list_for_each(it, (&self_in->bounding_boxes)) {
bounding_box_lnk_data_t *lnk_data = (bounding_box_lnk_data_t *) it->data;
mp_obj_t rect = mp_obj_new_tuple(4, (mp_obj_t []) {mp_obj_new_int(lnk_data->rect.x),
mp_obj_new_int(lnk_data->rect.y),
mp_obj_new_int(lnk_data->rect.w),
mp_obj_new_int(lnk_data->rect.h)});
mp_obj_t o = mp_obj_new_tuple(2, (mp_obj_t []) {rect, mp_obj_new_float(lnk_data->score)});
mp_obj_list_append(list->items[lnk_data->label_index], o);
}
return list;
}
static MP_DEFINE_CONST_FUN_OBJ_KW(py_ml_nms_get_bounding_boxes_obj, 1, py_ml_nms_get_bounding_boxes);
mp_obj_t py_ml_nms_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *all_args) {
enum { ARG_window_w, ARG_window_h, ARG_roi };
static const mp_arg_t allowed_args[] = {
{ MP_QSTR_window_w, MP_ARG_INT | MP_ARG_REQUIRED, {.u_int = 0 } },
{ MP_QSTR_window_h, MP_ARG_INT | MP_ARG_REQUIRED, {.u_int = 0 } },
{ MP_QSTR_roi, MP_ARG_OBJ | MP_ARG_REQUIRED, {.u_rom_obj = MP_ROM_NONE} },
};
mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
mp_arg_parse_all_kw_array(n_args, n_kw, all_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
// Extract the ROI manually as we do not have an image to validate against.
mp_obj_t *roi_obj;
mp_obj_get_array_fixed_n(args[ARG_roi].u_obj, 4, &roi_obj);
rectangle_t roi = {
.x = mp_obj_get_int(roi_obj[0]),
.y = mp_obj_get_int(roi_obj[1]),
.w = mp_obj_get_int(roi_obj[2]),
.h = mp_obj_get_int(roi_obj[3])
};
if ((roi.w < 1) || (roi.h < 1)) {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Invalid ROI dimensions!"));
}
py_ml_nms_obj_t *model = m_new_obj(py_ml_nms_obj_t);
model->base.type = &py_ml_nms_type;
model->window_w = args[ARG_window_w].u_int;
model->window_h = args[ARG_window_h].u_int;
model->roi = roi;
list_init(&model->bounding_boxes, sizeof(bounding_box_lnk_data_t));
return MP_OBJ_FROM_PTR(model);
}
static const mp_rom_map_elem_t py_ml_nms_locals_table[] = {
{ MP_ROM_QSTR(MP_QSTR_add_bounding_box), MP_ROM_PTR(&py_ml_nms_add_bounding_box_obj) },
{ MP_ROM_QSTR(MP_QSTR_get_bounding_boxes), MP_ROM_PTR(&py_ml_nms_get_bounding_boxes_obj) },
};
static MP_DEFINE_CONST_DICT(py_ml_nms_locals_dict, py_ml_nms_locals_table);
MP_DEFINE_CONST_OBJ_TYPE(
py_ml_nms_type,
MP_QSTR_tf_nms,
MP_TYPE_FLAG_NONE,
make_new, py_ml_nms_make_new,
locals_dict, &py_ml_nms_locals_dict
);
#endif // IMLIB_ENABLE_TFLM