def multiclass_nms(prediction, class_num, keep_top_k, nms_threshold):
selected_indices = {}
num_det = 0
for c in range(class_num):
if c not in prediction.keys():
continue
cls_dets = prediction[c]
all_scores = np.zeros(len(cls_dets))
for i in range(all_scores.shape[0]):
all_scores[i] = cls_dets[i][4]
indices = nms(cls_dets, all_scores, 0.0, nms_threshold, -1, False, 1.0)
selected_indices[c] = indices
num_det += len(indices)
score_index = []
for c, indices in selected_indices.items():
for idx in indices:
score_index.append((prediction[c][idx][4], c, idx))
sorted_score_index = sorted(
score_index, key=lambda tup: tup[0], reverse=True)
if keep_top_k > -1 and num_det > keep_top_k:
sorted_score_index = sorted_score_index[:keep_top_k]
num_det = keep_top_k
nmsed_outs = []
for s, c, idx in sorted_score_index:
xmin = prediction[c][idx][0]
ymin = prediction[c][idx][1]
xmax = prediction[c][idx][2]
ymax = prediction[c][idx][3]
nmsed_outs.append([c + 1, s, xmin, ymin, xmax, ymax])
return nmsed_outs, num_det
def proposal_for_one_image(im_info, all_anchors, variances, bbox_deltas,
scores, pre_nms_topN, post_nms_topN, nms_thresh,
min_size, eta):
# Transpose and reshape predicted bbox transformations to get them
# into the same order as the anchors:
# - bbox deltas will be (4 * A, H, W) format from conv output
# - transpose to (H, W, 4 * A)
# - reshape to (H * W * A, 4) where rows are ordered by (H, W, A)
# in slowest to fastest order to match the enumerated anchors
bbox_deltas = bbox_deltas.transpose((1, 2, 0)).reshape(-1, 4)
all_anchors = all_anchors.reshape(-1, 4)
variances = variances.reshape(-1, 4)
# Same story for the scores:
# - scores are (A, H, W) format from conv output
# - transpose to (H, W, A)
# - reshape to (H * W * A, 1) where rows are ordered by (H, W, A)
# to match the order of anchors and bbox_deltas
scores = scores.transpose((1, 2, 0)).reshape(-1, 1)
# sort all (proposal, score) pairs by score from highest to lowest
# take top pre_nms_topN (e.g. 6000)
if pre_nms_topN <= 0 or pre_nms_topN >= len(scores):
order = np.argsort(-scores.squeeze())
else:
# Avoid sorting possibly large arrays;
# First partition to get top K unsorted
# and then sort just those
inds = np.argpartition(-scores.squeeze(), pre_nms_topN)[:pre_nms_topN]
order = np.argsort(-scores[inds].squeeze())
order = inds[order]
scores = scores[order, :]
bbox_deltas = bbox_deltas[order, :]
all_anchors = all_anchors[order, :]
proposals = box_coder(all_anchors, bbox_deltas, variances)
# clip proposals to image (may result in proposals with zero area
# that will be removed in the next step)
proposals = clip_tiled_boxes(proposals, im_info[:2])
# remove predicted boxes with height or width < min_size
keep = filter_boxes(proposals, min_size, im_info)
if len(keep) == 0:
proposals = np.zeros((1, 4)).astype('float32')
scores = np.zeros((1, 1)).astype('float32')
return proposals, scores
proposals = proposals[keep, :]
scores = scores[keep, :]
# apply loose nms (e.g. threshold = 0.7)
# take post_nms_topN (e.g. 1000)
# return the top proposals
if nms_thresh > 0:
keep = nms(boxes=proposals,
scores=scores,
nms_threshold=nms_thresh,
eta=eta)
if post_nms_topN > 0 and post_nms_topN < len(keep):
keep = keep[:post_nms_topN]
proposals = proposals[keep, :]
scores = scores[keep, :]
return proposals, scores