Cannot generate an instance segmentation point cloud using easy setup

[Flora-Liao]

Hello, I'm trying to use the easy setup to generate an instance segmentation result for a point cloud(ScannetV2 scene0011_00)
I use the code provided in the easy setup example and change the path of the checkpoint to the pre-trained checkpoint fo scannet200
The following is my code:

import torch
from mask3d import get_model, load_mesh, prepare_data, map_output_to_pointcloud, save_colorized_mesh 
from omegaconf import OmegaConf, DictConfig
def main():
    #cfg = OmegaConf.load('mask3d/conf/config_base_instance_segmentation.yaml')
    #print(OmegaConf.to_yaml(cfg))
    # Load the pre-trained model
    #model = get_model('checkpoint/scannet200/scannet200_benchmark.ckpt')
    model = get_model('checkpoint/scannet200/scannet200_benchmark.ckpt')
    print("Finish get model")
    print(" ")
    print("Start eval")
    model.eval()
    
    # Set the device to GPU if available, otherwise CPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    
    # Load the input point cloud data
    pointcloud_file = 'data/scans/scene0011_00/scene0011_00_vh_clean.ply'
    mesh = load_mesh(pointcloud_file)
    
    # Prepare the data for the model
    data, points, colors, features, unique_map, inverse_map = prepare_data(mesh, device)
    
    # Run the model with the prepared data
    with torch.no_grad():
        outputs = model(data, raw_coordinates=features)
    #print("Output:")
    #print(outputs)
    # Map the model output to the point cloud
    labels = map_output_to_pointcloud(mesh, outputs, inverse_map)
    #print("labels")
    #print(labels)
    # Save the colorized mesh with the labels
    save_colorized_mesh(mesh, labels, 'data/pcl_result.ply', colormap='scannet200')

if __name__ == "__main__":
    main()

However, when I use open3d to visualize my pcl_result.ply, it seems like a semantic segmentation result!

Is there any method to change my point cloud result to an instance segmentation result?
Thank you very much!

[Flora-Liao]

I successfully solved the issue by modifying the map_output_to_pointcloud and save_colorized_mesh!

[Andyshen555]

Hi Liao, can you please provide details on how you have modified the two functions?

[Flora-Liao]

Hello,
I modified the code in "mask3d/init.py" as following:

def map_output_to_pointcloud(mesh, 
                             outputs, 
                             inverse_map, 
                             label_space='scannet200',
                             confidence_threshold=0.9):
    
    # parse predictions
    logits = outputs["pred_logits"]
    masks = outputs["pred_masks"]

    # reformat predictions
    logits = logits[0].detach().cpu()
    masks = masks[0].detach().cpu()

    labels = []
    confidences = []
    instance_masks = []

    for i in range(len(logits)):
        p_labels = torch.softmax(logits[i], dim=-1)
        p_masks = torch.sigmoid(masks[:, i])
        l = torch.argmax(p_labels, dim=-1)
        c_label = torch.max(p_labels)
        m = p_masks > 0.5
        c_m = p_masks[m].sum() / (m.sum() + 1e-8)
        c = c_label * c_m
        if l < 200 and c > confidence_threshold:
            labels.append(l.item())
            confidences.append(c.item())
            instance_masks.append(m[inverse_map])  # mapping the mask back to the original point cloud
    
    # save labelled mesh
    mesh_labelled = o3d.geometry.TriangleMesh()
    mesh_labelled.vertices = mesh.vertices
    mesh_labelled.triangles = mesh.triangles

    labels_mapped = np.zeros((len(mesh.vertices), 1))

    instance_id = 1  # Starting instance ID

    for i, (l, c, m) in enumerate(
        sorted(zip(labels, confidences, instance_masks), reverse=False)):
        
        if label_space == 'scannet200':
            label_offset = 2
            if l == 0:
                l = -1 + label_offset
            else:
                l = int(l) + label_offset
                        
        labels_mapped[m == 1] = instance_id
        instance_id += 1
        
    return labels_mapped

def save_colorized_mesh(mesh, labels_mapped, output_file, colormap='scannet'):
    
    # Generate unique colors for each instance
    unique_labels = np.unique(labels_mapped)
    np.random.seed(42)  # For reproducibility
    colors = np.random.rand(len(unique_labels), 3)
    
    # Map colors to the mesh
    vertex_colors = np.zeros((len(mesh.vertices), 3))
    for i, label in enumerate(unique_labels):
        if label != 0:  # Skip background
            vertex_colors[labels_mapped[:, 0] == label] = colors[i]
    
    mesh.vertex_colors = o3d.utility.Vector3dVector(vertex_colors)
    o3d.io.write_triangle_mesh(output_file, mesh)