Merge pull request #2582 from kwagyeman/kwabena/keep_dims

scripts/libraries: Simplify YOLO post-processing using keepdims.
This commit is contained in:
Ibrahim Abdelkader 2025-01-27 08:29:32 +02:00 committed by GitHub
commit ae1ada5db6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 21 additions and 20 deletions

View File

@ -72,17 +72,18 @@ class yolo_v2_postprocess:
_YOLO_V2_SCORE = const(4)
_YOLO_V2_CLASSES = const(5)
def __init__(self, threshold=0.6, anchors=None):
def __init__(self, threshold=0.6, anchors=None, nms_threshold=0.1, nms_sigma=0.1):
self.threshold = threshold
if anchors is not None:
self.anchors = anchors
else:
self.anchors = anchors
if self.anchors is None:
self.anchors = np.array([[0.98830, 3.36060],
[2.11940, 5.37590],
[3.05200, 9.13360],
[5.55170, 9.30660],
[9.72600, 11.1422]], dtype=np.float)
[9.72600, 11.1422]])
self.anchors_len = len(self.anchors)
self.nms_threshold = nms_threshold
self.nms_sigma = nms_sigma
def __call__(self, model, inputs, outputs):
ob, oh, ow, oc = model.output_shape[0]
@ -95,8 +96,8 @@ class yolo_v2_postprocess:
return a - (b * (a // b))
def softmax(x):
e_x = np.exp(x - np.max(x))
return e_x / np.sum(e_x)
e_x = np.exp(x - np.max(x, axis=1, keepdims=True))
return e_x / np.sum(e_x, axis=1, keepdims=True)
# Reshape the output to a 2D array
colum_outputs = outputs[0].reshape((oh * ow * self.anchors_len,
@ -119,18 +120,13 @@ class yolo_v2_postprocess:
bb_anchors = mod(score_indices, self.anchors_len)
# Get the anchor box information
bb_a_array = [self.anchors[i] for i in bb_anchors.tolist()]
bb_a_array = np.array(bb_a_array)
bb_a_array = np.take(self.anchors, bb_anchors, axis=0)
# Get the score information
bb_scores = sigmoid(bb[:, _YOLO_V2_SCORE])
# Get the class information
bb_classes = []
for i in range(len(score_indices)):
s = softmax(bb[i, _YOLO_V2_CLASSES:])
bb_classes.append(np.argmax(s))
bb_classes = np.array(bb_classes, dtype=np.uint16)
bb_classes = np.argmax(softmax(bb[:, _YOLO_V2_CLASSES:]), axis=1)
# Compute the bounding box information
x_center = (bb_cols + sigmoid(bb[:, _YOLO_V2_TX])) / ow
@ -152,7 +148,7 @@ class yolo_v2_postprocess:
x_center[i] + (w_rel[i] / 2),
y_center[i] + (h_rel[i] / 2),
bb_scores[i], bb_classes[i])
return nms.get_bounding_boxes()
return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)
class yolo_v5_postprocess:
@ -163,8 +159,10 @@ class yolo_v5_postprocess:
_YOLO_V5_SCORE = const(4)
_YOLO_V5_CLASSES = const(5)
def __init__(self, threshold=0.6):
def __init__(self, threshold=0.6, nms_threshold=0.1, nms_sigma=0.1):
self.threshold = threshold
self.nms_threshold = nms_threshold
self.nms_sigma = nms_sigma
def __call__(self, model, inputs, outputs):
oh, ow, oc = model.output_shape[0]
@ -188,8 +186,7 @@ class yolo_v5_postprocess:
bb_scores = bb[:, _YOLO_V5_SCORE]
# Get the class information
bb_classes = [np.argmax(bb[x, _YOLO_V5_CLASSES:]) for x in range(bb.shape[0])]
bb_classes = np.array(bb_classes, dtype=np.uint16)
bb_classes = np.argmax(bb[:, _YOLO_V5_CLASSES:], axis=1)
# Compute the bounding box information
x_center = bb[:, _YOLO_V5_CX]
@ -208,4 +205,4 @@ class yolo_v5_postprocess:
for i in range(len(bb)):
nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i],
bb_scores[i], bb_classes[i])
return nms.get_bounding_boxes()
return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)

@ -1 +1 @@
Subproject commit 303e8d790acc6e996c6851f00fa98122b3f85000
Subproject commit 20f7259a471fe180ad64279a6970eb991537cffb

View File

@ -454,6 +454,7 @@ FIRM_OBJ += $(addprefix $(BUILD)/$(MICROPY_DIR)/modules/ulab/code/,\
numpy/stats.o \
numpy/transform.o \
numpy/vector.o \
scipy/integrate/integrate.o \
scipy/linalg/linalg.o \
scipy/optimize/optimize.o \
scipy/scipy.o \

View File

@ -402,6 +402,7 @@ FIRM_OBJ += $(addprefix $(BUILD)/$(MICROPY_DIR)/modules/ulab/code/,\
numpy/stats.o \
numpy/transform.o \
numpy/vector.o \
scipy/integrate/integrate.o \
scipy/linalg/linalg.o \
scipy/optimize/optimize.o \
scipy/scipy.o \

View File

@ -276,6 +276,7 @@ if(MICROPY_PY_ULAB)
${MICROPY_ULAB_DIR}/code/numpy/stats.c
${MICROPY_ULAB_DIR}/code/numpy/transform.c
${MICROPY_ULAB_DIR}/code/numpy/vector.c
${MICROPY_ULAB_DIR}/code/scipy/integrate/integrate.c
${MICROPY_ULAB_DIR}/code/scipy/linalg/linalg.c
${MICROPY_ULAB_DIR}/code/scipy/optimize/optimize.c
${MICROPY_ULAB_DIR}/code/scipy/scipy.c

View File

@ -478,6 +478,7 @@ FIRM_OBJ += $(addprefix $(BUILD)/$(MICROPY_DIR)/modules/ulab/code/,\
numpy/stats.o \
numpy/transform.o \
numpy/vector.o \
scipy/integrate/integrate.o \
scipy/linalg/linalg.o \
scipy/optimize/optimize.o \
scipy/scipy.o \