def _postprocess(self, width, height, confidences, boxes, prob_threshold, iou_threshold=0.3, top_k=-1): boxes = boxes[0] confidences = confidences[0] picked_box_probs = [] picked_labels = [] for class_index in range(1, confidences.shape[1]): probs = confidences[:, class_index] mask = probs > prob_threshold probs = probs[mask] if probs.shape[0] == 0: continue subset_boxes = boxes[mask, :] box_probs = np.concatenate( [subset_boxes, probs.reshape(-1, 1)], axis=1) box_probs = box_utils.hard_nms(box_probs, iou_threshold=iou_threshold, top_k=top_k, ) picked_box_probs.append(box_probs) picked_labels.extend([class_index] * box_probs.shape[0]) if not picked_box_probs: return np.array([]), np.array([]), np.array([]) picked_box_probs = np.concatenate(picked_box_probs) picked_box_probs[:, 0] *= width picked_box_probs[:, 1] *= height picked_box_probs[:, 2] *= width picked_box_probs[:, 3] *= height return picked_box_probs[:, :4].astype(np.int32), np.array(picked_labels), picked_box_probs[:, 4]
def predict(width, height, confidences, box, prob_threshold, iou_threshold=0.5, top_k=-1): box = box[0] confidences = confidences[0] picked_box_probs = [] picked_labels = [] for class_index in range(1, confidences.shape[1]): probs = confidences[:, class_index] mask = probs > prob_threshold probs = probs[mask] if probs.shape[0] == 0: continue subset_box = box[mask, :] box_probs = np.concatenate([subset_box, probs.reshape(-1, 1)], axis=1) box_probs = box_utils.hard_nms( box_probs, iou_threshold=iou_threshold, top_k=top_k, ) picked_box_probs.append(box_probs) picked_labels.extend([class_index] * box_probs.shape[0]) if not picked_box_probs: return np.array([]), np.array([]), np.array([]) picked_box_probs = np.concatenate(picked_box_probs) w_center, h_center = ( picked_box_probs[:, 0] + picked_box_probs[:, 2]) / 2. * width, ( picked_box_probs[:, 1] + picked_box_probs[:, 3]) / 2. * height # global w h = h_center[0] w = w_center[0] boxw1, boxw2 = 200, 200 boxh1, boxh2 = 250, 150 img_w, img_h = width, height if w > img_w - boxw2: w = img_w - boxw2 if w < boxw1: w = boxw1 if h > img_h - boxh2: h = img_h - boxh2 if h < boxh1: h = boxh1 h_center = h w_center = w picked_box_probs[:, 0] = w_center - boxw1 picked_box_probs[:, 1] = h_center - boxh1 picked_box_probs[:, 2] = w_center + boxw2 picked_box_probs[:, 3] = h_center + boxh2 return picked_box_probs[:, :4].astype( np.int32), np.array(picked_labels), picked_box_probs[:, 4]