Add group slice with DBSCAN clustering colors to a grouped dataset in FiftyOne

dbscan_cluster_fiftyone.py

import open3d as o3d
import matplotlib.pyplot as plt

import fiftyone.core.utils as fou

def add_cluster_samples(dataset, eps = 0.5, min_points = 500):
    cluster_samples = []
    
    dataset.group_slice = "pcd"
    filepaths, groups = dataset.values(["filepath", "group"])
    
    with fou.ProgressBar(total=len(filepaths)) as pb:
        for fp, group in pb(zip(filepaths, groups)):
    
            sample = dataset[fp]
            pcd = o3d.io.read_point_cloud(fp)
            labels = np.array(
              pcd.cluster_dbscan(
                eps=eps, 
                min_points=min_points
              )
            )
            max_label = labels.max()
            classes = np.arange(-1, max_label+1)
            cluster_clouds = [
                pcd.select_by_index(np.where(labels == c)[0])
                for c in classes
            ]

            for cc in cluster_clouds[1:]:
                colors = plt.get_cmap("tab20")(labels / (max_label if max_label > 0 else 1))
                colors[labels < 0] = 0
                pcd.colors = o3d.utility.Vector3dVector(colors[:, :3])

            new_fp = "dbscan_cluster/" + sample.filename.replace(".", "_dbscan.")
            o3d.io.write_point_cloud(new_fp, pcd)

            new_sample = fo.Sample(filepath=new_fp)
            new_sample["group"] = group.element("pcd_cluster")
            cluster_samples.append(new_sample)

    dataset.add_samples(cluster_samples)