From ab319d83b569e4dd5bbfff4fa799c7a380cbb441 Mon Sep 17 00:00:00 2001 From: unknown Date: Tue, 25 May 2021 13:44:56 +0900 Subject: [PATCH] Fix number of classes in nms, Change for-loop to numpy calc --- core/ground_truth.py | 24 ++++++++---------------- utils/nms.py | 2 +- 2 files changed, 9 insertions(+), 17 deletions(-) diff --git a/core/ground_truth.py b/core/ground_truth.py index ba8a7dc..2d2357c 100644 --- a/core/ground_truth.py +++ b/core/ground_truth.py @@ -96,21 +96,15 @@ def __get_valid_boxes(boxes): def __label_positive_and_negative_predicted_boxes(self, box_true, box_pred): box_true_coord = box_true[..., :4] box_true_class = box_true[..., -1] - box_pred_assigned = np.zeros_like(box_pred, dtype=np.float32) iou_outside = [] for i in range(box_true_coord.shape[0]): - iou_inside = [] - for j in range(box_pred.shape[0]): - iou = IOU(box_1=box_true_coord[i], box_2=box_pred[j]).calculate_iou() - iou_inside.append(iou) + iou_inside = IOU(box_1=box_true_coord[i], box_2=box_pred).calculate_iou() iou_outside.append(iou_inside) iou_array = np.array(iou_outside, dtype=np.float32) # shape: (num_of_true_boxes, total_num_of_default_boxes) iou_max = np.max(iou_array, axis=0) max_index = np.argmax(iou_array, axis=0) - max_index_class = np.zeros_like(max_index, dtype=np.float32) - for k in range(max_index.shape[0]): - max_index_class[k] = box_true_class[max_index[k]] - box_pred_assigned[k] = self.__get_offset(box_true=box_true_coord[max_index[k]], box_pred=box_pred[k]) + max_index_class = box_true_class[max_index] + box_pred_assigned = self.__get_offset(box_true=box_true_coord[max_index], box_pred=box_pred) pos_boolean = np.where(iou_max > self.iou_threshold, 1.0, 0.0) # 1 for positive, 0 for negative pos_class_index = max_index_class * pos_boolean pos_class_index = pos_class_index.reshape((-1, 1)) @@ -119,13 +113,11 @@ def __label_positive_and_negative_predicted_boxes(self, box_true, box_pred): @staticmethod def __get_offset(box_true, box_pred): - d_cx, d_cy, d_w, d_h = box_pred - g_cx, g_cy, g_w, g_h = box_true - g_cx = (g_cx - d_cx) / d_w - g_cy = (g_cy - d_cy) / d_h - g_w = np.log(g_w / d_w) - g_h = np.log(g_h / d_h) - return np.stack([g_cx, g_cy, g_w, g_h], axis=0) + box_true[:, 0] = (box_true[:, 0] - box_pred[:, 0]) / box_pred[:, 2] + box_true[:, 1] = (box_true[:, 1] - box_pred[:, 1]) / box_pred[:, 3] + box_true[:, 2] = np.log(box_true[:, 2] / box_pred[:, 2]) + box_true[:, 3] = np.log(box_true[:, 3] / box_pred[:, 3]) + return box_true def generate_gt_boxes(self): true_boxes = self.___transform_true_boxes() # shape: (batch_size, MAX_BOXES_PER_IMAGE, 5) diff --git a/utils/nms.py b/utils/nms.py index 15c3d0c..8a63d3c 100644 --- a/utils/nms.py +++ b/utils/nms.py @@ -6,7 +6,7 @@ class NMS(object): def __init__(self): super(NMS, self).__init__() self.max_box_num = MAX_BOX_NUM - self.num_class = NUM_CLASSES + 1 + self.num_class = NUM_CLASSES def nms(self, boxes, box_scores): mask = box_scores >= CONFIDENCE_THRESHOLD