diff --git a/src/omv/imlib/imlib.h b/src/omv/imlib/imlib.h index 3bb736729..903ef322c 100644 --- a/src/omv/imlib/imlib.h +++ b/src/omv/imlib/imlib.h @@ -145,6 +145,12 @@ typedef struct rectangle { int16_t h; } rectangle_t; +typedef struct bounding_box_lnk_data { + rectangle_t rect; + float score; + int label_index; +} bounding_box_lnk_data_t; + void rectangle_init(rectangle_t *ptr, int x, int y, int w, int h); void rectangle_copy(rectangle_t *dst, rectangle_t *src); bool rectangle_equal_fast(rectangle_t *ptr0, rectangle_t *ptr1); @@ -152,6 +158,9 @@ bool rectangle_overlap(rectangle_t *ptr0, rectangle_t *ptr1); void rectangle_intersected(rectangle_t *dst, rectangle_t *src); void rectangle_united(rectangle_t *dst, rectangle_t *src); float rectangle_iou(rectangle_t *r1, rectangle_t *r2); +void rectangle_nms_add_bounding_box(list_t *bounding_boxes, bounding_box_lnk_data_t *box); +int rectangle_nms_get_bounding_boxes(list_t *bounding_boxes, float threshold, float sigma); +void rectangle_map_bounding_boxes(list_t *bounding_boxes, int window_w, int window_h, rectangle_t *roi); ///////////////// // Color Stuff // diff --git a/src/omv/imlib/rectangle.c b/src/omv/imlib/rectangle.c index 9f3dfcec4..611995c50 100644 --- a/src/omv/imlib/rectangle.c +++ b/src/omv/imlib/rectangle.c @@ -128,3 +128,86 @@ float rectangle_iou(rectangle_t *r1, rectangle_t *r2) { int rect_union = (r1->w * r1->h) + (r2->w * r2->h) - rect_intersection; return ((float) rect_intersection) / ((float) rect_union); } + +// Adds a bounding box to the list of bounding boxes in descending order of score. +void rectangle_nms_add_bounding_box(list_t *bounding_boxes, bounding_box_lnk_data_t *box) { + // Insertion sort bounding boxes by score. + list_lnk_t *it = bounding_boxes->head; + for (; it; it = it->next) { + if (box->score > ((bounding_box_lnk_data_t *) it->data)->score) { + list_insert(bounding_boxes, it, box); + break; + } + } + + if (!it) { + list_push_back(bounding_boxes, box); + } +} + +// Soft non-max supress the list of bounding boxes. Returns the maximum label index of the new list. +int rectangle_nms_get_bounding_boxes(list_t *bounding_boxes, float threshold, float sigma) { + // Soft non-max suppression with a Gaussian is used below, as this provides the best results. + // A Gaussian is used to apply a soft score penalty to overlapping boxes. On loop entry, + // "bounding_boxes" is sorted, but after each iteration, the next highest score must be picked + // again, given that the score penalty changes the order. + float sigma_scale = (sigma > 0.0f) ? (-1.0f / sigma) : 0.0f; + + list_t nms_bounding_boxes; + list_init(&nms_bounding_boxes, sizeof(bounding_box_lnk_data_t)); + + int max_label_index = 0; + + // The first detection has the higest score since the list is sorted. + list_lnk_t *max_it = bounding_boxes->head; + while (list_size(bounding_boxes)) { + bounding_box_lnk_data_t lnk_data; + memcpy(&lnk_data, max_it->data, bounding_boxes->data_len); + list_move_back(&nms_bounding_boxes, bounding_boxes, max_it); + + float max_score = 0.0f; + for (list_lnk_t *it = bounding_boxes->head; it; ) { + bounding_box_lnk_data_t *lnk_data2 = list_get_data(it); + + // Advance to next now as "it" will be invalid if we remove the current item. + list_lnk_t *old_it = it; + it = it->next; + + float iou = rectangle_iou(&lnk_data.rect, &lnk_data2->rect); + // Do not use fast_expf() as it does not output 1 when it's input is 0. + // This will cause the scores of non-overlapping bounding boxes to decay. + lnk_data2->score *= expf(sigma_scale * iou * iou); + + if (lnk_data2->score < threshold) { + list_remove(bounding_boxes, old_it, NULL); + } else if (lnk_data2->score > max_score) { + max_score = lnk_data2->score; + max_it = old_it; + } + } + + // Find the maximum label index for the output list. + max_label_index = IM_MAX(lnk_data.label_index, max_label_index); + } + + // Set the original list pointers to equal the new list. + memcpy(bounding_boxes, &nms_bounding_boxes, sizeof(list_t)); + return max_label_index; +} + +void rectangle_map_bounding_boxes(list_t *bounding_boxes, int window_w, int window_h, rectangle_t *roi) { + float x_scale = roi->w / ((float) window_w); + float y_scale = roi->h / ((float) window_h); + // MAX == KeepAspectRatioByExpanding - MIN == KeepAspectRatio + float scale = IM_MIN(x_scale, y_scale); + int x_offset = fast_floorf((roi->w - (window_w * scale)) / 2.0f) + roi->x; + int y_offset = fast_floorf((roi->h - (window_h * scale)) / 2.0f) + roi->y; + + list_for_each(it, bounding_boxes) { + rectangle_t *rect = &((bounding_box_lnk_data_t *) it->data)->rect; + rect->x = fast_floorf((rect->x * scale) + x_offset); + rect->y = fast_floorf((rect->y * scale) + y_offset); + rect->w = fast_floorf(rect->w * scale); + rect->h = fast_floorf(rect->h * scale); + } +} diff --git a/src/omv/modules/py_tf.c b/src/omv/modules/py_tf.c index 2b8504a1d..2171154c5 100644 --- a/src/omv/modules/py_tf.c +++ b/src/omv/modules/py_tf.c @@ -157,14 +157,11 @@ MP_DEFINE_CONST_OBJ_TYPE( // TF Model Output Object. typedef struct py_tf_model_output_obj { mp_obj_base_t base; - rectangle_t *roi; void *model_output; libtf_parameters_t *params; // Pre-compute for lookup speed. size_t output_size; mp_obj_t rect; - // Convenience stuff. - list_t bounding_boxes; } py_tf_model_output_obj_t; STATIC void py_tf_model_output_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) { @@ -234,164 +231,12 @@ STATIC mp_obj_t py_tf_model_output_subscr(mp_obj_t self_in, mp_obj_t index, mp_o return MP_OBJ_NULL; // op not supported } -typedef struct py_tf_model_output_bounding_box_lnk_data { - rectangle_t rect; - float score; - int label_index; -} py_tf_model_output_bounding_box_lnk_data_t; - -// This convenience function is designed to collect bounding boxes, allowing for a non-maximal -// suppression of them later. It is not necessary to use this function to parse model output. -// The use of mp_arg_parse_all() is deliberately avoided here to ensure this method remains fast. -STATIC mp_obj_t py_tf_model_output_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_tf_model_output_obj_t *self_in = (py_tf_model_output_obj_t *) pos_args[ARG_self]; - - py_tf_model_output_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->params->input_width)); - float ymin = IM_CLAMP(mp_obj_get_float(pos_args[ARG_ymin]), 0.0f, (float) (self_in->params->input_height)); - float xmax = IM_CLAMP(mp_obj_get_float(pos_args[ARG_xmax]), 0.0f, (float) (self_in->params->input_width)); - float ymax = IM_CLAMP(mp_obj_get_float(pos_args[ARG_ymax]), 0.0f, (float) (self_in->params->input_height)); - - 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]); - - // Insertion sort bounding boxes by score. - list_lnk_t *it = self_in->bounding_boxes.head; - for (; it; it = it->next) { - if (lnk_data.score > ((py_tf_model_output_bounding_box_lnk_data_t *) it->data)->score) { - list_insert(&self_in->bounding_boxes, it, &lnk_data); - break; - } - } - - if (!it) { - list_push_back(&self_in->bounding_boxes, &lnk_data); - } - } - } - - return mp_const_none; -} -STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(py_tf_model_output_add_bounding_box_obj, 7, 7, - py_tf_model_output_add_bounding_box); - -// This function performs non-maximal suppression on bounding boxes collected using the add_bounding_box -// function and returns the final list of bounding boxes, mapped to the image ROI and separated by label -// index into their own lists. Note that using this function to parse the model output is optional. -STATIC mp_obj_t py_tf_model_output_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_tf_model_output_obj_t *self_in = (py_tf_model_output_obj_t *) 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); - - // Soft non-max suppression with a Gaussian is used below, as this provides the best results. - // A Gaussian is used to apply a soft score penalty to overlapping boxes. On loop entry, - // "bounding_boxes" is sorted, but after each iteration, the next highest score must be picked - // again, given that the score penalty changes the order. - - float sigma_scale = (sigma > 0.0f) ? (-1.0f / sigma) : 0.0f; - - list_t nms_bounding_boxes; - list_init(&nms_bounding_boxes, sizeof(py_tf_model_output_bounding_box_lnk_data_t)); - - int max_label = 0; - - // The first detection has the higest score since the list is sorted. - list_lnk_t *max_it = self_in->bounding_boxes.head; - while (list_size(&self_in->bounding_boxes)) { - py_tf_model_output_bounding_box_lnk_data_t lnk_data; - list_remove(&self_in->bounding_boxes, max_it, &lnk_data); - list_push_back(&nms_bounding_boxes, &lnk_data); - - float max_score = 0.0f; - for (list_lnk_t *it = self_in->bounding_boxes.head; it; ) { - py_tf_model_output_bounding_box_lnk_data_t *lnk_data2 = list_get_data(it); - - // Advance to next now as "it" will be invalid if we remove the current item. - list_lnk_t *old_it = it; - it = it->next; - - float iou = rectangle_iou(&lnk_data.rect, &lnk_data2->rect); - // Do not use fast_expf() as it does not output 1 when it's input is 0. - // This will cause the scores of non-overlapping bounding boxes to decay. - lnk_data2->score *= expf(sigma_scale * iou * iou); - - if (lnk_data2->score < threshold) { - list_remove(&self_in->bounding_boxes, old_it, NULL); - } else if (lnk_data2->score > max_score) { - max_score = lnk_data2->score; - max_it = old_it; - } - } - - // Find the maximum label index for the output list. - max_label = IM_MAX(lnk_data.label_index, max_label); - } - - memcpy(&self_in->bounding_boxes, &nms_bounding_boxes, sizeof(list_t)); - - // Create a list per class label. - mp_obj_t list = mp_obj_new_list(max_label + 1, NULL); - for (size_t i = 0; i <= max_label; i++) { - ((mp_obj_list_t *) list)->items[i] = mp_obj_new_list(0, NULL); - } - - float x_scale = self_in->roi->w / ((float) self_in->params->input_width); - float y_scale = self_in->roi->h / ((float) self_in->params->input_height); - // MAX == KeepAspectRatioByExpanding - MIN == KeepAspectRatio - float scale = IM_MIN(x_scale, y_scale); - int x_offset = fast_floorf((self_in->roi->w - (self_in->params->input_width * scale)) / 2.0f) + self_in->roi->x; - int y_offset = fast_floorf((self_in->roi->h - (self_in->params->input_height * scale)) / 2.0f) + self_in->roi->y; - - size_t len = list_size(&nms_bounding_boxes); - for (size_t i = 0; i < len; i++) { - py_tf_model_output_bounding_box_lnk_data_t lnk_data; - list_pop_front(&nms_bounding_boxes, &lnk_data); - py_tf_classification_obj_t *o = m_new_obj(py_tf_classification_obj_t); - o->base.type = &py_tf_classification_type; - o->x = mp_obj_new_int(fast_floorf(lnk_data.rect.x * scale) + x_offset); - o->y = mp_obj_new_int(fast_floorf(lnk_data.rect.y * scale) + y_offset); - o->w = mp_obj_new_int(fast_floorf(lnk_data.rect.w * scale)); - o->h = mp_obj_new_int(fast_floorf(lnk_data.rect.h * scale)); - o->output = mp_obj_new_float(lnk_data.score); - mp_obj_list_append(((mp_obj_list_t *) list)->items[lnk_data.label_index], o); - } - - return list; -} -STATIC MP_DEFINE_CONST_FUN_OBJ_KW(py_tf_model_output_get_bounding_boxes_obj, 1, py_tf_model_output_get_bounding_boxes); - -STATIC const mp_rom_map_elem_t py_tf_model_output_locals_table[] = { - { MP_ROM_QSTR(MP_QSTR_add_bounding_box), MP_ROM_PTR(&py_tf_model_output_add_bounding_box_obj) }, - { MP_ROM_QSTR(MP_QSTR_get_bounding_boxes), MP_ROM_PTR(&py_tf_model_output_get_bounding_boxes_obj) }, -}; - -STATIC MP_DEFINE_CONST_DICT(py_tf_model_output_locals_dict, py_tf_model_output_locals_table); - STATIC MP_DEFINE_CONST_OBJ_TYPE( py_tf_model_output_type, MP_QSTR_tf_model_output, MP_TYPE_FLAG_NONE, attr, py_tf_model_output_attr, - subscr, py_tf_model_output_subscr, - locals_dict, &py_tf_model_output_locals_dict + subscr, py_tf_model_output_subscr ); // TF Input/Output callback functions. @@ -680,7 +525,6 @@ STATIC void py_tf_predict_output_callback(void *callback_data, py_tf_predict_callback_data_t *arg = (py_tf_predict_callback_data_t *) callback_data; py_tf_model_output_obj_t *o = m_new_obj(py_tf_model_output_obj_t); o->base.type = &py_tf_model_output_type; - o->roi = arg->roi; o->model_output = model_output; o->params = params; o->output_size = params->output_height * params->output_width * params->output_channels; @@ -688,7 +532,6 @@ STATIC void py_tf_predict_output_callback(void *callback_data, mp_obj_new_int(arg->roi->y), mp_obj_new_int(arg->roi->w), mp_obj_new_int(arg->roi->h)}); - list_init(&o->bounding_boxes, sizeof(py_tf_model_output_bounding_box_lnk_data_t)); *(arg->out) = mp_call_function_2(arg->callback, arg->model, o); } @@ -1119,6 +962,8 @@ STATIC MP_DEFINE_CONST_OBJ_TYPE( locals_dict, &py_tf_model_locals_dict ); +extern const mp_obj_type_t py_tf_nms_type; + STATIC const mp_rom_map_elem_t py_tf_globals_dict_table[] = { { MP_ROM_QSTR(MP_QSTR___name__), MP_OBJ_NEW_QSTR(MP_QSTR_tf) }, { MP_ROM_QSTR(MP_QSTR_SCALE_NONE), MP_ROM_INT(PY_TF_SCALE_NONE) }, @@ -1126,6 +971,7 @@ STATIC const mp_rom_map_elem_t py_tf_globals_dict_table[] = { { MP_ROM_QSTR(MP_QSTR_SCALE_S1_1), MP_ROM_INT(PY_TF_SCALE_S1_1) }, { MP_ROM_QSTR(MP_QSTR_SCALE_S128_127), MP_ROM_INT(PY_TF_SCALE_S128_127) }, { MP_ROM_QSTR(MP_QSTR_Model), MP_ROM_PTR(&py_tf_model_type) }, + { MP_ROM_QSTR(MP_QSTR_NMS), MP_ROM_PTR(&py_tf_nms_type) }, { MP_ROM_QSTR(MP_QSTR_load), MP_ROM_PTR(&py_tf_model_type) }, { MP_ROM_QSTR(MP_QSTR_load_builtin_model), MP_ROM_PTR(&py_tf_model_type) }, }; diff --git a/src/omv/modules/py_tf_nms.c b/src/omv/modules/py_tf_nms.c new file mode 100644 index 000000000..8a46efaa6 --- /dev/null +++ b/src/omv/modules/py_tf_nms.c @@ -0,0 +1,143 @@ +/* + * This file is part of the OpenMV project. + * + * Copyright (c) 2013-2024 Ibrahim Abdelkader + * Copyright (c) 2013-2024 Kwabena W. Agyeman + * + * This work is licensed under the MIT license, see the file LICENSE for details. + * + * Python NMS class. + */ +#include "imlib_config.h" + +#ifdef IMLIB_ENABLE_TF +#include "py/runtime.h" +#include "py_helper.h" + +// TF NMS Object. +typedef struct py_tf_nms_obj { + mp_obj_base_t base; + int window_w; + int window_h; + rectangle_t roi; + list_t bounding_boxes; +} py_tf_nms_obj_t; + +const mp_obj_type_t py_tf_nms_type; + +// The use of mp_arg_parse_all() is deliberately avoided here to ensure this method remains fast. +STATIC mp_obj_t py_tf_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_tf_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_tf_nms_add_bounding_box_obj, 7, 7, py_tf_nms_add_bounding_box); + +STATIC mp_obj_t py_tf_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_tf_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_tf_nms_get_bounding_boxes_obj, 1, py_tf_nms_get_bounding_boxes); + +mp_obj_t py_tf_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_tf_nms_obj_t *model = m_new_obj(py_tf_nms_obj_t); + model->base.type = &py_tf_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_tf_nms_locals_table[] = { + { MP_ROM_QSTR(MP_QSTR_add_bounding_box), MP_ROM_PTR(&py_tf_nms_add_bounding_box_obj) }, + { MP_ROM_QSTR(MP_QSTR_get_bounding_boxes), MP_ROM_PTR(&py_tf_nms_get_bounding_boxes_obj) }, +}; + +STATIC MP_DEFINE_CONST_DICT(py_tf_nms_locals_dict, py_tf_nms_locals_table); + +MP_DEFINE_CONST_OBJ_TYPE( + py_tf_nms_type, + MP_QSTR_tf_nms, + MP_TYPE_FLAG_NONE, + make_new, py_tf_nms_make_new, + locals_dict, &py_tf_nms_locals_dict + ); + +#endif // IMLIB_ENABLE_TF