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]
