modules/py_tf: Remove detect() and segment() in favor of predict().

This commit is contained in:
Kwabena W. Agyeman 2024-06-20 14:29:30 -07:00 committed by iabdalkader
parent 3863c38228
commit 3e37f46db4
2 changed files with 100 additions and 218 deletions

View File

@ -14,10 +14,10 @@ import math
sensor.reset() # Reset and initialize the sensor.
sensor.set_pixformat(sensor.RGB565) # Set pixel format to RGB565 (or GRAYSCALE)
sensor.set_framesize(sensor.QVGA) # Set frame size to QVGA (320x240)
sensor.set_windowing((240, 240)) # Set 240x240 window.
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
labels, net = tf.Model("fomo_face_detection")
@ -36,18 +36,38 @@ colors = [ # Add more colors if you are detecting more than 7 types of classes
(255, 255, 255),
]
# 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 callback then returns a
# list per class which each contain a list of (rect, score) tuples representing the detected
# objects.
def fomo_callback(model, rect):
out = model.output[0]
oh, ow, oc = model.output_shape
nms = tf.NMS(ow, oh, rect)
for i in range(oc):
img = out.get_image(i)
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()
while True:
clock.tick()
img = sensor.snapshot()
# detect() returns all objects found in the image (splitted out per class already)
# we skip class index 0, as that is the background, and then draw circles of the center
# of our objects
for i, detection_list in enumerate(
net.detect(img, thresholds=[(math.ceil(min_confidence * 255), 255)])
net.predict(img, callback=fomo_callback)
):
if i == 0:
continue # background class

View File

@ -120,11 +120,84 @@ 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
}
STATIC mp_obj_t py_tf_model_output_get_image(uint n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
enum { ARG_channel, ARG_roi, ARG_scale };
static const mp_arg_t allowed_args[] = {
{ MP_QSTR_channel, MP_ARG_INT | MP_ARG_REQUIRED, {.u_int = 0} },
{ MP_QSTR_roi, MP_ARG_OBJ | MP_ARG_KW_ONLY, {.u_rom_obj = MP_ROM_NONE} },
{ MP_QSTR_scale, MP_ARG_INT | MP_ARG_KW_ONLY, {.u_int = PY_TF_SCALE_0_1} },
};
// Parse args.
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 = MP_OBJ_TO_PTR(pos_args[0]);
image_t temp = {.w = self->params->output_width, .h = self->params->output_height};
rectangle_t roi = py_helper_arg_to_roi(args[ARG_roi].u_obj, &temp);
image_t img = {
.w = roi.w,
.h = roi.h,
.pixfmt = PIXFORMAT_GRAYSCALE,
.pixels = xalloc(roi.w * roi.h)
};
int channel = args[ARG_channel].u_int;
int shift = (self->params->output_datatype == LIBTF_DATATYPE_INT8) ? PY_TF_GRAYSCALE_MID : 0;
float fscale = 1.0f, fadd = 0.0f;
switch (args[ARG_scale].u_int) {
case PY_TF_SCALE_0_1: // convert 0->1 to 0->255
fscale = 255.0f;
break;
case PY_TF_SCALE_S1_1: // convert -1->1 to 0->255
fscale = 127.5f;
fadd = 127.5f;
break;
case PY_TF_SCALE_S128_127: // convert -128->127 to 0->255
fadd = 128.0f;
break;
case PY_TF_SCALE_NONE: // convert 0->255 to 0->255
default:
break;
}
for (int y = 0; y < roi.h; y++) {
int row_index = (y + roi.y) * self->params->output_width * self->params->output_channels;
uint8_t *row_ptr = IMAGE_COMPUTE_GRAYSCALE_PIXEL_ROW_PTR(&img, y);
for (int x = 0; x < roi.w; x++) {
int index = row_index + ((x + roi.x) * self->params->output_channels) + channel;
if (self->params->output_datatype == LIBTF_DATATYPE_FLOAT) {
float mo = (((float *) self->model_output)[index] * fscale) + fadd;
IMAGE_PUT_GRAYSCALE_PIXEL_FAST(row_ptr, x, fast_floorf(mo));
} else {
uint8_t mo = ((uint8_t *) self->model_output)[index] ^ shift;
IMAGE_PUT_GRAYSCALE_PIXEL_FAST(row_ptr, x, mo);
}
}
}
return py_image_from_struct(&img);
}
STATIC MP_DEFINE_CONST_FUN_OBJ_KW(py_tf_model_output_get_image_obj, 1, py_tf_model_output_get_image);
STATIC const mp_rom_map_elem_t py_tf_model_output_locals_dict_table[] = {
{ MP_ROM_QSTR(MP_QSTR_get_image), MP_ROM_PTR(&py_tf_model_output_get_image_obj) },
};
STATIC MP_DEFINE_CONST_DICT(py_tf_model_output_locals_dict, py_tf_model_output_locals_dict_table);
STATIC MP_DEFINE_CONST_OBJ_TYPE(
py_tf_model_output_type,
MP_QSTR_tf_model_output,
MP_TYPE_FLAG_NONE,
subscr, py_tf_model_output_subscr
subscr, py_tf_model_output_subscr,
locals_dict, &py_tf_model_output_locals_dict
);
// TF Input/Output callback functions.
@ -361,45 +434,6 @@ STATIC void py_tf_regression_input_callback(void *callback_data,
}
}
STATIC void py_tf_segment_output_callback(void *callback_data,
void *model_output,
libtf_parameters_t *params) {
mp_obj_t *arg = (mp_obj_t *) callback_data;
int shift = (params->output_datatype == LIBTF_DATATYPE_INT8) ? PY_TF_GRAYSCALE_MID : 0;
*arg = mp_obj_new_list(params->output_channels, NULL);
for (int i = 0, ii = params->output_channels; i < ii; i++) {
image_t img = {
.w = params->output_width,
.h = params->output_height,
.pixfmt = PIXFORMAT_GRAYSCALE,
.pixels = xalloc(params->output_width * params->output_height * sizeof(uint8_t))
};
((mp_obj_list_t *) *arg)->items[i] = py_image_from_struct(&img);
for (int y = 0, yy = params->output_height, xx = params->output_width; y < yy; y++) {
int row = y * xx * ii;
uint8_t *row_ptr = IMAGE_COMPUTE_GRAYSCALE_PIXEL_ROW_PTR(&img, y);
for (int x = 0; x < xx; x++) {
int col = x * ii;
if (params->output_datatype == LIBTF_DATATYPE_FLOAT) {
IMAGE_PUT_GRAYSCALE_PIXEL_FAST(row_ptr, x,
((float *) model_output)[row + col + i] * PY_TF_GRAYSCALE_RANGE);
} else {
IMAGE_PUT_GRAYSCALE_PIXEL_FAST(row_ptr, x,
((uint8_t *) model_output)[row + col + i] ^ shift);
}
}
}
}
}
typedef struct py_tf_predict_callback_data {
mp_obj_t model;
rectangle_t *roi;
@ -452,176 +486,6 @@ STATIC void py_tf_model_print(const mp_print_t *print, mp_obj_t self_in, mp_prin
(double) self->params.output_scale, self->params.output_zero_point);
}
STATIC mp_obj_t py_tf_model_segment(uint n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
enum { ARG_roi, ARG_scale, ARG_mean, ARG_stdev };
static const mp_arg_t allowed_args[] = {
{ MP_QSTR_roi, MP_ARG_OBJ | MP_ARG_KW_ONLY, {.u_rom_obj = MP_ROM_NONE} },
{ MP_QSTR_scale, MP_ARG_INT | MP_ARG_KW_ONLY, {.u_int = PY_TF_SCALE_0_1} },
{ MP_QSTR_mean, MP_ARG_OBJ | MP_ARG_KW_ONLY, {.u_rom_obj = MP_ROM_NONE} },
{ MP_QSTR_stdev, MP_ARG_OBJ | MP_ARG_KW_ONLY, {.u_rom_obj = MP_ROM_NONE} },
};
// Parse args.
mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
mp_arg_parse_all(n_args - 2, pos_args + 2, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
image_t *image = py_helper_arg_to_image(pos_args[1], ARG_IMAGE_ANY);
rectangle_t roi = py_helper_arg_to_roi(args[ARG_roi].u_obj, image);
fb_alloc_mark();
py_tf_alloc_log_buffer();
py_tf_model_obj_t *model = MP_OBJ_TO_PTR(pos_args[0]);
uint8_t *tensor_arena = fb_alloc(model->params.tensor_arena_size, FB_ALLOC_PREFER_SPEED | FB_ALLOC_CACHE_ALIGN);
py_tf_input_callback_data_t py_tf_input_callback_data = {
.img = image,
.roi = &roi,
.scale = args[ARG_scale].u_int,
.mean = {0.0f, 0.0f, 0.0f},
.stdev = {1.0f, 1.0f, 1.0f}
};
py_helper_arg_to_float_array(args[ARG_mean].u_obj, py_tf_input_callback_data.mean, 3);
py_helper_arg_to_float_array(args[ARG_stdev].u_obj, py_tf_input_callback_data.stdev, 3);
mp_obj_t py_tf_segment_output_callback_data;
if (libtf_invoke(model->data,
tensor_arena,
&model->params,
py_tf_input_callback,
&py_tf_input_callback_data,
py_tf_segment_output_callback,
&py_tf_segment_output_callback_data) != 0) {
// Note can't use MP_ERROR_TEXT here.
mp_raise_msg(&mp_type_OSError, (mp_rom_error_text_t) py_tf_log_buffer);
}
fb_alloc_free_till_mark();
return py_tf_segment_output_callback_data;
}
STATIC MP_DEFINE_CONST_FUN_OBJ_KW(py_tf_model_segment_obj, 2, py_tf_model_segment);
STATIC mp_obj_t py_tf_model_detect(uint n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
enum { ARG_roi, ARG_thresholds, ARG_invert, ARG_scale, ARG_mean, ARG_stdev };
static const mp_arg_t allowed_args[] = {
{ MP_QSTR_roi, MP_ARG_OBJ | MP_ARG_KW_ONLY, {.u_rom_obj = MP_ROM_NONE} },
{ MP_QSTR_thresholds, MP_ARG_OBJ | MP_ARG_KW_ONLY, {.u_rom_obj = MP_ROM_NONE} },
{ MP_QSTR_invert, MP_ARG_INT | MP_ARG_KW_ONLY, {.u_bool = false } },
{ MP_QSTR_scale, MP_ARG_INT | MP_ARG_KW_ONLY, {.u_int = PY_TF_SCALE_0_1} },
{ MP_QSTR_mean, MP_ARG_OBJ | MP_ARG_KW_ONLY, {.u_rom_obj = MP_ROM_NONE} },
{ MP_QSTR_stdev, MP_ARG_OBJ | MP_ARG_KW_ONLY, {.u_rom_obj = MP_ROM_NONE} },
};
// Parse args.
mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
mp_arg_parse_all(n_args - 2, pos_args + 2, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
image_t *image = py_helper_arg_to_image(pos_args[1], ARG_IMAGE_ANY);
rectangle_t roi = py_helper_arg_to_roi(args[ARG_roi].u_obj, image);
bool invert = args[ARG_invert].u_int;
fb_alloc_mark();
py_tf_alloc_log_buffer();
py_tf_model_obj_t *model = MP_OBJ_TO_PTR(pos_args[0]);
uint8_t *tensor_arena = fb_alloc(model->params.tensor_arena_size, FB_ALLOC_PREFER_SPEED | FB_ALLOC_CACHE_ALIGN);
py_tf_input_callback_data_t py_tf_input_callback_data = {
.img = image,
.roi = &roi,
.scale = args[ARG_scale].u_int,
.mean = {0.0f, 0.0f, 0.0f},
.stdev = {1.0f, 1.0f, 1.0f}
};
py_helper_arg_to_float_array(args[ARG_mean].u_obj, py_tf_input_callback_data.mean, 3);
py_helper_arg_to_float_array(args[ARG_stdev].u_obj, py_tf_input_callback_data.stdev, 3);
mp_obj_t py_tf_segment_output_callback_data;
if (libtf_invoke(model->data,
tensor_arena,
&model->params,
py_tf_input_callback,
&py_tf_input_callback_data,
py_tf_segment_output_callback,
&py_tf_segment_output_callback_data) != 0) {
// Note can't use MP_ERROR_TEXT here.
mp_raise_msg(&mp_type_OSError, (mp_rom_error_text_t) py_tf_log_buffer);
}
list_t thresholds;
list_init(&thresholds, sizeof(color_thresholds_list_lnk_data_t));
py_helper_arg_to_thresholds(args[ARG_thresholds].u_obj, &thresholds);
if (!list_size(&thresholds)) {
color_thresholds_list_lnk_data_t lnk_data;
lnk_data.LMin = PY_TF_GRAYSCALE_MID;
lnk_data.LMax = PY_TF_GRAYSCALE_RANGE;
lnk_data.AMin = COLOR_A_MIN;
lnk_data.AMax = COLOR_A_MAX;
lnk_data.BMin = COLOR_B_MIN;
lnk_data.BMax = COLOR_B_MAX;
list_push_back(&thresholds, &lnk_data);
}
mp_obj_list_t *img_list = (mp_obj_list_t *) py_tf_segment_output_callback_data;
mp_obj_list_t *out_list = mp_obj_new_list(img_list->len, NULL);
float fscale = 1.f / PY_TF_GRAYSCALE_RANGE;
for (int i = 0, ii = img_list->len; i < ii; i++) {
image_t *img = py_image_cobj(img_list->items[i]);
float x_scale = roi.w / ((float) img->w);
float y_scale = roi.h / ((float) img->h);
// MAX == KeepAspectRatioByExpanding - MIN == KeepAspectRatio
float scale = IM_MIN(x_scale, y_scale);
int x_offset = fast_floorf((roi.w - (img->w * scale)) / 2.0f) + roi.x;
int y_offset = fast_floorf((roi.h - (img->h * scale)) / 2.0f) + roi.y;
list_t out;
imlib_find_blobs(&out, img, &((rectangle_t) {0, 0, img->w, img->h}), 1, 1,
&thresholds, invert, 1, 1, false, 0,
NULL, NULL, NULL, NULL, 0, 0);
mp_obj_list_t *objects_list = mp_obj_new_list(list_size(&out), NULL);
for (int j = 0, jj = list_size(&out); j < jj; j++) {
find_blobs_list_lnk_data_t lnk_data;
list_pop_front(&out, &lnk_data);
histogram_t hist;
hist.LBinCount = PY_TF_GRAYSCALE_RANGE + 1;
hist.ABinCount = 0;
hist.BBinCount = 0;
hist.LBins = fb_alloc(hist.LBinCount * sizeof(float), FB_ALLOC_NO_HINT);
hist.ABins = NULL;
hist.BBins = NULL;
imlib_get_histogram(&hist, img, &lnk_data.rect, &thresholds, invert, NULL);
statistics_t stats;
imlib_get_statistics(&stats, img->pixfmt, &hist);
fb_free(); // fb_alloc(hist.LBinCount * sizeof(float), FB_ALLOC_NO_HINT);
mp_obj_t rect = mp_obj_new_tuple(4, (mp_obj_t []) {
mp_obj_new_int(fast_floorf(lnk_data.rect.x * scale) + x_offset),
mp_obj_new_int(fast_floorf(lnk_data.rect.y * scale) + y_offset),
mp_obj_new_int(fast_floorf(lnk_data.rect.w * scale)),
mp_obj_new_int(fast_floorf(lnk_data.rect.h * scale))
});
objects_list->items[j] = mp_obj_new_tuple(2, (mp_obj_t []) {
rect, mp_obj_new_float(stats.LMean * fscale)
});
}
out_list->items[i] = objects_list;
}
fb_alloc_free_till_mark();
return out_list;
}
STATIC MP_DEFINE_CONST_FUN_OBJ_KW(py_tf_model_detect_obj, 2, py_tf_model_detect);
STATIC mp_obj_t py_tf_model_predict(uint n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
enum { ARG_roi, ARG_callback, ARG_scale, ARG_mean, ARG_stdev };
static const mp_arg_t allowed_args[] = {
@ -846,8 +710,6 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_1(py_tf_model_deinit_obj, py_tf_model_deinit);
STATIC const mp_rom_map_elem_t py_tf_model_locals_dict_table[] = {
{ MP_ROM_QSTR(MP_QSTR___del__), MP_ROM_PTR(&py_tf_model_deinit_obj) },
{ MP_ROM_QSTR(MP_QSTR_segment), MP_ROM_PTR(&py_tf_model_segment_obj) },
{ MP_ROM_QSTR(MP_QSTR_detect), MP_ROM_PTR(&py_tf_model_detect_obj) },
{ MP_ROM_QSTR(MP_QSTR_regression), MP_ROM_PTR(&py_tf_model_predict_obj) },
{ MP_ROM_QSTR(MP_QSTR_predict), MP_ROM_PTR(&py_tf_model_predict_obj) },
};