Unable to replicate performance on DenseCorr3D

[factoryofthesun]

Hi thank you for the work and code! I have attempted to reproduce the evaluation results on the DenseCorr3D benchmark using the code provided in example.ipynb. I use the predicted surface_map, surface_map_inv mappings and compute normalized semantic geodesic error and AUC w/ 10% threshold with the below code. However, the values I get with this are geodesic error 10.1, auc 0.50 for total and geodesic error 6.84, auc 0.52 for holdout. I do not perform any normalization other than centering bounding box to the origin, as in the example notebook. Could you please advise?

def calculate_semantic_geodesic_error(dist_x, dist_y, groups_x, groups_y, p2p_x, p2p_y, return_mean=True):
    """
    Calculate the geodesic error between predicted correspondence and gt correspondence

    Args:
        dist_x (np.ndarray): Geodesic distance matrix of shape x. shape [Vx, Vx]
        dist_y (np.ndarray): Geodesic distance matrix of shape y. shape [Vy, Vy]
        groups_x (list): Group indexing into shape x
        groups_y (list): Group indexing into shape y
        p2p_x (np.ndarray): Point-to-point map (shape y -> shape x). shape [Vy]
        p2p_y (np.ndarray): Point-to-point map (shape x -> shape y). shape [Vx]
        return_mean (bool, optional): Average the geodesic error. Default True.
    Returns:
        avg_geodesic_error (np.ndarray): Average geodesic error.
    """
    geo_errs_x = [] # Error for y->x map
    geo_errs_y = [] # Error for x->y map
    # Loop over groups to get the geodesic errors
    for groupi, gt_y in enumerate(groups_y):
        gt_x = groups_x[groupi]

        if len(gt_x) == 0 or len(gt_y) == 0:
            continue

        pred_x = p2p_y[gt_y]
        gt_x_dists = dist_x[:, gt_x]
        group_err_x = np.min(gt_x_dists[pred_x], axis=1)
        geo_errs_x.append(group_err_x)

        pred_y = p2p_x[gt_x]
        gt_y_dists = dist_y[:, gt_y]
        group_err_y = np.min(gt_y_dists[pred_y], axis=1)
        geo_errs_y.append(group_err_y)

    geo_errs_x = np.concatenate(geo_errs_x)
    geo_errs_y = np.concatenate(geo_errs_y)

    if return_mean:
        return geo_errs_x.mean(), geo_errs_y.mean()
    else:
        return geo_errs_x, geo_errs_y

def plot_pck(geo_err, threshold=0.10, steps=40):
    """
    plot pck curve and compute auc.
    Args:
        geo_err (np.ndarray): geodesic error list.
        threshold (float, optional): threshold upper bound. Default 0.15.
        steps (int, optional): number of steps between [0, threshold]. Default 30.
    Returns:
        auc (float): area under curve.
        fig (matplotlib.pyplot.figure): pck curve.
        pcks (np.ndarray): pcks.
    """
    assert threshold > 0 and steps > 0
    geo_err = np.ravel(geo_err)
    thresholds = np.linspace(0., threshold, steps)
    pcks = []
    for i in range(thresholds.shape[0]):
        thres = thresholds[i]
        pck = np.mean((geo_err <= thres).astype(float))
        pcks.append(pck)
    pcks = np.array(pcks)
    # compute auc
    auc = np.trapezoid(pcks, np.linspace(0., 1., steps))

    return auc, pcks