def parse_predictions()

in wypr/evaluation/ap_helper.py [0:0]


def parse_predictions(
        end_points, num_class, remove_empty_box=False,
        use_3d_nms=True, cls_nms=True,
        nms_iou=0.25, use_old_type_nms=False,
        conf_thresh=0):
    """ Parse predictions to OBB parameters and suppress overlapping boxes """
    proposals = end_points['proposals'].cpu().numpy()
    bsize = proposals.shape[0]
    box_label_mask = end_points['proposal_mask']
    roi_scores = [_logits.detach().cpu().numpy() for _logits in end_points['roi_logits']]
    pred_sem_cls = [np.argmax(_score, -1) for _score in roi_scores]
    nonempty_rois = end_points['non_empty_roi_idx']
    nonempty_box_mask = [np.ones(len(nonempty_rois[_i])) for _i in range(bsize)]

    pred_corners_3d = [np.ones((len(_rois), 8, 3)) for _rois in nonempty_rois]
    for i in range(bsize):
        K_i = nonempty_rois[i].shape[0]
        for j in range(K_i):
            pred_corners_3d[i][j] = get_3d_box(proposals[i, j])

    # Remove predicted boxes without any point within them..
    if remove_empty_box:
        batch_pc = end_points['point_clouds'].cpu().numpy()[:,:,0:3] # B,N,3
        for i in range(bsize):
            pc = batch_pc[i,:,:] # (N,3)
            K_i = nonempty_rois[i].shape[0]
            for j in range(K_i):
                pc_in_box, inds = extract_pc_in_box3d(pc, pred_corners_3d[i][j])
                if len(pc_in_box) < 5:
                    nonempty_box_mask[i][j] = 0

    # NMS
    if use_3d_nms and not cls_nms:
        pred_mask = [np.zeros(len(nonempty_rois[_i])) for _i in range(bsize)]
        for i in range(bsize):
            K_i = nonempty_rois[i].shape[0]
            boxes_3d_with_prob = np.zeros((K_i, 7))
            for j in range(K_i):
                boxes_3d_with_prob[j,0] = np.min(pred_corners_3d[i][j,:,0])
                boxes_3d_with_prob[j,1] = np.min(pred_corners_3d[i][j,:,1])
                boxes_3d_with_prob[j,2] = np.min(pred_corners_3d[i][j,:,2])
                boxes_3d_with_prob[j,3] = np.max(pred_corners_3d[i][j,:,0])
                boxes_3d_with_prob[j,4] = np.max(pred_corners_3d[i][j,:,1])
                boxes_3d_with_prob[j,5] = np.max(pred_corners_3d[i][j,:,2])
                boxes_3d_with_prob[j,6] = roi_scores[i][j, pred_sem_cls[i][j]]
            nonempty_box_inds = np.where(nonempty_box_mask[i] == 1)[0]
            pick = nms_3d_faster(boxes_3d_with_prob[nonempty_box_mask[i]==1],
                nms_iou, use_old_type_nms)
            assert(len(pick)>0)
            pred_mask[i][nonempty_box_inds[pick]] = 1
    elif use_3d_nms and cls_nms:
        pred_mask = [np.zeros(len(nonempty_rois[_i])) for _i in range(bsize)]
        for i in range(bsize):
            K_i = nonempty_rois[i].shape[0]
            boxes_3d_with_prob = np.zeros((K_i, 8))
            for j in range(K_i):
                boxes_3d_with_prob[j,0] = np.min(pred_corners_3d[i][j,:,0])
                boxes_3d_with_prob[j,1] = np.min(pred_corners_3d[i][j,:,1])
                boxes_3d_with_prob[j,2] = np.min(pred_corners_3d[i][j,:,2])
                boxes_3d_with_prob[j,3] = np.max(pred_corners_3d[i][j,:,0])
                boxes_3d_with_prob[j,4] = np.max(pred_corners_3d[i][j,:,1])
                boxes_3d_with_prob[j,5] = np.max(pred_corners_3d[i][j,:,2])
                boxes_3d_with_prob[j,6] = roi_scores[i][j, pred_sem_cls[i][j]]
                boxes_3d_with_prob[j,7] = pred_sem_cls[i][j] 
            nonempty_box_inds = np.where(nonempty_box_mask[i] == 1)[0]
            pick = nms_3d_faster_samecls(
                boxes_3d_with_prob[nonempty_box_mask[i]==1],
                nms_iou, use_old_type_nms)
            assert(len(pick)>0)
            pred_mask[i][nonempty_box_inds[pick]] = 1
    else:
        pred_mask = [np.ones(len(nonempty_rois[_i])) for _i in range(bsize)]

    batch_props = []
    for i in range(bsize):
        num_rois_i = nonempty_rois[i].shape[0] 
        cur_list = []
        for j in range(num_rois_i):
            c = pred_sem_cls[i][j]
            if pred_mask[i][j] == 1 and roi_scores[i][j, c] > conf_thresh and c != num_class:
                cur_list += [(c, proposals[i, j], roi_scores[i][j, c])]
        batch_props.append(cur_list)
    return batch_props