|
| 1 | +#include "py/dynruntime.h" |
| 2 | +#include "model.h" |
| 3 | +#include <string.h> |
| 4 | + |
| 5 | +typedef struct _dc_obj_t { |
| 6 | + mp_obj_base_t base; |
| 7 | + bool model_state; |
| 8 | + uint8_t model_in_dim; |
| 9 | + uint8_t model_out_dim; |
| 10 | +} dc_obj_t; |
| 11 | + |
| 12 | +mp_obj_t get_model_input_dim(mp_obj_t self_in){ |
| 13 | + dc_obj_t *self = MP_OBJ_TO_PTR(self_in); |
| 14 | + return MP_OBJ_NEW_SMALL_INT(self->model_in_dim); |
| 15 | +} |
| 16 | + |
| 17 | +mp_obj_t get_model_output_dim(mp_obj_t self_in){ |
| 18 | + dc_obj_t *self = MP_OBJ_TO_PTR(self_in); |
| 19 | + return MP_OBJ_NEW_SMALL_INT(self->model_out_dim); |
| 20 | +} |
| 21 | + |
| 22 | +mp_obj_t init(mp_obj_t self_in){ |
| 23 | + dc_obj_t *self = MP_OBJ_TO_PTR(self_in); |
| 24 | + self->model_state = true; |
| 25 | + self->model_in_dim = IMAI_DATA_IN_COUNT; |
| 26 | + self->model_out_dim = IMAI_DATA_OUT_COUNT; |
| 27 | + IMAI_init(); |
| 28 | + return mp_const_none; |
| 29 | +} |
| 30 | + |
| 31 | +mp_obj_t enqueue(mp_obj_t self_in, const mp_obj_t data_in_obj){ |
| 32 | + dc_obj_t *self = MP_OBJ_TO_PTR(self_in); |
| 33 | + |
| 34 | + // Check if model is initialized |
| 35 | + if(!self->model_state){ |
| 36 | + mp_raise_ValueError("Model should be initialized first."); |
| 37 | + } |
| 38 | + |
| 39 | + float data_in[self->model_in_dim]; |
| 40 | + mp_obj_t *data_in_items; |
| 41 | + size_t len; |
| 42 | + mp_obj_get_array(data_in_obj, &len, &data_in_items); |
| 43 | + |
| 44 | + if (len != self->model_in_dim) { |
| 45 | + mp_raise_ValueError("data_in must be a list of floats with size matching to input dimensions to model. Check using get_model_input_dim()."); |
| 46 | + } |
| 47 | + |
| 48 | + for (int i = 0; i < self->model_in_dim; i++) { |
| 49 | + data_in[i] = mp_obj_get_float(data_in_items[i]); |
| 50 | + } |
| 51 | + int result = IMAI_enqueue(data_in); |
| 52 | + return MP_OBJ_NEW_SMALL_INT(result); |
| 53 | +} |
| 54 | + |
| 55 | +mp_obj_t dequeue(mp_obj_t self_in, mp_obj_t data_out_obj) { |
| 56 | + dc_obj_t *self = MP_OBJ_TO_PTR(self_in); |
| 57 | + |
| 58 | + mp_buffer_info_t buf_info; |
| 59 | + mp_get_buffer(data_out_obj, &buf_info, MP_BUFFER_WRITE); |
| 60 | + float *data_out = (float *)buf_info.buf; |
| 61 | + size_t len = buf_info.len / sizeof(float); |
| 62 | + |
| 63 | + if (len != self->model_out_dim) { |
| 64 | + mp_raise_ValueError("data_out must be a list of floats with size matching to output dimensions of model. Check using get_model_output_dim()."); |
| 65 | + } |
| 66 | + |
| 67 | + int result = IMAI_dequeue(data_out); |
| 68 | + if (result == 0) { |
| 69 | + return MP_OBJ_NEW_SMALL_INT(result); |
| 70 | + } else if (result == -1) { |
| 71 | + return MP_OBJ_NEW_SMALL_INT(result); |
| 72 | + } else if (result == -2) { |
| 73 | + mp_raise_ValueError(MP_ERROR_TEXT("Internal memory allocation error")); |
| 74 | + } |
| 75 | + return MP_OBJ_NEW_SMALL_INT(result); |
| 76 | +} |
| 77 | + |
| 78 | + |
| 79 | +static const mp_obj_fun_builtin_fixed_t init_obj = { |
| 80 | + .base = { &mp_type_fun_builtin_1 }, |
| 81 | + .fun._1 = (mp_fun_1_t)init, |
| 82 | +}; |
| 83 | + |
| 84 | + |
| 85 | +static const mp_obj_fun_builtin_fixed_t enqueue_obj = { |
| 86 | + .base = { &mp_type_fun_builtin_2 }, |
| 87 | + .fun._2 = (mp_fun_2_t)enqueue, |
| 88 | +}; |
| 89 | + |
| 90 | +static const mp_obj_fun_builtin_fixed_t dequeue_obj = { |
| 91 | + .base = { &mp_type_fun_builtin_2 }, |
| 92 | + .fun._2 = (mp_fun_2_t)dequeue, |
| 93 | +}; |
| 94 | + |
| 95 | +static const mp_obj_fun_builtin_fixed_t get_model_input_dim_obj = { |
| 96 | + .base = { &mp_type_fun_builtin_1 }, |
| 97 | + .fun._1 = (mp_fun_1_t)get_model_input_dim, |
| 98 | +}; |
| 99 | + |
| 100 | +static const mp_obj_fun_builtin_fixed_t get_model_output_dim_obj = { |
| 101 | + .base = { &mp_type_fun_builtin_1 }, |
| 102 | + .fun._1 = (mp_fun_1_t)get_model_output_dim, |
| 103 | +}; |
0 commit comments