DataSet Modifications & Checks

datasets/shapenet.py

@@ -94,6 +94,30 @@ class CustomShapeNet(InMemoryDataset):
     def num_classes(self):
         return len(self.categories) if self.poly_as_plane else (len(self.categories) - 2)
 
+    @property
+    def class_map_all(self):
+        return {0: 0,
+                1: 1,
+                2: None,
+                3: 2,
+                4: 3,
+                5: None,
+                6: 4,
+                7: None
+                }
+
+    @property
+    def class_map_poly_as_plane(self):
+        return {0: 0,
+                1: 1,
+                2: None,
+                3: 2,
+                4: 2,
+                5: None,
+                6: 2,
+                7: None
+                }
+
     def _load_dataset(self):
         data, slices = None, None
         filepath = self.processed_paths[0]
@@ -154,6 +178,7 @@ class CustomShapeNet(InMemoryDataset):
         datasets = defaultdict(list)
         path_to_clouds = self.raw_dir / self.mode
         found_clouds = list(path_to_clouds.glob('*.xyz'))
+        class_map = self.class_map_all if not self.poly_as_plane else self.class_map_poly_as_plane
         if len(found_clouds):
             for pointcloud in tqdm(found_clouds):
                 if self.cluster_type not in pointcloud.name:
@@ -171,7 +196,7 @@ class CustomShapeNet(InMemoryDataset):
                                 raise ValueError('Check the Input!!!!!!')
                             # Expand the values from the csv by fake labels if non are provided.
                             vals = vals + [0] * (8 - len(vals))
-
+                            vals[-2] = float(class_map[int(vals[-2])])
                             src[vals[-1]].append(vals)
 
                 # Switch from un-pickable Defaultdict to Standard Dict
@@ -193,10 +218,6 @@ class CustomShapeNet(InMemoryDataset):
                         if all([x == 0 for x in tensor]):
                             continue
                         tensor = tensor.unsqueeze(0)
-                    if self.poly_as_plane:
-                        tensor[:, -2][tensor[:, -2] == float(self.classes.Plane)] = 2.0
-                        tensor[:, -2][tensor[:, -2] == float(self.classes.Box)] = 2.0
-                        tensor[:, -2][tensor[:, -2] == float(self.classes.Polytope)] = 2.0
                     src[key] = tensor
 
                 for key, values in src.items():