mirror of
https://github.com/openmv/openmv.git
synced 2025-11-04 14:49:50 +08:00
scripts/libraries: Transpose YOLOV8 output.
The YOLOV8 tensor is transposed by default which is slow to access.
This commit is contained in:
parent
fafa8271ec
commit
09330d92ba
@ -258,28 +258,28 @@ class yolo_v8_postprocess:
|
|||||||
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
|
||||||
column_outputs = outputs[0].reshape((oh * (_YOLO_V8_CLASSES + class_count), oc))
|
row_outputs = outputs[0].reshape((oh * (_YOLO_V8_CLASSES + class_count), oc)).T
|
||||||
|
|
||||||
# Threshold all the scores
|
# Threshold all the scores
|
||||||
score_indices = np.max(dequantize(column_outputs[_YOLO_V8_CLASSES:, :], dt, zp, s), axis=0)
|
score_indices = np.max(dequantize(row_outputs[:, _YOLO_V8_CLASSES:], dt, zp, s), axis=1)
|
||||||
score_indices = np.nonzero(score_indices > self.threshold)[0]
|
score_indices = np.nonzero(score_indices > self.threshold)[0]
|
||||||
if not len(score_indices):
|
if not len(score_indices):
|
||||||
return _NO_DETECTION
|
return _NO_DETECTION
|
||||||
|
|
||||||
# Get the bounding boxes that have a valid score
|
# Get the bounding boxes that have a valid score
|
||||||
bb = dequantize(np.take(column_outputs, score_indices, axis=1), dt, zp, s)
|
bb = dequantize(np.take(row_outputs, score_indices, axis=0), dt, zp, s)
|
||||||
|
|
||||||
# Get the score information
|
# Get the score information
|
||||||
bb_scores = np.max(bb[_YOLO_V8_CLASSES:, :], axis=0)
|
bb_scores = np.max(bb[:, _YOLO_V8_CLASSES:], axis=1)
|
||||||
|
|
||||||
# Get the class information
|
# Get the class information
|
||||||
bb_classes = np.argmax(bb[_YOLO_V8_CLASSES:, :], axis=0)
|
bb_classes = np.argmax(bb[:, _YOLO_V8_CLASSES:], axis=1)
|
||||||
|
|
||||||
# Compute the bounding box information
|
# Compute the bounding box information
|
||||||
x_center = bb[_YOLO_V8_CX, :]
|
x_center = bb[:, _YOLO_V8_CX]
|
||||||
y_center = bb[_YOLO_V8_CY, :]
|
y_center = bb[:, _YOLO_V8_CY]
|
||||||
w_rel = bb[_YOLO_V8_CW, :]
|
w_rel = bb[:, _YOLO_V8_CW]
|
||||||
h_rel = bb[_YOLO_V8_CH, :]
|
h_rel = bb[:, _YOLO_V8_CH]
|
||||||
|
|
||||||
return ml.utils.box_nms(x_center, y_center, w_rel, h_rel, bb_scores, bb_classes,
|
return ml.utils.box_nms(x_center, y_center, w_rel, h_rel, bb_scores, bb_classes,
|
||||||
model.input_shape[0][1:3], inputs[0].roi,
|
model.input_shape[0][1:3], inputs[0].roi,
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user