diff --git a/scripts/libraries/ml/ml/postprocessing.py b/scripts/libraries/ml/ml/postprocessing.py index 78ae062c7..c63cee0a9 100644 --- a/scripts/libraries/ml/ml/postprocessing.py +++ b/scripts/libraries/ml/ml/postprocessing.py @@ -69,19 +69,17 @@ class yolo_v2_postprocess: _YOLO_V2_SCORE = const(4) _YOLO_V2_CLASSES = const(5) - def __init__(self, score_threshold=0.6, anchors=None, - nms_threshold=0.1, nms_sigma=0.1): + def __init__(self, score_threshold=0.6, anchors=None): self.score_threshold = score_threshold - self.anchors = anchors - if self.anchors is None: + if anchors is not None: + self.anchors = anchors + else: self.anchors = np.array([[0.98830, 3.36060], [2.11940, 5.37590], [3.05200, 9.13360], [5.55170, 9.30660], - [9.72600, 11.1422]]) + [9.72600, 11.1422]], dtype=np.float) 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] @@ -94,39 +92,42 @@ class yolo_v2_postprocess: return a - (b * (a // b)) def softmax(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)) + e_x = np.exp(x - np.max(x)) + return e_x / np.sum(e_x) # 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 - 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): + 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): return [] # Get the bounding boxes that have a valid score - bb = np.take(colum_outputs, valid_score_indices, axis=0) + bb = np.take(colum_outputs, score_indices, axis=0) # Extract rows, columns, and anchor indices - 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) + 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) # Get the anchor box information - bb_a_array = np.take(self.anchors, bb_anchors, axis=0) + bb_a_array = [self.anchors[i] for i in bb_anchors.tolist()] + bb_a_array = np.array(bb_a_array) # Get the score information bb_scores = sigmoid(bb[:, _YOLO_V2_SCORE]) # Get the class information - bb_classes = np.argmax(softmax(bb[:, _YOLO_V2_CLASSES:]), axis=1) + 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) # Compute the bounding box information x_center = (bb_cols + sigmoid(bb[:, _YOLO_V2_TX])) / ow @@ -148,4 +149,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(self.nms_threshold, self.nms_sigma) + return nms.get_bounding_boxes()