def nms(dets, thresh): """Dispatch to either CPU or GPU NMS implementations. used in both inference (keep_out is 1D) and 'proposal_layer' (keep_out is 2D) Args: dets: [bs, N, 4] thresh: nms threshold Returns: keep_out: ndarray (bs, min_keep_num) """ bs = dets.size(0) keep = [] min_keep_num = 1e10 for i in range(bs): curr_sample_keep = pth_nms(dets[i, :], thresh) keep.append(curr_sample_keep) if len(curr_sample_keep) < min_keep_num: min_keep_num = len(curr_sample_keep) keep_out = np.zeros((bs, min_keep_num), dtype=np.int32) for i in range(bs): keep_out[i, :] = keep[i][:min_keep_num].cpu().numpy() return keep_out
def nms(dets, thresh): "Dispatch to either CPU or GPU NMS implementations.\ Accept dets as tensor" "" return pth_nms(dets, thresh)
def nms(dets, thresh): """Dispatch to either CPU or GPU NMS implementations. Accept dets as tensor""" return pth_nms(dets, thresh)
def box_nms(bboxes, scores, threshold=0.5, mode='union'): dets = torch.cat([bboxes, scores.unsqueeze(1)], dim=1) return pth_nms(dets, threshold)