Minor cleanup , slight path optimization

Author: harrydobbs

smart_tree/cli.py

@@ -29,7 +29,7 @@ def main(cfg: DictConfig):
             pipeline.process_cloud(Path(f"{cfg.directory}/{p}"))
 
     else:
-        print("Please Supply a path or Directory")
+        print("Please supply a path or directory to point clouds.")
 
 
 if __name__ == "__main__":

smart_tree/conf/tree-dataset.yaml

@@ -69,25 +69,25 @@ skeletonizer:
   K: 16
   min_connection_length: 0.02 
   minimum_graph_vertices: 32
-  max_number_components: 100
+  max_number_components: 8
   voxel_downsample: False
   edge_non_linear: None
 
 
 pipeline:
 
   preprocessing:
-    #  Scale:
-    #      _target_:  smart_tree.dataset.augmentations.Scale
-    #      min_scale: 0.1
-    #      max_scale: 0.1
+    # Scale:
+    #       _target_:  smart_tree.dataset.augmentations.Scale
+    #       min_scale: 0.6
+    #       max_scale: 0.6
     # Scale:
     #   _target_:  smart_tree.dataset.augmentations.Scale
     #   min_scale: 1
     #   max_scale: 1
-    VoxelDownsample:
-       _target_:  smart_tree.dataset.augmentations.VoxelDownsample
-       voxel_size : 0.01
+    # VoxelDownsample:
+    #     _target_:  smart_tree.dataset.augmentations.VoxelDownsample
+    #     voxel_size : 0.01
     #FixedRotate:
     #   _target_:  smart_tree.dataset.augmentations.FixedRotate
     #   xyz: [0, 0, 90]
@@ -108,6 +108,7 @@ pipeline:
 
   repair_skeletons : True
   smooth_skeletons : True
-  prune_skeletons : False # need to fix as some skeletons don't know their start I guess?
+  kernel_size: 10
+  prune_skeletons : False 
   min_skeleton_radius : 0.05 #0.005
   min_skeleton_length : 0.02

smart_tree/data_types/cloud.py

@@ -27,7 +27,8 @@ def to_labelled_cld(self, radii, direction, class_l) -> LabelledCloud:
         return LabelledCloud(self.xyz, self.rgb, radii * direction, class_l)
 
     def to_o3d_cld(self):
-        return o3d_cloud(self.xyz, colours=self.rgb)
+        cpu_cld = self.to_device("cpu")
+        return o3d_cloud(cpu_cld.xyz, colours=cpu_cld.rgb)
 
     def filter(self, mask):
         return Cloud(self.xyz[mask], self.rgb[mask])
@@ -50,8 +51,7 @@ def cat(self):
         )
 
     def view(self):
-        cpu_cld = self.to_device("cpu")
-        o3d_viewer([cpu_cld.to_o3d_cld()])
+        o3d_viewer([self.to_o3d_cld()])
 
     def voxel_down_sample(self, voxel_size):
         idx = voxel_downsample(self.xyz, voxel_size)
@@ -193,8 +193,8 @@ def medial_pts(self):
     @staticmethod
     def from_numpy(xyz, rgb, vector, class_l):
         return LabelledCloud(
-            torch.from_numpy(xyz).float(),  # float64 -> these data types are stupid...
-            torch.from_numpy(rgb).float(),  # float64
-            torch.from_numpy(vector).float(),  # float32
-            torch.from_numpy(class_l).float(),  # int64
+            torch.from_numpy(xyz),  # float64 -> these data types are stupid...
+            torch.from_numpy(rgb),  # float64
+            torch.from_numpy(vector),  # float32
+            torch.from_numpy(class_l),  # int64
         )

smart_tree/data_types/tree.py

@@ -109,8 +109,10 @@ def smooth(self, kernel_size=10):
         for branch in self.branches.values():
             if branch.radii.shape[0] >= kernel_size:
                 branch.radii = np.convolve(
-                    branch.radii.ravel(), kernel, mode="same"
-                ).reshape(-1, 1)
+                    branch.radii.reshape(-1),
+                    kernel,
+                    mode="same",
+                )
 
 
 @dataclass
@@ -125,9 +127,9 @@ def repair(self):
         for skeleton in self.skeletons:
             skeleton.repair()
 
-    def smooth(self):
+    def smooth(self, kernel_size=10):
         for skeleton in self.skeletons:
-            skeleton.smooth()
+            skeleton.smooth(kernel_size=kernel_size)
 
     def to_o3d_lineset(self):
         return o3d_merge_linesets(

smart_tree/pipeline.py

@@ -2,7 +2,9 @@
 
 import numpy as np
 import torch
+import time
 
+from copy import deepcopy
 from .data_types.cloud import LabelledCloud, Cloud
 from .data_types.tree import TreeSkeleton, DisjointTreeSkeleton
 from hydra.utils import instantiate
@@ -66,11 +68,11 @@ def __init__(
 
     def process_cloud(self, path: Path):
         # Load point cloud
-        cloud: Cloud = load_cloud(path)
+        cloud: Cloud = load_cloud(path).to_device(self.device)
         cloud = self.preprocessing(cloud)
 
         # Run point cloud through model to predict class, radius, direction
-        lc: LabelledCloud = self.inferer.forward(cloud).to_device("cuda")
+        lc: LabelledCloud = self.inferer.forward(cloud).to_device(self.device)
         if self.view_model_output:
             lc.view(self.cmap)
 
@@ -79,6 +81,8 @@ def process_cloud(self, path: Path):
 
         # Run the branch cloud through skeletonization algorithm, then post process
         skeleton: DisjointTreeSkeleton = self.skeletonizer.forward(branch_cloud)
+        original_skeleton = deepcopy(skeleton)
+
         self.post_process(skeleton)
 
         # View skeletonization results
@@ -89,13 +93,14 @@ def process_cloud(self, path: Path):
                     skeleton.to_o3d_lineset(),
                     skeleton.to_o3d_tube(colour=False),
                     cloud.to_o3d_cld(),
+                    original_skeleton.to_o3d_tube(),
                 ],
                 line_width=5,
             )
 
         if self.save_outputs:
-            save_o3d_mesh("skeleton.ply", skeleton.to_o3d_lineset())
-            save_o3d_lineset("mesh.ply", skeleton.to_o3d_tube())
+            save_o3d_lineset("skeleton.ply", skeleton.to_o3d_lineset())
+            save_o3d_mesh("mesh.ply", skeleton.to_o3d_tube())
             save_o3d_cloud("cloud.ply", cloud.to_o3d_cld())
 
     def post_process(self, skeleton: DisjointTreeSkeleton):
@@ -109,7 +114,7 @@ def post_process(self, skeleton: DisjointTreeSkeleton):
             skeleton.repair()
 
         if self.smooth_skeletons:
-            skeleton.smooth()
+            skeleton.smooth(kernel_size=30)
 
     @staticmethod
     def from_cfg(inferer, skeletonizer, cfg):

smart_tree/skeleton/path copy.py

@@ -13,15 +13,14 @@
 from ..util.visualizer.view import o3d_viewer
 
 
-def trace_route(preds, idx, end_points):
-    # cpu_preds = preds.cpu().numpy()
+def trace_route(preds, idx, termination_pts):
     path = []
 
-    while idx >= 0 and idx not in end_points:
+    while idx >= 0 and idx not in termination_pts:
         path.append(idx)
         idx = preds[idx]
 
-    return preds.new_tensor(path, dtype=torch.long).flip(0)
+    return preds.new_tensor(path, dtype=torch.long).flip(0), idx
 
 
 def select_path_points(
@@ -84,7 +83,7 @@ def sample_tree(
     selection_mask = preds > 0
     distances[~selection_mask] = -1
 
-    end_points = torch.tensor([], device=torch.device("cuda"))
+    termination_pts = torch.tensor([], device=torch.device("cuda"))
 
     branch_id = 0
 
@@ -101,10 +100,10 @@ def sample_tree(
             pts_sampled = f"{100 * (1.0 - ((distances > 0).sum().item() / medial_pts.shape[0])):.2f}"
             pbar.set_postfix_str(f"Sampling Graph: {pts_sampled} %")
 
-        path_vertices_idx = trace_route(
+        path_vertices_idx, termination_idx = trace_route(
             preds,
             farthest,
-            end_points,
+            termination_pts,
         )
 
         idx_points, idx_path = select_path_points(
@@ -116,7 +115,9 @@ def sample_tree(
         distances[idx_points] = -1
         distances[idx_path] = -1
 
-        end_points = torch.unique(torch.cat((end_points, idx_points, idx_path)))
+        termination_pts = torch.unique(
+            torch.cat((termination_pts, idx_point, idx_path))
+        )
 
         if len(path_vertices_idx) < 2:
             continue
@@ -125,7 +126,7 @@ def sample_tree(
             branch_id,
             xyz=medial_pts[path_vertices_idx].cpu().numpy(),
             radii=medial_radii[path_vertices_idx].cpu().numpy(),
-            parent_id=find_branch_parent(int(path_vertices_idx[0]), idx_lookup),
+            parent_id=-1,  # find_branch_parent(termination_idx, idx_lookup),
             child_id=-1,
         )
 

smart_tree/skeleton/path.py

@@ -50,7 +50,7 @@ def forward(self, cloud: LabelledCloud):
         cloud.to_device(self.device)
 
         if self.voxel_downsample != False:
-            cloud = cloud.medial_voxel_down_sample(self.voxel_downsample)
+            cloud = cloud.voxel_down_sample(0.01)
 
         mask = outlier_removal(cloud.medial_pts, cloud.radii.unsqueeze(1))
         cloud = cloud.filter(mask)