V1.0.0 beta

Author: harrydobbs

smart_tree/conf/tree-dataset.yaml

@@ -66,7 +66,7 @@ model_inference:
   batch_size : 4
 
 skeletonizer:
-  K: 16
+  K: 8
   min_connection_length: 0.02 
   minimum_graph_vertices: 32
   max_number_components: 8
@@ -108,7 +108,6 @@ pipeline:
 
   repair_skeletons : True
   smooth_skeletons : True
-  kernel_size: 10
-  prune_skeletons : False 
-  min_skeleton_radius : 0.05 #0.005
+  prune_skeletons : True 
+  min_skeleton_radius : 0.1 #0.005
   min_skeleton_length : 0.02

smart_tree/data_types/tree.py

@@ -34,7 +34,7 @@ def to_o3d_linesets(self) -> List:
         return [b.to_o3d_lineset() for b in self.branches.values()]
 
     def to_o3d_lineset(self, colour=(0, 0, 0)):
-        return o3d_merge_linesets(self.to_o3d_linesets())
+        return o3d_merge_linesets(self.to_o3d_linesets(), colour=colour)
 
     def to_o3d_tubes(self) -> List:
         return [b.to_o3d_tube() for b in self.branches.values()]
@@ -94,7 +94,8 @@ def prune(self, min_radius, min_length, root_id=0):
 
             if branch.parent_id in branches_to_keep:
                 if branch.length > min_length and (
-                    (branch.radii[0] > min_radius) or branch.radii[-1] > min_radius
+                    max(branch.radii[0], branch.radii[-1])
+                    > min_radius  # We don't know which end of the branch is the start for skeletons that aren't connected...
                 ):
                     branches_to_keep[branch_id] = branch
 
@@ -120,8 +121,10 @@ class DisjointTreeSkeleton:
     skeletons: List[TreeSkeleton]
 
     def prune(self, min_radius, min_length):
-        for skeleton in self.skeletons:
-            skeleton.prune(min_radius=min_radius, min_length=min_length)
+        self.skeletons[0].prune(
+            min_radius=min_radius,
+            min_length=min_length,
+        )  # Can only prune the first skeleton as we don't know the root points for all the other skeletons...
 
     def repair(self):
         for skeleton in self.skeletons:

smart_tree/skeleton/path.py

@@ -80,15 +80,20 @@ def sample_tree(
     Surface Points: The point the medial pts got projected from..
     """
 
+    branch_id = 0
+
+    branches = {}
+
     selection_mask = preds > 0
     distances[~selection_mask] = -1
 
     termination_pts = torch.tensor([], device=torch.device("cuda"))
-
-    branch_id = 0
-
-    idx_lookup = {}
-    branches = {}
+    branch_ids = torch.full(
+        (medial_pts.shape[0],),
+        -1,
+        device=torch.device("cuda"),
+        dtype=int,
+    )
 
     while True:
         farthest = distances.argmax().item()
@@ -100,23 +105,25 @@ 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} %")
 
+        """ Traces the path of the futhrest point until it converges with allocated points """
         path_vertices_idx, termination_idx = trace_route(
             preds,
             farthest,
             termination_pts,
         )
 
+        """ Gets the points around that path """
         idx_points, idx_path = select_path_points(
             medial_pts,
             medial_pts[path_vertices_idx],
             medial_radii[path_vertices_idx],
         )
 
         distances[idx_points] = -1
-        distances[idx_path] = -1
+        distances[path_vertices_idx] = -1
 
         termination_pts = torch.unique(
-            torch.cat((termination_pts, idx_points, idx_path))
+            torch.cat((termination_pts, idx_points, path_vertices_idx))
         )
 
         if len(path_vertices_idx) < 2:
@@ -126,11 +133,13 @@ def sample_tree(
             branch_id,
             xyz=medial_pts[path_vertices_idx].cpu().numpy(),
             radii=medial_radii[path_vertices_idx].cpu().numpy(),
-            parent_id=-1,  # find_branch_parent(termination_idx, idx_lookup),
+            parent_id=int(branch_ids[termination_idx].item()),
             child_id=-1,
         )
 
-        idx_lookup[branch_id] = idx_path.cpu().tolist() + idx_points.cpu().tolist()
+        branch_ids[path_vertices_idx] = branch_id
+        branch_ids[idx_points] = branch_id
+
         branch_id += 1
 
     return branches

smart_tree/util/mesh/geometries.py

@@ -41,7 +41,7 @@ def o3d_merge_linesets(line_sets, colour=[0, 0, 0]):
 
     return o3d.geometry.LineSet(
         o3d.utility.Vector3dVector(points), o3d.utility.Vector2iVector(idxs)
-    )
+    ).paint_uniform_color(colour)
 
 
 def points_to_edge_idx(points):