scripts/libraries: Add yolo v5 post processing support.

This commit is contained in:
Kwabena W. Agyeman 2024-12-30 00:05:40 -08:00
parent 47cede4ce0
commit 2b98a4e963

View File

@ -150,3 +150,60 @@ class yolo_v2_postprocess:
y_center[i] + (h_rel[i] / 2), y_center[i] + (h_rel[i] / 2),
bb_scores[i], bb_classes[i]) bb_scores[i], bb_classes[i])
return nms.get_bounding_boxes() return nms.get_bounding_boxes()
class yolo_v5_postprocess:
_YOLO_V5_CX = const(0)
_YOLO_V5_CY = const(1)
_YOLO_V5_CW = const(2)
_YOLO_V5_CH = const(3)
_YOLO_V5_SCORE = const(4)
_YOLO_V5_CLASSES = const(5)
_NO_DETECTION = const(())
def __init__(self, threshold=0.6):
self.threshold = threshold
def __call__(self, model, inputs, outputs):
oh, ow, oc = model.output_shape[0]
class_count = oc - _YOLO_V5_CLASSES
# Reshape the output to a 2D array
colum_outputs = outputs[0].reshape((oh * ow, _YOLO_V5_CLASSES + class_count))
# Threshold all the scores
score_indices = colum_outputs[:, _YOLO_V5_SCORE]
score_indices = np.nonzero(score_indices > self.threshold)
if isinstance(score_indices, tuple):
score_indices = score_indices[0]
if not len(score_indices):
return _NO_DETECTION
# Get the bounding boxes that have a valid score
bb = np.take(colum_outputs, score_indices, axis=0)
# Get the score information
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)
# Compute the bounding box information
x_center = bb[:, _YOLO_V5_CX]
y_center = bb[:, _YOLO_V5_CY]
w_rel = bb[:, _YOLO_V5_CW] * 0.5
h_rel = bb[:, _YOLO_V5_CH] * 0.5
# Scale the bounding boxes to have enough integer precision for NMS
ib, ih, iw, ic = model.input_shape[0]
xmin = (x_center - w_rel) * iw
ymin = (y_center - h_rel) * ih
xmax = (x_center + w_rel) * iw
ymax = (y_center + h_rel) * ih
nms = NMS(iw, ih, inputs[0].roi)
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()