added full pipeline

main_pipeline.py

@@ -20,10 +20,16 @@ from ml_lib.utils.model_io import SavedLightningModels
 # Datasets
 from datasets.shapenet import ShapeNetPartSegDataset
 from models import PointNet2
-from utils.pointcloud import cluster_cubes, append_onehotencoded_type, label2color
+from utils.pointcloud import cluster_cubes, append_onehotencoded_type, label2color, hierarchical_clustering, \
-from utils.project_settings import GlobalVar
+    write_clusters, cluster2Color, cluster_dbscan
+from utils.project_settings import GlobalVar, DataClass
 
 
+class DisplayMode(DataClass):
+    Clusters = 0,
+    Types = 1,
+    Nothing = 2
+
 def restore_logger_and_model(log_dir):
     model = SavedLightningModels.load_checkpoint(models_root_path=log_dir, model=PointNet2, n=-1)
     model = model.restore()
@@ -54,39 +60,89 @@ def predict_prim_type(input_pc, model):
 
 if __name__ == '__main__':
 
-    # input_pc_path = Path('data') / 'pc' / 'test.xyz'
+    display_mode = DisplayMode.Clusters
+    grid_cluster_max_pts = 3000#8192
+    grid_clusters = [3,3,3]
+    type_cluster_eps = 0.1
+    type_cluster_min_pts = 50
 
-    model_path = Path('output') / 'PN2' / 'PN_14628b734c5b651b013ad9e36c406934' / 'version_0'
+    model_path = Path('output') / 'PN2' / 'PN_9843bf499399786cfd58fe79fa1b3db8' / 'version_0'
-    # config_filename = 'config.ini'
-    # config = ThisConfig()
-    # config.read_file((Path(model_path) / config_filename).open('r'))
     loaded_model = restore_logger_and_model(model_path)
     loaded_model.eval()
 
-     #input_pc = read_pointcloud(input_pc_path, ' ', False)
-
-    # input_pc = normalize_pointcloud(input_pc)
-
-    # TEST DATASET
     transforms = Compose([NormalizeScale(), ])
     test_dataset = ShapeNetPartSegDataset('data', mode=GlobalVar.data_split.predict, collate_per_segment=False,
                                           refresh=True, transform=transforms)
 
-    grid_clusters = cluster_cubes(test_dataset[0], [1, 1, 1], max_points_per_cluster=8192)
+    grid_clusters = cluster_cubes(test_dataset[0], grid_clusters, max_points_per_cluster=grid_cluster_max_pts)
 
     ps.init()
 
+    # ========================== Grid Clustering ==========================
+
+    grid_cluster_pcs = []
+
     for i, grid_cluster_pc in enumerate(grid_clusters):
 
         print("Cluster pointcloud size: {}".format(grid_cluster_pc.shape[0]))
 
         pc_with_prim_type = predict_prim_type(grid_cluster_pc, loaded_model)
 
-        # pc_with_prim_type = polytopes_to_planes(pc_with_prim_type)
+        # Re-Map Primitive type for 1-hot-encoding.
+        pc_with_prim_type[:, 6][pc_with_prim_type[:, 6] == 0.0] = 0.0 # Sphere
+        pc_with_prim_type[:, 6][pc_with_prim_type[:, 6] == 1.0] = 1.0 # Cylinder
+        pc_with_prim_type[:, 6][pc_with_prim_type[:, 6] == 3.0] = 2.0  # Box
+        pc_with_prim_type[:, 6][pc_with_prim_type[:, 6] == 4.0] = 2.0  # Polytope
+        pc_with_prim_type[:, 6][pc_with_prim_type[:, 6] == 6.0] = 2.0  # Plane
+
         pc_with_prim_type = append_onehotencoded_type(pc_with_prim_type)
 
-        pc = ps.register_point_cloud("points_" + str(i), pc_with_prim_type[:, :3], radius=0.01)
+        grid_cluster_pcs.append(pc_with_prim_type)
-        pc.add_color_quantity("prim types", label2color(pc_with_prim_type[:, 6].astype(np.int64)), True)
+
+    # Merge grid cluster pcs together
+    final_pc = np.concatenate(grid_cluster_pcs)
+
+    # ========================== DBSCAN Clustering ==========================
+
+    print("DB Scan on point cloud " + str(final_pc.shape))
+    total_clusters = []
+
+    clusters = cluster_dbscan(final_pc, [0, 1, 2, 3, 4, 5], eps=type_cluster_eps,
+                                 min_samples=type_cluster_min_pts)
+    print("Pre-clustering done. Clusters: ", len(clusters))
+
+    for cluster in clusters:
+
+        print("2nd level clustering ..")
+
+        prim_types_in_cluster = len(np.unique(cluster[:, 6], axis=0))
+        if prim_types_in_cluster == 1:
+            print("No need for 2nd level clustering since there is only a single primitive type in the cluster.")
+            total_clusters.append(cluster)
+        else:
+            sub_clusters = cluster_dbscan(cluster, [0, 1, 2, 7, 8, 9], eps=type_cluster_eps,
+                                             min_samples=type_cluster_min_pts)
+            print("Sub clusters: ", len(sub_clusters))
+            total_clusters.extend(sub_clusters)
+
+    result_clusters = list(filter(lambda c: c.shape[0] > type_cluster_min_pts, total_clusters))
+
+    for cluster in result_clusters:
+        print("Cluster: ", cluster.shape[0])
+
+    write_clusters('clusters.txt', result_clusters, 6)
+
+    # ========================== Result visualization ==========================
+    if display_mode == DisplayMode.Types:
+
+        pc = ps.register_point_cloud("points_" + str(i), final_pc[:, :3], radius=0.01)
+        pc.add_color_quantity("prim types", label2color(final_pc[:, 6].astype(np.int64)), True)
+
+    elif display_mode == DisplayMode.Clusters:
+
+        for i, result_cluster in enumerate(result_clusters):
+            pc = ps.register_point_cloud("points_" + str(i), result_cluster[:, :3], radius=0.01)
+            pc.add_color_quantity("prim types", cluster2Color(result_cluster,i), True)
 
     ps.show()
 

utils/pointcloud.py

@@ -1,7 +1,7 @@
 import numpy as np
 from sklearn.cluster import DBSCAN
 
-import open3d as o3d
+#import open3d as o3d
 
 from pyod.models.lof import LOF
 
@@ -52,6 +52,8 @@ def label2color(labels):
     num = labels.shape[0]
     colors = np.zeros((num, 3))
 
+    print(labels)
+
     minl, maxl = np.min(labels), np.max(labels)
     for l in range(minl, maxl + 1):
         colors[labels == l, :] = mini_color_table(l)
@@ -189,19 +191,6 @@ def cluster_dbscan(data, selected_indices, eps, min_samples=5, metric='euclidean
 
     return npClusters
 
-
-def draw_clusters(clusters):
-    clouds = []
-
-    for cluster_idx, cluster in enumerate(clusters):
-        cloud = o3d.PointCloud()
-        cloud.points = o3d.Vector3dVector(cluster[:, :3])
-        cloud.colors = o3d.Vector3dVector(cluster2Color(cluster, cluster_idx))
-        clouds.append(cloud)
-
-    o3d.draw_geometries(clouds)
-
-
 def write_clusters(path, clusters, type_column=6):
     file = open(path, "w")
     file.write(str(len(clusters)) + "\n")
@@ -216,7 +205,7 @@ def write_clusters(path, clusters, type_column=6):
                 return 2
             elif t == 1:
                 return 1
-            elif t == 3:
+            elif t == 2:
                 return 4
             return t
 
@@ -228,15 +217,6 @@ def write_clusters(path, clusters, type_column=6):
         np.savetxt(file, cluster[:, :6], header=str(len(cluster)) + ' ' + str(6), comments='')
 
 
-def draw_sample_data(sample_data, colored_normals=False):
-    cloud = o3d.PointCloud()
-    cloud.points = o3d.Vector3dVector(sample_data[:, :3])
-    cloud.colors = \
-        o3d.Vector3dVector(label2color(sample_data[:, 6].astype(int)) if not colored_normals else sample_data[:, 3:6])
-
-    o3d.draw_geometries([cloud])
-
-
 def normalize_pointcloud(pc, factor=1.0):
     max = pc.max(axis=0)
     min = pc.min(axis=0)
@@ -284,7 +264,7 @@ def split_outliers(pc, columns):
 
 def append_onehotencoded_type(data, factor=1.0):
     types = data[:, 6].astype(int)
-    res = np.zeros((len(types), 8))
+    res = np.zeros((len(types), 3))
     res[np.arange(len(types)), types] = factor
 
     return np.column_stack((data, res))