mirror of
https://github.com/openmv/openmv.git
synced 2025-09-26 23:09:13 +08:00
scripts/libraries: Add support for yolov8 post-processing.
This commit is contained in:
parent
4e21df977b
commit
3e46eee35a
@ -206,3 +206,57 @@ class yolo_v5_postprocess:
|
|||||||
nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i],
|
nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i],
|
||||||
bb_scores[i], bb_classes[i])
|
bb_scores[i], bb_classes[i])
|
||||||
return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)
|
return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)
|
||||||
|
|
||||||
|
|
||||||
|
class yolo_v8_postprocess:
|
||||||
|
_YOLO_V8_CX = const(0)
|
||||||
|
_YOLO_V8_CY = const(1)
|
||||||
|
_YOLO_V8_CW = const(2)
|
||||||
|
_YOLO_V8_CH = const(3)
|
||||||
|
_YOLO_V8_CLASSES = const(4)
|
||||||
|
|
||||||
|
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]
|
||||||
|
class_count = ow - _YOLO_V8_CLASSES
|
||||||
|
|
||||||
|
# Reshape the output to a 2D array
|
||||||
|
column_outputs = outputs[0].reshape((oh * (_YOLO_V8_CLASSES + class_count), oc))
|
||||||
|
|
||||||
|
# Threshold all the scores
|
||||||
|
score_indices = np.max(column_outputs[_YOLO_V8_CLASSES:, :], axis=0)
|
||||||
|
score_indices = np.nonzero(score_indices > self.threshold)[0]
|
||||||
|
if not len(score_indices):
|
||||||
|
return _NO_DETECTION
|
||||||
|
|
||||||
|
# Get the bounding boxes that have a valid score
|
||||||
|
bb = np.take(column_outputs, score_indices, axis=1)
|
||||||
|
|
||||||
|
# Get the score information
|
||||||
|
bb_scores = np.max(bb[_YOLO_V8_CLASSES:, :], axis=0)
|
||||||
|
|
||||||
|
# Get the class information
|
||||||
|
bb_classes = np.argmax(bb[_YOLO_V8_CLASSES:, :], axis=0)
|
||||||
|
|
||||||
|
# Compute the bounding box information
|
||||||
|
x_center = bb[_YOLO_V8_CX, :]
|
||||||
|
y_center = bb[_YOLO_V8_CY, :]
|
||||||
|
w_rel = bb[_YOLO_V8_CW, :] * 0.5
|
||||||
|
h_rel = bb[_YOLO_V8_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(bb.shape[1]):
|
||||||
|
nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i],
|
||||||
|
bb_scores[i], bb_classes[i])
|
||||||
|
return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)
|
||||||
|
Loading…
Reference in New Issue
Block a user