def write_detection_results_nms(output_dir, gt_dirs, det_results, threshold=cfg.TEST.THRESH): nms_results = {} for idx in det_results: for class_type in det_results[idx]: dets = np.array(det_results[idx][class_type], dtype=np.float32) # scores = dets[:, -1] # keep = (scores > 0.001).nonzero()[0] # print(len(scores), len(keep)) # dets = dets[keep] if len(dets) > 1: dets_for_nms = dets[:, 4:][:, [0, 1, 2, 5, 4, 3, 6, 7]] keep = cube_nms_np(dets_for_nms, threshold) # print(len(dets_for_nms), len(keep)) dets_keep = dets[keep] else: dets_keep = dets if idx not in nms_results: nms_results[idx] = {} # if class_type not in nms_results[idx]: # nms_results[idx][class_type] = [] nms_results[idx][class_type] = dets_keep write_detection_results(output_dir, gt_dirs, nms_results)
def make_nms(det_results, threshold=cfg.TEST.THRESH): nms_results = {} for idx in det_results: for class_type in det_results[idx]: dets = np.array(det_results[idx][class_type], dtype=np.float32) if len(dets) > 1: dets_for_nms = dets[:, 4:][:, [0, 1, 2, 5, 4, 3, 6, 7]] keep = cube_nms_np(dets_for_nms, threshold) dets_keep = dets[keep] else: dets_keep = dets if idx not in nms_results: nms_results[idx] = {} nms_results[idx][class_type] = dets_keep return nms_results