modules/py_ml: Update ML API to support multi-input models.

This commit is contained in:
iabdalkader 2024-07-07 16:46:05 +03:00
parent de0d46fa68
commit 70b89f4744
3 changed files with 79 additions and 255 deletions

View File

@ -56,15 +56,15 @@ static bool ml_backend_valid_dataype(TfLiteType type) {
type == kTfLiteFloat32);
}
static py_ml_dtype_t ml_backend_map_dtype(TfLiteType type) {
static char ml_backend_map_dtype(TfLiteType type) {
if (type == kTfLiteUInt8) {
return PY_ML_DTYPE_UINT8;
return 'B';
} else if (type == kTfLiteInt8) {
return PY_ML_DTYPE_INT8;
return 'b';
} else if (type == kTfLiteInt16) {
return PY_ML_DTYPE_INT16;
return 'h';
} else {
return PY_ML_DTYPE_FLOAT;
return 'f';
}
}

View File

@ -25,32 +25,6 @@
#include "tflm_builtin_models.h"
#include "ulab/code/ndarray.h"
#define PY_ML_GRAYSCALE_RANGE ((COLOR_GRAYSCALE_MAX) -(COLOR_GRAYSCALE_MIN))
#define PY_ML_GRAYSCALE_MID (((PY_ML_GRAYSCALE_RANGE) +1) / 2)
static const char *py_ml_map_dtype(py_ml_dtype_t dtype) {
if (dtype == PY_ML_DTYPE_UINT8) {
return "uint8";
} else if (dtype == PY_ML_DTYPE_INT8) {
return "int8";
} else if (dtype == PY_ML_DTYPE_INT16) {
return "int16";
} else {
return "float";
}
}
// TF Input/Output callback functions.
typedef mp_obj_t py_ml_output_data_t;
typedef struct _py_ml_input_callback_data {
void *data;
rectangle_t roi;
py_ml_scale_t scale;
float mean[3];
float stdev[3];
} py_ml_input_data_t;
static size_t py_ml_tuple_sum(mp_obj_tuple_t *o) {
if (o->len < 1) {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Unexpected tensor shape"));
@ -63,175 +37,67 @@ static size_t py_ml_tuple_sum(mp_obj_tuple_t *o) {
return size;
}
static void py_ml_tuple_hwc(mp_obj_tuple_t *o, size_t *h, size_t *w, size_t *c) {
if (o->len != 1 || ((mp_obj_tuple_t *) MP_OBJ_TO_PTR(o->items[0]))->len != 4) {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Unexpected tensor shape"));
}
o = MP_OBJ_TO_PTR(o->items[0]);
*h = mp_obj_get_int(o->items[1]);
*w = mp_obj_get_int(o->items[2]);
*c = mp_obj_get_int(o->items[3]);
}
// TF Input/Output callback functions.
static void py_ml_input_callback(py_ml_model_obj_t *model, void *arg) {
// TODO we assume that there's a single input.
void *model_input = ml_backend_get_input(model, 0);
py_ml_input_data_t *input_data = (py_ml_input_data_t *) arg;
mp_obj_list_t *input_list = MP_OBJ_TO_PTR(*((mp_obj_t *) arg));
// TODO we assume that the input shape is (1, h, w, c)
size_t input_height = 0, input_width = 0, input_channels = 0;
py_ml_tuple_hwc(model->input_shape, &input_height, &input_width, &input_channels);
for (size_t i = 0; i < model->inputs_size; i++) {
void *input_buffer = ml_backend_get_input(model, i);
size_t input_size = py_ml_tuple_sum(MP_OBJ_TO_PTR(model->input_shape->items[i]));
mp_obj_tuple_t *input_shape = MP_OBJ_TO_PTR(model->input_shape->items[i]);
mp_obj_t input_arg = input_list->items[i];
int shift = (model->input_dtype == PY_ML_DTYPE_INT8) ? PY_ML_GRAYSCALE_MID : 0;
float fscale = 1.0f, fadd = 0.0f;
if (mp_obj_is_callable(input_arg)) {
// Input is a callable. Call the object and pass the tensor buffer and dtype.
mp_obj_t fargs[3] = {
mp_obj_new_bytearray_by_ref(input_size, input_buffer),
MP_OBJ_FROM_PTR(input_shape),
mp_obj_new_int(model->input_dtype)
};
mp_call_function_n_kw(input_arg, 3, 0, fargs);
} else if (MP_OBJ_IS_TYPE(input_arg, &ulab_ndarray_type)) {
// Input is an ndarry. The input is converted and copied to the tensor buffer.
ndarray_obj_t *input_array = MP_OBJ_TO_PTR(input_arg);
switch (input_data->scale) {
case PY_ML_SCALE_0_1: // convert 0->255 to 0->1
fscale = 1.0f / 255.0f;
break;
case PY_ML_SCALE_S1_1: // convert 0->255 to -1->1
fscale = 2.0f / 255.0f;
fadd = -1.0f;
break;
case PY_ML_SCALE_S128_127: // convert 0->255 to -128->127
fadd = -128.0f;
break;
case PY_ML_SCALE_NONE: // convert 0->255 to 0->255
default:
break;
}
float fscale_r = fscale, fadd_r = fadd;
float fscale_g = fscale, fadd_g = fadd;
float fscale_b = fscale, fadd_b = fadd;
// To normalize the input image we need to subtract the mean and divide by the standard deviation.
// We can do this by applying the normalization to fscale and fadd outside the loop.
// Red
fadd_r = (fadd_r - input_data->mean[0]) / input_data->stdev[0];
fscale_r /= input_data->stdev[0];
// Green
fadd_g = (fadd_g - input_data->mean[1]) / input_data->stdev[1];
fscale_g /= input_data->stdev[1];
// Blue
fadd_b = (fadd_b - input_data->mean[2]) / input_data->stdev[2];
fscale_b /= input_data->stdev[2];
// Grayscale -> Y = 0.299R + 0.587G + 0.114B
float mean = (input_data->mean[0] * 0.299f) + (input_data->mean[1] * 0.587f) + (input_data->mean[2] * 0.114f);
float std = (input_data->stdev[0] * 0.299f) + (input_data->stdev[1] * 0.587f) + (input_data->stdev[2] * 0.114f);
fadd = (fadd - mean) / std;
fscale /= std;
image_t dst_img;
dst_img.w = input_width;
dst_img.h = input_height;
dst_img.data = (uint8_t *) model_input;
if (input_channels == 1) {
dst_img.pixfmt = PIXFORMAT_GRAYSCALE;
} else if (input_channels == 3) {
dst_img.pixfmt = PIXFORMAT_RGB565;
} else {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Expected model input channels to be 1 or 3!"));
}
imlib_draw_image(&dst_img, input_data->data, 0, 0, 1.0f, 1.0f, &input_data->roi,
-1, 256, NULL, NULL, IMAGE_HINT_BILINEAR | IMAGE_HINT_CENTER |
IMAGE_HINT_SCALE_ASPECT_EXPAND | IMAGE_HINT_BLACK_BACKGROUND, NULL, NULL, NULL);
int size = (input_width * input_height) - 1; // must be int per countdown loop
if (input_channels == 1) {
// GRAYSCALE
if (model->input_dtype == PY_ML_DTYPE_FLOAT) {
// convert u8 -> f32
uint8_t *model_input_u8 = (uint8_t *) model_input;
float *model_input_f32 = (float *) model_input;
for (; size >= 0; size -= 1) {
model_input_f32[size] = (model_input_u8[size] * fscale) + fadd;
if (input_array->ndim != input_shape->len) {
mp_raise_msg(&mp_type_ValueError,
MP_ERROR_TEXT("Input shape does not match the model input shape"));
}
} else {
if (shift) {
// convert u8 -> s8
uint8_t *model_input_8 = (uint8_t *) model_input;
#if (__ARM_ARCH > 6)
for (; size >= 3; size -= 4) {
*((uint32_t *) (model_input_8 + size - 3)) ^= 0x80808080;
}
#endif
for (; size >= 0; size -= 1) {
model_input_8[size] ^= PY_ML_GRAYSCALE_MID;
for (size_t i = 0; i < input_array->ndim; i++) {
if (input_array->shape[i] != mp_obj_get_int(input_shape->items[i])) {
mp_raise_msg(&mp_type_ValueError,
MP_ERROR_TEXT("Input shape does not match the model input shape"));
}
}
}
} else if (input_channels == 3) {
// RGB888
int rgb_size = size * 3; // must be int per countdown loop
if (model->input_dtype == PY_ML_DTYPE_FLOAT) {
uint16_t *model_input_u16 = (uint16_t *) model_input;
float *model_input_f32 = (float *) model_input;
for (; size >= 0; size -= 1, rgb_size -= 3) {
int pixel = model_input_u16[size];
model_input_f32[rgb_size] = (COLOR_RGB565_TO_R8(pixel) * fscale_r) + fadd_r;
model_input_f32[rgb_size + 1] = (COLOR_RGB565_TO_G8(pixel) * fscale_g) + fadd_g;
model_input_f32[rgb_size + 2] = (COLOR_RGB565_TO_B8(pixel) * fscale_b) + fadd_b;
if (model->input_dtype == 'f') {
float *model_input_float = (float *) input_buffer;
for (size_t i = 0; i < input_array->len; i++) {
float value = ndarray_get_float_index(input_array->array, input_array->dtype, i);
model_input_float[i] = value;
}
} else if (model->input_dtype == 'b') {
int8_t *model_input_8 = (int8_t *) input_buffer;
for (size_t i = 0; i < input_array->len; i++) {
float value = ndarray_get_float_index(input_array->array, input_array->dtype, i);
model_input_8[i] = (int8_t) ((value / model->input_scale) + model->input_zero_point);
}
} else if (model->input_dtype == 'B') {
uint8_t *model_input_8 = (uint8_t *) input_buffer;
for (size_t i = 0; i < input_array->len; i++) {
float value = ndarray_get_float_index(input_array->array, input_array->dtype, i);
model_input_8[i] = (uint8_t) ((value / model->input_scale) + model->input_zero_point);
}
} else {
int16_t *model_input_16 = (int16_t *) input_buffer;
for (size_t i = 0; i < input_array->len; i++) {
float value = ndarray_get_float_index(input_array->array, input_array->dtype, i);
model_input_16[i] = (int16_t) ((value / model->input_scale) + model->input_zero_point);
}
}
} else {
uint16_t *model_input_u16 = (uint16_t *) model_input;
uint8_t *model_input_8 = (uint8_t *) model_input;
for (; size >= 0; size -= 1, rgb_size -= 3) {
int pixel = model_input_u16[size];
model_input_8[rgb_size] = COLOR_RGB565_TO_R8(pixel) ^ shift;
model_input_8[rgb_size + 1] = COLOR_RGB565_TO_G8(pixel) ^ shift;
model_input_8[rgb_size + 2] = COLOR_RGB565_TO_B8(pixel) ^ shift;
}
}
}
}
static void py_ml_input_callback_regression(py_ml_model_obj_t *model, void *arg) {
// TODO we assume that there's a single input.
void *model_input = ml_backend_get_input(model, 0);
py_ml_input_data_t *input_data = (py_ml_input_data_t *) arg;
mp_obj_tuple_t *input_shape = MP_OBJ_TO_PTR(model->input_shape->items[0]);
ndarray_obj_t *input_array = MP_OBJ_TO_PTR(*((mp_obj_t *) input_data->data));
if (input_array->ndim != input_shape->len) {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Input shape does not match the model input shape"));
}
for (size_t i = 0; i < input_array->ndim; i++) {
if (input_array->shape[i] != mp_obj_get_int(input_shape->items[i])) {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Input shape does not match the model input shape"));
}
}
if (model->input_dtype == PY_ML_DTYPE_FLOAT) {
float *model_input_float = (float *) model_input;
for (size_t i = 0; i < input_array->len; i++) {
float value = ndarray_get_float_index(input_array->array, input_array->dtype, i);
model_input_float[i] = value;
}
} else if (model->input_dtype == PY_ML_DTYPE_INT8) {
int8_t *model_input_8 = (int8_t *) model_input;
for (size_t i = 0; i < input_array->len; i++) {
float value = ndarray_get_float_index(input_array->array, input_array->dtype, i);
model_input_8[i] = (int8_t) ((value / model->input_scale) + model->input_zero_point);
}
} else if (model->input_dtype == PY_ML_DTYPE_UINT8) {
uint8_t *model_input_8 = (uint8_t *) model_input;
for (size_t i = 0; i < input_array->len; i++) {
float value = ndarray_get_float_index(input_array->array, input_array->dtype, i);
model_input_8[i] = (uint8_t) ((value / model->input_scale) + model->input_zero_point);
}
} else {
int16_t *model_input_16 = (int16_t *) model_input;
for (size_t i = 0; i < input_array->len; i++) {
float value = ndarray_get_float_index(input_array->array, input_array->dtype, i);
model_input_16[i] = (int16_t) ((value / model->input_scale) + model->input_zero_point);
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Unsupported input type"));
}
}
}
@ -243,16 +109,16 @@ static void py_ml_output_callback(py_ml_model_obj_t *model, void *arg) {
size_t size = py_ml_tuple_sum(MP_OBJ_TO_PTR(model->output_shape->items[i]));
mp_obj_tuple_t *output = MP_OBJ_TO_PTR(mp_obj_new_tuple(size, NULL));
if (model->output_dtype == PY_ML_DTYPE_FLOAT) {
if (model->output_dtype == 'f') {
for (size_t j = 0; j < size; j++) {
output->items[j] = mp_obj_new_float(((float *) model_output)[j]);
}
} else if (model->output_dtype == PY_ML_DTYPE_INT8) {
} else if (model->output_dtype == 'b') {
for (size_t j = 0; j < size; j++) {
float v = (((int8_t *) model_output)[j] - model->output_zero_point);
output->items[j] = mp_obj_new_float(v * model->output_scale);
}
} else if (model->output_dtype == PY_ML_DTYPE_UINT8) {
} else if (model->output_dtype == 'B') {
for (size_t j = 0; j < size; j++) {
float v = (((uint8_t *) model_output)[j] - model->output_zero_point);
output->items[j] = mp_obj_new_float(v * model->output_scale);
@ -265,7 +131,7 @@ static void py_ml_output_callback(py_ml_model_obj_t *model, void *arg) {
}
output_list->items[i] = MP_OBJ_FROM_PTR(output);
}
*((py_ml_output_data_t *) arg) = MP_OBJ_FROM_PTR(output_list);
*((mp_obj_t *) arg) = MP_OBJ_FROM_PTR(output_list);
}
// TF Model Object.
@ -274,21 +140,18 @@ static const mp_obj_type_t py_ml_model_type;
static void py_ml_model_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind) {
py_ml_model_obj_t *self = MP_OBJ_TO_PTR(self_in);
mp_printf(print,
"{size: %d, ram: %d, inputs_size: %d, input_dtype: %s, input_scale: %f, input_zero_point: %d, "
"outputs_size: %d output_dtype: %s, output_scale: %f, output_zero_point: %d}",
self->size, self->memory_size, self->inputs_size, py_ml_map_dtype(self->input_dtype),
(double) self->input_scale, self->input_zero_point, self->outputs_size, py_ml_map_dtype(self->output_dtype),
"{size: \"%d\", ram: \"%d\","
" inputs_size: \"%d\", input_dtype: \"%c\", input_scale: \"%f\", input_zero_point: \"%d\","
" outputs_size: \"%d\" output_dtype: \"%c\", output_scale: \"%f\", output_zero_point: \"%d\"}",
self->size, self->memory_size, self->inputs_size, self->input_dtype,
(double) self->input_scale, self->input_zero_point, self->outputs_size, self->output_dtype,
(double) self->output_scale, self->output_zero_point);
}
static mp_obj_t py_ml_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 };
enum { ARG_callback };
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_callback, 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_ML_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.
@ -297,41 +160,22 @@ static mp_obj_t py_ml_model_predict(uint n_args, const mp_obj_t *pos_args, mp_ma
py_ml_model_obj_t *model = MP_OBJ_TO_PTR(pos_args[0]);
py_ml_input_data_t input_data = {
.scale = args[ARG_scale].u_int,
.mean = {0.0f, 0.0f, 0.0f},
.stdev = {1.0f, 1.0f, 1.0f}
};
mp_obj_t input_data = pos_args[1];
ml_backend_input_callback_t input_callback = py_ml_input_callback;
py_ml_output_data_t output_data;
mp_obj_t output_data;
ml_backend_output_callback_t output_callback = py_ml_output_callback;
if (MP_OBJ_IS_TYPE(pos_args[1], &ulab_ndarray_type)) {
input_data.data = (void *) &pos_args[1];
input_callback = py_ml_input_callback_regression;
} else if (MP_OBJ_IS_TYPE(pos_args[1], &py_image_type)) {
input_data.data = py_helper_arg_to_image(pos_args[1], ARG_IMAGE_ANY);
input_data.roi = py_helper_arg_to_roi(args[ARG_roi].u_obj, input_data.data);
py_helper_arg_to_float_array(args[ARG_mean].u_obj, input_data.mean, 3);
py_helper_arg_to_float_array(args[ARG_stdev].u_obj, input_data.stdev, 3);
} else {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Unsupported input type"));
if (!MP_OBJ_IS_TYPE(input_data, &mp_type_list)) {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Unsupported input type. Expected a list"));
}
ml_backend_run_inference(model, input_callback, &input_data, output_callback, &output_data);
if (args[ARG_callback].u_obj != mp_const_none) {
mp_obj_t rect = mp_obj_new_tuple(4, (mp_obj_t []) { mp_obj_new_int(input_data.roi.x),
mp_obj_new_int(input_data.roi.y),
mp_obj_new_int(input_data.roi.w),
mp_obj_new_int(input_data.roi.h) });
mp_obj_t fun_args[3] = { MP_OBJ_FROM_PTR(model), output_data, rect };
if (!MP_OBJ_IS_TYPE(pos_args[1], &py_image_type)) {
output_data = mp_call_function_n_kw(args[ARG_callback].u_obj, 2, 0, fun_args);
} else {
output_data = mp_call_function_n_kw(args[ARG_callback].u_obj, 3, 0, fun_args);
}
// Pass model, inputs, outputs to the post-processing callback.
mp_obj_t fargs[3] = { MP_OBJ_FROM_PTR(model), pos_args[1], output_data };
output_data = mp_call_function_n_kw(args[ARG_callback].u_obj, 3, 0, fargs);
}
return output_data;
@ -340,7 +184,7 @@ static MP_DEFINE_CONST_FUN_OBJ_KW(py_ml_model_predict_obj, 2, py_ml_model_predic
static void py_ml_model_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
py_ml_model_obj_t *self = MP_OBJ_TO_PTR(self_in);
const char *str;
if (dest[0] == MP_OBJ_NULL) {
// Load attribute.
switch (attr) {
@ -354,8 +198,7 @@ static void py_ml_model_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
dest[0] = MP_OBJ_FROM_PTR(self->input_shape);
break;
case MP_QSTR_input_dtype:
str = py_ml_map_dtype(self->input_dtype);
dest[0] = mp_obj_new_str(str, strlen(str));
dest[0] = mp_obj_new_str(&self->input_dtype, 1);
break;
case MP_QSTR_input_scale:
dest[0] = mp_obj_new_float(self->input_scale);
@ -367,8 +210,7 @@ static void py_ml_model_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
dest[0] = MP_OBJ_FROM_PTR(self->output_shape);
break;
case MP_QSTR_output_dtype:
str = py_ml_map_dtype(self->output_dtype);
dest[0] = mp_obj_new_str(str, strlen(str));
dest[0] = mp_obj_new_str(&self->output_dtype, 1);
break;
case MP_QSTR_output_scale:
dest[0] = mp_obj_new_float(self->output_scale);
@ -490,10 +332,6 @@ 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) },
{ MP_ROM_QSTR(MP_QSTR_SCALE_NONE), MP_ROM_INT(PY_ML_SCALE_NONE) },
{ MP_ROM_QSTR(MP_QSTR_SCALE_0_1), MP_ROM_INT(PY_ML_SCALE_0_1) },
{ MP_ROM_QSTR(MP_QSTR_SCALE_S1_1), MP_ROM_INT(PY_ML_SCALE_S1_1) },
{ MP_ROM_QSTR(MP_QSTR_SCALE_S128_127), MP_ROM_INT(PY_ML_SCALE_S128_127) },
};
static MP_DEFINE_CONST_DICT(py_ml_globals_dict, py_ml_globals_dict_table);

View File

@ -10,20 +10,6 @@
*/
#ifndef __PY_ML_H__
#define __PY_ML_H__
typedef enum {
PY_ML_SCALE_NONE,
PY_ML_SCALE_0_1,
PY_ML_SCALE_S1_1,
PY_ML_SCALE_S128_127
} py_ml_scale_t;
typedef enum py_ml_dtype {
PY_ML_DTYPE_INT8,
PY_ML_DTYPE_UINT8,
PY_ML_DTYPE_INT16,
PY_ML_DTYPE_FLOAT
} py_ml_dtype_t;
// TF Model Object.
typedef struct py_ml_model_obj {
mp_obj_base_t base;
@ -35,12 +21,12 @@ typedef struct py_ml_model_obj {
mp_obj_tuple_t *input_shape;
float input_scale;
int input_zero_point;
py_ml_dtype_t input_dtype;
char input_dtype;
size_t outputs_size;
mp_obj_tuple_t *output_shape;
float output_scale;
int output_zero_point;
py_ml_dtype_t output_dtype;
char output_dtype;
void *state; // Private context for the backend.
} py_ml_model_obj_t;