visualization looks off for the correspondence task

[yuefuturew]

I tried to add a function to visualize the correspondence map just like the paper, with different colors showing different corresponded parts. However, I got this for the ground truth VTS:

I don't think the GT data should look like this, but I'm not able to identify the issue with the visualization code either. Here is what I did:

1. I let the test function return a pair of shape + their gt and pred correspondence:

return {
        'mean_loss': mean_loss,
        'mean_geodesic_error': mean_geodesic_error,
        'verts1': verts1_orig.detach().cpu().numpy(),  # Source mesh (shape [N1, 3])
        'faces1': faces1.detach().cpu().numpy(),
        'verts2': verts2.detach().cpu().numpy(),   # Target mesh (shape [N2, 3])
        'faces2': faces2.detach().cpu().numpy(),
        'pred_vts2on1': vts2on1.detach().cpu().numpy(), # Predicted: Indices in verts2 for each verts1 point
        'gt_vts1': vts1.detach().cpu().numpy()          # Ground truth: Indices in verts1 (if applicable)
    }

2. Called with GT:

    results['verts1'], results['faces1'],
    results['verts2'], results['faces2'],
    results['gt_vts1'], n_clusters=8)

3. Here is the visualization function:

def visualize_clustered_correspondences(
    verts1, faces1, verts2, faces2, pred_vts2on1,
    n_clusters=10, opacity=0.8
):
    def to_polydata(verts, faces):
        vtk_faces = np.hstack([np.full((faces.shape[0],1),3), faces]).ravel()
        return pv.PolyData(verts, vtk_faces)
    
    mesh1 = to_polydata(verts1, faces1)
    mesh2 = to_polydata(verts2, faces2)

    kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(verts1)
    labels1 = kmeans.labels_.astype(np.int32)

    mesh1['part'] = labels1
    # to handle different verts length
    mesh2['part'] = np.concatenate([labels1[pred_vts2on1],np.full(verts2.shape[0] - labels1.shape[0], -1, dtype=np.int32)])

    z_offset = 1.5 * max(mesh1.length, mesh2.length)  # 150% of largest dimension
    mesh2 = mesh2.translate([0, 0, -z_offset])

    p = pv.Plotter()
    p.enable_depth_peeling()

    for mesh in (mesh1, mesh2):
        mesh_colored = mesh.color_labels(
            scalars='part',
            colors='tab20',
            coloring_mode='index',
            negative_indexing=True,
            inplace=False
        )
        p.add_mesh(mesh_colored, rgb=True, lighting=False, opacity=0.8)

    p.show()

I would really appreciate if anyone can tell me what I did wrong.