mirror of
https://github.com/openmv/openmv.git
synced 2025-11-04 14:49:50 +08:00
Merge pull request #2690 from kwagyeman/kwabena/add_yolo_v8_support
scripts/libraries: Add support for more yolo post-processing.
This commit is contained in:
commit
c0a0c32d06
@ -64,6 +64,8 @@ class fomo_postprocess:
|
||||
return nms.get_bounding_boxes()
|
||||
|
||||
|
||||
# This is a lightweight version of the tiny yolo v2 object detection algorithm.
|
||||
# It was optimized to work well on embedded devices with limited computational resources.
|
||||
class yolo_v2_postprocess:
|
||||
_YOLO_V2_TX = const(0)
|
||||
_YOLO_V2_TY = const(1)
|
||||
@ -100,19 +102,17 @@ class yolo_v2_postprocess:
|
||||
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,
|
||||
row_outputs = outputs[0].reshape((oh * ow * self.anchors_len,
|
||||
_YOLO_V2_CLASSES + class_count))
|
||||
|
||||
# Threshold all the scores
|
||||
score_indices = sigmoid(colum_outputs[:, _YOLO_V2_SCORE])
|
||||
score_indices = np.nonzero(score_indices > self.threshold)
|
||||
if isinstance(score_indices, tuple):
|
||||
score_indices = score_indices[0]
|
||||
score_indices = sigmoid(row_outputs[:, _YOLO_V2_SCORE])
|
||||
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(colum_outputs, score_indices, axis=0)
|
||||
bb = np.take(row_outputs, score_indices, axis=0)
|
||||
|
||||
# Extract rows, columns, and anchor indices
|
||||
bb_rows = score_indices // (ow * self.anchors_len)
|
||||
@ -142,7 +142,7 @@ class yolo_v2_postprocess:
|
||||
h_rel = h_rel * ih
|
||||
|
||||
nms = NMS(iw, ih, inputs[0].roi)
|
||||
for i in range(len(bb)):
|
||||
for i in range(bb.shape[0]):
|
||||
nms.add_bounding_box(x_center[i] - (w_rel[i] / 2),
|
||||
y_center[i] - (h_rel[i] / 2),
|
||||
x_center[i] + (w_rel[i] / 2),
|
||||
@ -151,6 +151,19 @@ class yolo_v2_postprocess:
|
||||
return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)
|
||||
|
||||
|
||||
# This is a lightweight version of the YOLO (You Only Look Once) object detection algorithm.
|
||||
# It is designed to work well on embedded devices with limited computational resources.
|
||||
class yolo_lc_postprocess(yolo_v2_postprocess):
|
||||
def __init__(self, threshold=0.6, anchors=None, nms_threshold=0.1, nms_sigma=0.1):
|
||||
if anchors is None:
|
||||
anchors = np.array([[0.076023, 0.258508],
|
||||
[0.163031, 0.413531],
|
||||
[0.234769, 0.702585],
|
||||
[0.427054, 0.715892],
|
||||
[0.748154, 0.857092]])
|
||||
super().__init__(threshold, anchors, nms_threshold, nms_sigma)
|
||||
|
||||
|
||||
class yolo_v5_postprocess:
|
||||
_YOLO_V5_CX = const(0)
|
||||
_YOLO_V5_CY = const(1)
|
||||
@ -169,18 +182,16 @@ class yolo_v5_postprocess:
|
||||
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))
|
||||
row_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]
|
||||
score_indices = row_outputs[:, _YOLO_V5_SCORE]
|
||||
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(colum_outputs, score_indices, axis=0)
|
||||
bb = np.take(row_outputs, score_indices, axis=0)
|
||||
|
||||
# Get the score information
|
||||
bb_scores = bb[:, _YOLO_V5_SCORE]
|
||||
@ -202,7 +213,61 @@ class yolo_v5_postprocess:
|
||||
ymax = (y_center + h_rel) * ih
|
||||
|
||||
nms = NMS(iw, ih, inputs[0].roi)
|
||||
for i in range(len(bb)):
|
||||
for i in range(bb.shape[0]):
|
||||
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)
|
||||
|
||||
|
||||
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