Speed up skeleton repair

Author: harrydobbs

smart_tree/conf/tree-dataset.yaml

@@ -66,10 +66,10 @@ model_inference:
   batch_size : 4
 
 skeletonizer:
-  K: 8
+  K: 16
   min_connection_length: 0.02 
   minimum_graph_vertices: 32
-  max_number_components: 8
+  max_number_components: 256
   voxel_downsample: False
   edge_non_linear: None
 
@@ -107,7 +107,7 @@ pipeline:
 
 
   repair_skeletons : True
-  smooth_skeletons : True
+  smooth_skeletons : False
   prune_skeletons : True 
-  min_skeleton_radius : 0.1 #0.005
-  min_skeleton_length : 0.02
+  min_skeleton_radius : 0.001 
+  min_skeleton_length : 0.002

smart_tree/data_types/branch.py

@@ -33,7 +33,7 @@ def to_o3d_lineset(self, colour=(0, 0, 0)):
         return o3d_path(self.xyz, colour)
 
     def to_o3d_tube(self):
-        return o3d_tube_mesh(self.xyz, self.radii, self.colour)
+        return o3d_tube_mesh(self.xyz.numpy(), self.radii.numpy(), self.colour)
 
     def to_tubes(self, colour=(1, 0, 0)) -> List[Tube]:
         a_, b_, r1_, r2_ = (

smart_tree/data_types/tree.py

@@ -2,10 +2,12 @@
 from dataclasses import dataclass
 from typing import Dict, List
 
-import numpy as np
 import torch
+import numpy as np
+import torch.nn.functional as F
+
 
-from ..util.math.queries import pts_to_nearest_tube
+from ..util.math.queries import pts_to_nearest_tube_gpu
 from ..util.mesh.geometries import (
     o3d_merge_clouds,
     o3d_merge_linesets,
@@ -72,13 +74,12 @@ def repair(self):
             parent_branch = self.branches[branch.parent_id]
             tubes = parent_branch.to_tubes()
 
-            v, idx, _ = pts_to_nearest_tube(branch.xyz[0].reshape(-1, 3), tubes)
+            v, idx, _ = pts_to_nearest_tube_gpu(branch.xyz[0].reshape(-1, 3), tubes)
 
-            # connection point...
-            connection_pt = branch.xyz[0].reshape(-1, 3) + v[0]
+            connection_pt = branch.xyz[0].reshape(-1, 3).cpu() + v[0].cpu()
 
-            branch.xyz = np.insert(branch.xyz, 0, connection_pt, axis=0)
-            branch.radii = np.insert(branch.radii, 0, branch.radii[0], axis=0)
+            branch.xyz = torch.cat((connection_pt, branch.xyz))
+            branch.radii = torch.cat((branch.radii[[0]], branch.radii))
 
     def prune(self, min_radius, min_length, root_id=0):
         """If a branch doesn't meet the initial radius threshold or length threshold we want to remove it and all
@@ -101,19 +102,17 @@ def prune(self, min_radius, min_length, root_id=0):
 
         self.branches = branches_to_keep
 
-    def smooth(self, kernel_size=10):
+    def smooth(self, kernel_size=5):
         """
         Smooths the skeleton radius.
         """
-        kernel = np.ones(kernel_size) / kernel_size
-
         for branch in self.branches.values():
-            if branch.radii.shape[0] >= kernel_size:
-                branch.radii = np.convolve(
-                    branch.radii.reshape(-1),
-                    kernel,
-                    mode="same",
-                )
+            if branch.radii.shape[0] > kernel_size:
+                branch.radii = F.conv1d(
+                    branch.radii.reshape(-1, 1, 1),
+                    torch.ones(1, 1, kernel_size),
+                    padding="same",
+                ).reshape(-1)
 
 
 @dataclass

smart_tree/data_types/tube.py

@@ -27,18 +27,24 @@ def to_numpy(self):
 
 @dataclass
 class CollatedTube:
-    a: np.array  # Nx3
-    b: np.array  # Nx3
-    r1: np.array  # N
-    r2: np.array  # N
+    a: torch.tensor  # Nx3
+    b: torch.tensor  # Nx3
+    r1: torch.tensor  # N
+    r2: torch.tensor  # N
+
+    def to_gpu(self, device=torch.device("cuda")):
+        self.a = self.a.to(device)
+        self.b = self.b.to(device)
+        self.r1 = self.r1.to(device)
+        self.r2 = self.r2.to(device)
 
 
 def collate_tubes(tubes: List[Tube]) -> CollatedTube:
-    a = np.concatenate([tube.a for tube in tubes]).reshape(-1, 3)
-    b = np.concatenate([tube.b for tube in tubes]).reshape(-1, 3)
+    a = torch.cat([tube.a for tube in tubes]).reshape(-1, 3)
+    b = torch.cat([tube.b for tube in tubes]).reshape(-1, 3)
 
-    r1 = np.asarray([tube.r1 for tube in tubes]).reshape(1, -1)
-    r2 = np.asarray([tube.r2 for tube in tubes]).reshape(1, -1)
+    r1 = torch.cat([tube.r1 for tube in tubes]).reshape(1, -1)
+    r2 = torch.cat([tube.r2 for tube in tubes]).reshape(1, -1)
 
     return CollatedTube(a, b, r1, r2)
 

smart_tree/pipeline.py

@@ -111,7 +111,7 @@ def post_process(self, skeleton: DisjointTreeSkeleton):
             skeleton.repair()
 
         if self.smooth_skeletons:
-            skeleton.smooth(kernel_size=30)
+            skeleton.smooth(kernel_size=27)
 
     @staticmethod
     def from_cfg(inferer, skeletonizer, cfg):

smart_tree/skeleton/path.py

@@ -112,27 +112,34 @@ def sample_tree(
             termination_pts,
         )
 
-        """ Gets the points around that path """
+        """ Gets the points around that path (and which path indexs they are close to) """
         idx_points, idx_path = select_path_points(
             medial_pts,
             medial_pts[path_vertices_idx],
             medial_radii[path_vertices_idx],
         )
 
+        """ Mark this points as allocated and as termination points """
         distances[idx_points] = -1
         distances[path_vertices_idx] = -1
-
         termination_pts = torch.unique(
-            torch.cat((termination_pts, idx_points, path_vertices_idx))
+            torch.cat(
+                (
+                    termination_pts,
+                    idx_points,
+                    path_vertices_idx,
+                )
+            )
         )
 
+        """ If the path has at least two points, save it as a branch """
         if len(path_vertices_idx) < 2:
             continue
 
         branches[branch_id] = BranchSkeleton(
             branch_id,
-            xyz=medial_pts[path_vertices_idx].cpu().numpy(),
-            radii=medial_radii[path_vertices_idx].cpu().numpy(),
+            xyz=medial_pts[path_vertices_idx].cpu(),
+            radii=medial_radii[path_vertices_idx].cpu(),
             parent_id=int(branch_ids[termination_idx].item()),
             child_id=-1,
         )

smart_tree/util/math/queries.py

@@ -106,14 +106,14 @@ def projection_to_distance_matrix_gpu(projections, pts):  # N x M x 3
 
 
 def pts_to_nearest_tube_gpu(
-    pts: np.array, tubes: List[Tube], device=torch.device("cuda")
+    pts: torch.tensor, tubes: List[Tube], device=torch.device("cuda")
 ):
     """Vectors from pt to the nearest tube"""
 
-    collated_tube = collate_tubes(tubes)
+    collated_tube_gpu = collate_tubes(tubes)
+    collated_tube_gpu.to_gpu()
 
-    collated_tube_gpu = collated_tube_to_gpu(collated_tube)
-    pts = torch.from_numpy(pts).float().to(device)
+    pts = pts.float().to(device)
 
     projections, t = points_to_collated_tube_projections_gpu(
         pts, collated_tube_gpu, device=torch.device("cuda")
@@ -128,9 +128,9 @@ def pts_to_nearest_tube_gpu(
     assert idx.shape[0] == pts.shape[0]
 
     return (
-        to_numpy(projections[torch.arange(pts.shape[0]), idx] - pts),
-        to_numpy(idx),
-        to_numpy(r[torch.arange(pts.shape[0]), idx]),
+        projections[torch.arange(pts.shape[0]), idx] - pts,
+        idx,
+        r[torch.arange(pts.shape[0]), idx],
     )