mirror of
https://github.com/openmv/openmv.git
synced 2025-11-04 14:49:50 +08:00
modules/py_ml: Optimize threshold using Helium.
The majority of the time spent processing the model is during the max, threshold, and nonzero operation. Combing these togheter in C using Helium provides a massive speedup for larger arrays.
This commit is contained in:
parent
a40976e1b3
commit
77950f03e1
154
lib/imlib/simd.h
154
lib/imlib/simd.h
@ -867,6 +867,26 @@ static inline v128_t vadd_s32(v128_t v0, v128_t v1) {
|
|||||||
#endif
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
|
static inline v128_t vadd_n_u32(v128_t v0, uint32_t x) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vaddq_n_u32(v0.u32, x);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.u32 = v0.u32 + x
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
|
static inline v128_t vadd_n_s32(v128_t v0, int32_t x) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vaddq_n_s32(v0.s32, x);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.s32 = v0.s32 + x
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
static inline v128_t vsub_u8(v128_t v0, v128_t v1) {
|
static inline v128_t vsub_u8(v128_t v0, v128_t v1) {
|
||||||
#if (__ARM_ARCH >= 8)
|
#if (__ARM_ARCH >= 8)
|
||||||
return (v128_t) vsubq_u8(v0.u8, v1.u8);
|
return (v128_t) vsubq_u8(v0.u8, v1.u8);
|
||||||
@ -923,6 +943,26 @@ static inline v128_t vsub_s16(v128_t v0, v128_t v1) {
|
|||||||
#endif
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
|
static inline v128_t vsub_n_u32(v128_t v0, uint32_t x) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vsubq_n_u32(v0.u32, x);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.u32 = v0.u32 - x
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
|
static inline v128_t vsub_n_s32(v128_t v0, int32_t x) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vsubq_n_s32(v0.s32, x);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.s32 = v0.s32 - x
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
#if (__ARM_ARCH >= 8)
|
#if (__ARM_ARCH >= 8)
|
||||||
#define vsli_u8(v0, v1, n) ((v128_t) vsliq_n_u8(v0.u8, v1.u8, n))
|
#define vsli_u8(v0, v1, n) ((v128_t) vsliq_n_u8(v0.u8, v1.u8, n))
|
||||||
#else
|
#else
|
||||||
@ -1229,6 +1269,16 @@ static inline v128_t vmul_n_s32(v128_t v0, int32_t x) {
|
|||||||
#endif
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
|
static inline v128_t vmul_n_f32(v128_t v0, float32_t x) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vmulq_n_f32(v0.f32, x);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.f32 = v0.f32 * x
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
static inline v128_t vmla_n_u16(v128_t v0, uint16_t x, v128_t v2) {
|
static inline v128_t vmla_n_u16(v128_t v0, uint16_t x, v128_t v2) {
|
||||||
#if (__ARM_ARCH >= 8)
|
#if (__ARM_ARCH >= 8)
|
||||||
return (v128_t) vmlaq_n_u16(v2.u16, v0.u16, x);
|
return (v128_t) vmlaq_n_u16(v2.u16, v0.u16, x);
|
||||||
@ -1373,6 +1423,50 @@ static inline int32_t vmladava_s16(v128_t v0, v128_t v1, int32_t acc) {
|
|||||||
#endif
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
|
static inline v128_t vcvt_f32_u32(v128_t v0) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vcvtq(v0.u32);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.f32 = { (float32_t) v0.u32[0] }
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
|
static inline v128_t vcvt_f32_s32(v128_t v0) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vcvtq(v0.s32);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.f32 = { (float32_t) v0.s32[0] }
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
|
static inline float vminv_f32_pred(v128_t v, float min, v128_predicate_t pred) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return vminnmvq_p_f32(min, v.f32, pred);
|
||||||
|
#else
|
||||||
|
if (pred > 0) {
|
||||||
|
min = (v.f32[0] < min) ? v.f32[0] : min;
|
||||||
|
}
|
||||||
|
|
||||||
|
return min;
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
|
static inline float vmaxv_f32_pred(v128_t v, float max, v128_predicate_t pred) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return vmaxnmvq_p_f32(max, v.f32, pred);
|
||||||
|
#else
|
||||||
|
if (pred > 0) {
|
||||||
|
max = (v.f32[0] > max) ? v.f32[0] : max;
|
||||||
|
}
|
||||||
|
|
||||||
|
return max;
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
static inline v128_t vldr_u8(const uint8_t *p) {
|
static inline v128_t vldr_u8(const uint8_t *p) {
|
||||||
#if (__ARM_ARCH >= 8)
|
#if (__ARM_ARCH >= 8)
|
||||||
return (v128_t) vldrbq_u8(p);
|
return (v128_t) vldrbq_u8(p);
|
||||||
@ -1404,6 +1498,30 @@ static inline v128_t vldr_u8_pred(const uint8_t *p, v128_predicate_t pred) {
|
|||||||
#endif
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
|
static inline v128_t vldr_u8_widen_u32_gather_pred(const uint8_t *p,
|
||||||
|
v128_t offsets,
|
||||||
|
v128_predicate_t pred) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vldrbq_gather_offset_z_u32(p, offsets.u32, pred);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.u32 = { *(p + offsets.u32[0]) }
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
|
static inline v128_t vldr_s8_widen_s32_gather_pred(const int8_t *p,
|
||||||
|
v128_t offsets,
|
||||||
|
v128_predicate_t pred) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vldrbq_gather_offset_z_s32(p, offsets.u32, pred);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.s32 = { *(p + offsets.u32[0]) }
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
static inline void vstr_u8(uint8_t *p, v128_t v0) {
|
static inline void vstr_u8(uint8_t *p, v128_t v0) {
|
||||||
#if (__ARM_ARCH >= 8)
|
#if (__ARM_ARCH >= 8)
|
||||||
vstrbq(p, v0.u8);
|
vstrbq(p, v0.u8);
|
||||||
@ -1455,6 +1573,30 @@ static inline v128_t vldr_u16_pred(const uint16_t *p, v128_predicate_t pred) {
|
|||||||
#endif
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
|
static inline v128_t vldr_u16_widen_u32_gather_pred(const uint16_t *p,
|
||||||
|
v128_t offsets,
|
||||||
|
v128_predicate_t pred) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vldrhq_gather_shifted_offset_z_u32(p, offsets.u32, pred);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.u32 = { *(p + offsets.u32[0]) }
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
|
static inline v128_t vldr_s16_widen_s32_gather_pred(const int16_t *p,
|
||||||
|
v128_t offsets,
|
||||||
|
v128_predicate_t pred) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vldrhq_gather_shifted_offset_z_s32(p, offsets.u32, pred);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.s32 = { *(p + offsets.u32[0]) }
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
static inline void vstr_u16(uint16_t *p, v128_t v0) {
|
static inline void vstr_u16(uint16_t *p, v128_t v0) {
|
||||||
#if (__ARM_ARCH >= 8)
|
#if (__ARM_ARCH >= 8)
|
||||||
vstrhq(p, v0.u16);
|
vstrhq(p, v0.u16);
|
||||||
@ -1570,6 +1712,18 @@ static inline v4x_rows_t vldr_u32_gather_pred_x4_unaligned(v4x_row_ptrs_t rowptr
|
|||||||
return rows;
|
return rows;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
static inline v128_t vldr_f32_gather_pred(const float32_t *p,
|
||||||
|
v128_t offsets,
|
||||||
|
v128_predicate_t pred) {
|
||||||
|
#if (__ARM_ARCH >= 8)
|
||||||
|
return (v128_t) vldrwq_gather_shifted_offset_z_f32(p, offsets.u32, pred);
|
||||||
|
#else
|
||||||
|
return (v128_t) {
|
||||||
|
.f32 = { *(p + offsets.u32[0]) }
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
static inline void vstr_f32_scatter(float32_t *p, v128_t offsets, v128_t v0) {
|
static inline void vstr_f32_scatter(float32_t *p, v128_t offsets, v128_t v0) {
|
||||||
#if (__ARM_ARCH >= 8)
|
#if (__ARM_ARCH >= 8)
|
||||||
vstrwq_scatter_shifted_offset(p, offsets.u32, v0.f32);
|
vstrwq_scatter_shifted_offset(p, offsets.u32, v0.f32);
|
||||||
|
|||||||
151
modules/py_ml.c
151
modules/py_ml.c
@ -50,6 +50,7 @@
|
|||||||
#include "file_utils.h"
|
#include "file_utils.h"
|
||||||
#include "py_ml.h"
|
#include "py_ml.h"
|
||||||
#include "ulab/code/ndarray.h"
|
#include "ulab/code/ndarray.h"
|
||||||
|
#include "simd.h"
|
||||||
|
|
||||||
#define IMLIB_ML_MODEL_ALIGN (OMV_CACHE_LINE_SIZE)
|
#define IMLIB_ML_MODEL_ALIGN (OMV_CACHE_LINE_SIZE)
|
||||||
|
|
||||||
@ -426,11 +427,159 @@ static MP_DEFINE_CONST_OBJ_TYPE(
|
|||||||
locals_dict, &py_ml_model_locals_dict
|
locals_dict, &py_ml_model_locals_dict
|
||||||
);
|
);
|
||||||
|
|
||||||
extern const mp_obj_type_t py_ml_nms_type;
|
// The function finds the maximum value in each row of a 2D ndarray and returns the indices
|
||||||
|
// of the rows where the maximum exceeds a threshold which is the most CPU intensive post-processing
|
||||||
|
// step. This function can handle regular and transposed ndarrays where the strides result in
|
||||||
|
// non-contiguous memory access patterns. TODO: Futher performance optimizations are possible for
|
||||||
|
// contiguous int8/uint8/int16/uint16 ndarrays by processing contiguous rows using contiguous
|
||||||
|
// memory operations and integer SIMD operations.
|
||||||
|
static mp_obj_t py_ml_threshold(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
|
||||||
|
OMV_PROFILE_START();
|
||||||
|
enum { ARG_output_scale, ARG_output_zero_point, ARG_threshold };
|
||||||
|
static const mp_arg_t allowed_args[] = {
|
||||||
|
{ MP_QSTR_output_scale, MP_ARG_OBJ, {.u_rom_obj = MP_ROM_NONE} },
|
||||||
|
{ MP_QSTR_output_zero_point, MP_ARG_INT, {.u_int = 0 } },
|
||||||
|
{ MP_QSTR_threshold, MP_ARG_OBJ, {.u_rom_obj = MP_ROM_NONE} }
|
||||||
|
};
|
||||||
|
|
||||||
|
if (!MP_OBJ_IS_TYPE(pos_args[0], &ulab_ndarray_type)) {
|
||||||
|
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Expected an ndarray."));
|
||||||
|
}
|
||||||
|
|
||||||
|
ndarray_obj_t *input = MP_OBJ_TO_PTR(pos_args[0]);
|
||||||
|
size_t height, width, row_stride, value_stride;
|
||||||
|
|
||||||
|
if (input->ndim == 1) {
|
||||||
|
height = input->shape[ULAB_MAX_DIMS - 1];
|
||||||
|
width = 1;
|
||||||
|
row_stride = input->strides[ULAB_MAX_DIMS - 1] / input->itemsize;
|
||||||
|
value_stride = 1;
|
||||||
|
} else if (input->ndim == 2) {
|
||||||
|
height = input->shape[ULAB_MAX_DIMS - 2];
|
||||||
|
width = input->shape[ULAB_MAX_DIMS - 1];
|
||||||
|
row_stride = input->strides[ULAB_MAX_DIMS - 2] / input->itemsize;
|
||||||
|
value_stride = input->strides[ULAB_MAX_DIMS - 1] / input->itemsize;
|
||||||
|
} else {
|
||||||
|
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Expected a 1D or 2D ndarray."));
|
||||||
|
}
|
||||||
|
|
||||||
|
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);
|
||||||
|
|
||||||
|
float output_scale = py_helper_arg_to_float(args[ARG_output_scale].u_obj, 1.0f);
|
||||||
|
int output_zero_point = args[ARG_output_zero_point].u_int;
|
||||||
|
float threshold = py_helper_arg_to_float(args[ARG_threshold].u_obj, 0.1f);
|
||||||
|
|
||||||
|
fb_alloc_mark();
|
||||||
|
uint16_t *output_array = (uint16_t *) fb_alloc(height * sizeof(uint16_t), FB_ALLOC_PREFER_SPEED);
|
||||||
|
v128_t offsets_start = vmul_n_u32(vidup_u32(0, 1), value_stride);
|
||||||
|
size_t offsets_inc = FLOAT32_VECTOR_SIZE * value_stride;
|
||||||
|
size_t count = 0;
|
||||||
|
|
||||||
|
if (input->dtype == 'f') {
|
||||||
|
for (size_t y = 0; y < height; y++) {
|
||||||
|
v128_t offsets = vadd_n_u32(offsets_start, y * row_stride);
|
||||||
|
float row_max = -FLT_MAX;
|
||||||
|
|
||||||
|
for (size_t m = 0; m < width; m += FLOAT32_VECTOR_SIZE) {
|
||||||
|
v128_predicate_t pred = vpredicate_32(width - m);
|
||||||
|
v128_t v = vldr_f32_gather_pred((float32_t *) input->array, offsets, pred);
|
||||||
|
row_max = vmaxv_f32_pred(v, row_max, pred);
|
||||||
|
offsets = vadd_n_u32(offsets, offsets_inc);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (row_max > threshold) {
|
||||||
|
output_array[count++] = y;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else if (input->dtype == 'b') {
|
||||||
|
for (size_t y = 0; y < height; y++) {
|
||||||
|
v128_t offsets = vadd_n_u32(offsets_start, y * row_stride);
|
||||||
|
float row_max = -FLT_MAX;
|
||||||
|
|
||||||
|
for (size_t x = 0; x < width; x += FLOAT32_VECTOR_SIZE) {
|
||||||
|
v128_predicate_t pred = vpredicate_32(width - x);
|
||||||
|
v128_t v = vldr_s8_widen_s32_gather_pred((int8_t *) input->array, offsets, pred);
|
||||||
|
v = vmul_n_f32(vcvt_f32_s32(vsub_n_s32(v, output_zero_point)), output_scale);
|
||||||
|
row_max = vmaxv_f32_pred(v, row_max, pred);
|
||||||
|
offsets = vadd_n_u32(offsets, offsets_inc);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (row_max > threshold) {
|
||||||
|
output_array[count++] = y;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else if (input->dtype == 'B') {
|
||||||
|
for (size_t y = 0; y < height; y++) {
|
||||||
|
v128_t offsets = vadd_n_u32(offsets_start, y * row_stride);
|
||||||
|
float row_max = -FLT_MAX;
|
||||||
|
|
||||||
|
for (size_t x = 0; x < width; x += FLOAT32_VECTOR_SIZE) {
|
||||||
|
v128_predicate_t pred = vpredicate_32(width - x);
|
||||||
|
v128_t v = vldr_u8_widen_u32_gather_pred((uint8_t *) input->array, offsets, pred);
|
||||||
|
v = vmul_n_f32(vcvt_f32_s32(vsub_n_s32(v, output_zero_point)), output_scale);
|
||||||
|
row_max = vmaxv_f32_pred(v, row_max, pred);
|
||||||
|
offsets = vadd_n_u32(offsets, offsets_inc);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (row_max > threshold) {
|
||||||
|
output_array[count++] = y;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else if (input->dtype == 'h') {
|
||||||
|
for (size_t y = 0; y < height; y++) {
|
||||||
|
v128_t offsets = vadd_n_u32(offsets_start, y * row_stride);
|
||||||
|
float row_max = -FLT_MAX;
|
||||||
|
|
||||||
|
for (size_t x = 0; x < width; x += FLOAT32_VECTOR_SIZE) {
|
||||||
|
v128_predicate_t pred = vpredicate_32(width - x);
|
||||||
|
v128_t v = vldr_s16_widen_s32_gather_pred((int16_t *) input->array, offsets, pred);
|
||||||
|
v = vmul_n_f32(vcvt_f32_s32(vsub_n_s32(v, output_zero_point)), output_scale);
|
||||||
|
row_max = vmaxv_f32_pred(v, row_max, pred);
|
||||||
|
offsets = vadd_n_u32(offsets, offsets_inc);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (row_max > threshold) {
|
||||||
|
output_array[count++] = y;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else if (input->dtype == 'H') {
|
||||||
|
for (size_t y = 0; y < height; y++) {
|
||||||
|
v128_t offsets = vadd_n_u32(offsets_start, y * row_stride);
|
||||||
|
float row_max = -FLT_MAX;
|
||||||
|
|
||||||
|
for (size_t x = 0; x < width; x += FLOAT32_VECTOR_SIZE) {
|
||||||
|
v128_predicate_t pred = vpredicate_32(width - x);
|
||||||
|
v128_t v = vldr_u16_widen_u32_gather_pred((uint16_t *) input->array, offsets, pred);
|
||||||
|
v = vmul_n_f32(vcvt_f32_s32(vsub_n_s32(v, output_zero_point)), output_scale);
|
||||||
|
row_max = vmaxv_f32_pred(v, row_max, pred);
|
||||||
|
offsets = vadd_n_u32(offsets, offsets_inc);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (row_max > threshold) {
|
||||||
|
output_array[count++] = y;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
mp_raise_ValueError(MP_ERROR_TEXT("Unsupported dtype"));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Copy the output array to a new ndarray.
|
||||||
|
size_t output_shape[ULAB_MAX_DIMS] = {};
|
||||||
|
output_shape[ULAB_MAX_DIMS - 1] = count;
|
||||||
|
ndarray_obj_t *output = ndarray_new_dense_ndarray(1, output_shape, NDARRAY_UINT16);
|
||||||
|
memcpy(output->array, output_array, count * sizeof(uint16_t));
|
||||||
|
fb_alloc_free_till_mark();
|
||||||
|
|
||||||
|
OMV_PROFILE_PRINT();
|
||||||
|
return MP_OBJ_FROM_PTR(output);
|
||||||
|
}
|
||||||
|
static MP_DEFINE_CONST_FUN_OBJ_KW(py_ml_threshold_obj, 1, py_ml_threshold);
|
||||||
|
|
||||||
static const mp_rom_map_elem_t py_ml_globals_dict_table[] = {
|
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___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_Model), MP_ROM_PTR(&py_ml_model_type) },
|
||||||
|
{ MP_ROM_QSTR(MP_QSTR_threshold), MP_ROM_PTR(&py_ml_threshold_obj) },
|
||||||
};
|
};
|
||||||
|
|
||||||
static MP_DEFINE_CONST_DICT(py_ml_globals_dict, py_ml_globals_dict_table);
|
static MP_DEFINE_CONST_DICT(py_ml_globals_dict, py_ml_globals_dict_table);
|
||||||
|
|||||||
@ -26,6 +26,7 @@
|
|||||||
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
||||||
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||||
|
import uml
|
||||||
import ml.utils
|
import ml.utils
|
||||||
from micropython import const
|
from micropython import const
|
||||||
from ulab import numpy as np
|
from ulab import numpy as np
|
||||||
@ -75,17 +76,6 @@ def softmax(x):
|
|||||||
return e_x / np.sum(e_x, axis=1, keepdims=True)
|
return e_x / np.sum(e_x, axis=1, keepdims=True)
|
||||||
|
|
||||||
|
|
||||||
def threshold(scores, threshold, scale, find_max=False, find_max_axis=1):
|
|
||||||
if scale > 0:
|
|
||||||
if find_max:
|
|
||||||
scores = np.max(scores, axis=find_max_axis)
|
|
||||||
return np.nonzero(scores > threshold)[0]
|
|
||||||
else:
|
|
||||||
if find_max:
|
|
||||||
scores = np.min(scores, axis=find_max_axis)
|
|
||||||
return np.nonzero(scores < threshold)[0]
|
|
||||||
|
|
||||||
|
|
||||||
class fomo_postprocess:
|
class fomo_postprocess:
|
||||||
def __init__(self, threshold=0.4, w_scale=1.414214, h_scale=1.414214,
|
def __init__(self, threshold=0.4, w_scale=1.414214, h_scale=1.414214,
|
||||||
nms_threshold=0.1, nms_sigma=0.001,
|
nms_threshold=0.1, nms_sigma=0.001,
|
||||||
@ -102,14 +92,13 @@ class fomo_postprocess:
|
|||||||
s = model.output_scale[0]
|
s = model.output_scale[0]
|
||||||
zp = model.output_zero_point[0]
|
zp = model.output_zero_point[0]
|
||||||
dt = model.output_dtype[0]
|
dt = model.output_dtype[0]
|
||||||
t = (self.threshold / s) + zp
|
|
||||||
|
|
||||||
# Reshape the output to a 2D array
|
# Reshape the output to a 2D array
|
||||||
row_outputs = outputs[0].reshape((oh * ow, oc))
|
row_outputs = outputs[0].reshape((oh * ow, oc))
|
||||||
|
|
||||||
# Threshold all the scores
|
# Threshold all the scores
|
||||||
score_indices = row_outputs[:, _FOMO_CLASSES:]
|
score_indices = row_outputs[:, _FOMO_CLASSES:]
|
||||||
score_indices = threshold(score_indices, t, s, find_max=True, find_max_axis=1)
|
score_indices = uml.threshold(score_indices, s, zp, self.threshold)
|
||||||
if not len(score_indices):
|
if not len(score_indices):
|
||||||
return _NO_DETECTION
|
return _NO_DETECTION
|
||||||
|
|
||||||
@ -169,7 +158,7 @@ class yolo_v2_postprocess:
|
|||||||
|
|
||||||
# Threshold all the scores
|
# Threshold all the scores
|
||||||
score_indices = row_outputs[:, _YOLO_V2_SCORE]
|
score_indices = row_outputs[:, _YOLO_V2_SCORE]
|
||||||
score_indices = threshold(score_indices, t, s)
|
score_indices = uml.threshold(score_indices, s, zp, logit(self.threshold))
|
||||||
if not len(score_indices):
|
if not len(score_indices):
|
||||||
return _NO_DETECTION
|
return _NO_DETECTION
|
||||||
|
|
||||||
@ -228,7 +217,6 @@ class yolo_v5_postprocess:
|
|||||||
s = model.output_scale[0]
|
s = model.output_scale[0]
|
||||||
zp = model.output_zero_point[0]
|
zp = model.output_zero_point[0]
|
||||||
dt = model.output_dtype[0]
|
dt = model.output_dtype[0]
|
||||||
t = (self.threshold / s) + zp
|
|
||||||
class_count = oc - _YOLO_V5_CLASSES
|
class_count = oc - _YOLO_V5_CLASSES
|
||||||
|
|
||||||
# Reshape the output to a 2D array
|
# Reshape the output to a 2D array
|
||||||
@ -236,7 +224,7 @@ class yolo_v5_postprocess:
|
|||||||
|
|
||||||
# Threshold all the scores
|
# Threshold all the scores
|
||||||
score_indices = row_outputs[:, _YOLO_V5_SCORE]
|
score_indices = row_outputs[:, _YOLO_V5_SCORE]
|
||||||
score_indices = threshold(score_indices, t, s)
|
score_indices = uml.threshold(score_indices, s, zp, self.threshold)
|
||||||
if not len(score_indices):
|
if not len(score_indices):
|
||||||
return _NO_DETECTION
|
return _NO_DETECTION
|
||||||
|
|
||||||
@ -273,7 +261,6 @@ class yolo_v8_postprocess:
|
|||||||
s = model.output_scale[0]
|
s = model.output_scale[0]
|
||||||
zp = model.output_zero_point[0]
|
zp = model.output_zero_point[0]
|
||||||
dt = model.output_dtype[0]
|
dt = model.output_dtype[0]
|
||||||
t = (self.threshold / s) + zp
|
|
||||||
class_count = ow - _YOLO_V8_CLASSES
|
class_count = ow - _YOLO_V8_CLASSES
|
||||||
|
|
||||||
# Reshape the output to a 2D array
|
# Reshape the output to a 2D array
|
||||||
@ -281,7 +268,7 @@ class yolo_v8_postprocess:
|
|||||||
|
|
||||||
# Threshold all the scores
|
# Threshold all the scores
|
||||||
score_indices = row_outputs[:, _YOLO_V8_CLASSES:]
|
score_indices = row_outputs[:, _YOLO_V8_CLASSES:]
|
||||||
score_indices = threshold(score_indices, t, s, find_max=True, find_max_axis=1)
|
score_indices = uml.threshold(score_indices, s, zp, self.threshold)
|
||||||
if not len(score_indices):
|
if not len(score_indices):
|
||||||
return _NO_DETECTION
|
return _NO_DETECTION
|
||||||
|
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user