scripts/libraries: Fully vectorize and cleanup yolo_v2 post-processing.

This commit is contained in:
Kwabena W. Agyeman 2024-11-24 14:29:48 -08:00
parent 37ec53ca8e
commit 5394a3d32f

View File

@ -69,17 +69,19 @@ class yolo_v2_postprocess:
_YOLO_V2_SCORE = const(4)
_YOLO_V2_CLASSES = const(5)
def __init__(self, score_threshold=0.6, anchors=None):
def __init__(self, score_threshold=0.6, anchors=None,
nms_threshold=0.1, nms_sigma=0.1):
self.score_threshold = score_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]
@ -92,42 +94,39 @@ 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)
max_x = np.max(x, axis=1)
e_x = np.exp(x - max_x.reshape((x.shape[0], 1)))
sum_e_x = np.sum(e_x, axis=1)
return e_x / sum_e_x.reshape((x.shape[0], 1))
# Reshape the output to a 2D array
colum_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.score_threshold)
if isinstance(score_indices, tuple):
score_indices = score_indices[0]
if not len(score_indices):
all_scores = sigmoid(colum_outputs[:, _YOLO_V2_SCORE])
valid_score_indices = np.nonzero(all_scores >= self.score_threshold)
if isinstance(valid_score_indices, tuple):
valid_score_indices = valid_score_indices[0]
if not len(valid_score_indices):
return []
# Get the bounding boxes that have a valid score
bb = np.take(colum_outputs, score_indices, axis=0)
bb = np.take(colum_outputs, valid_score_indices, axis=0)
# Extract rows, columns, and anchor indices
bb_rows = score_indices // (ow * self.anchors_len)
bb_cols = mod(score_indices // self.anchors_len, ow)
bb_anchors = mod(score_indices, self.anchors_len)
bb_rows = valid_score_indices // (ow * self.anchors_len)
bb_cols = mod(valid_score_indices // self.anchors_len, ow)
bb_anchors = mod(valid_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
@ -149,4 +148,4 @@ 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(self.nms_threshold, self.nms_sigma)