Fixed radius smoothing kernel dimensions

Author: harrydobbs

Diff for: smart_tree/conf/tree-dataset.yaml

@@ -11,7 +11,7 @@ training:
   dataset:
     _target_: smart_tree.dataset.dataset.TreeDataset
     voxel_size: 0.01
-    json_path: /local/uc-vision/smart-tree/smart_tree/conf/tree-split.json
+    json_path: smart_tree/conf/tree-split.json
     directory: /local/uc-vision/dataset/branches/
     blocking: True
     block_size: 4
@@ -55,10 +55,9 @@ training:
     _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
     mode: "min"
 
-
 model_inference: 
-  model_path: /local/uc-vision/smart-tree/smart_tree/model/weights/noble-elevator-58_model.pt
-  weights_path: /local/uc-vision/smart-tree/smart_tree/model/weights/noble-elevator-58_model_weights.pt
+  model_path: smart_tree/model/weights/noble-elevator-58_model.pt
+  weights_path: smart_tree/model/weights/noble-elevator-58_model_weights.pt
   voxel_size: 0.01
   block_size: 4
   buffer_size: 0.4
@@ -73,18 +72,14 @@ skeletonizer:
   voxel_downsample: False
   edge_non_linear: None
 
-
 pipeline:
 
   preprocessing:
     # 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
+    #        _target_:  smart_tree.dataset.augmentations.Scale
+    #        min_scale: 1
+    #        max_scale: 1
+
     VoxelDownsample:
          _target_:  smart_tree.dataset.augmentations.VoxelDownsample
          voxel_size : 0.01
@@ -100,14 +95,12 @@ pipeline:
   save_outputs : False
   branch_classes: [0]
   cmap:
-    - [0.325, 0.207, 0.039] # Trunk
-    - [0.290, 0.703, 0.254] # Foliage
-
-
-
+    - [0.450, 0.325, 0.164] # Trunk
+    - [0.541, 0.670, 0.164] # Foliage
 
   repair_skeletons : True
-  smooth_skeletons : False
+  smooth_skeletons : True
+  smooth_kernel_size: 21 # Needs to be odd
   prune_skeletons : True 
   min_skeleton_radius : 0.001 
   min_skeleton_length : 0.002

Diff for: smart_tree/data_types/tree.py

@@ -106,13 +106,14 @@ def smooth(self, kernel_size=5):
         """
         Smooths the skeleton radius.
         """
-        kernel = torch.ones(1, 1, kernel_size)
+        kernel = torch.ones(1, 1, kernel_size) / kernel_size
         for branch in self.branches.values():
-            branch.radii = F.conv1d(
-                branch.radii.unsqueeze(2),
-                kernel,
-                padding=(kernel_size - 1) // 2,
-            ).squeeze(2)
+            if branch.radii.shape[0] > kernel_size:
+                branch.radii = F.conv1d(
+                    branch.radii.reshape(1, 1, -1),
+                    kernel,
+                    padding="same",
+                ).reshape(-1)
 
 
 @dataclass

Diff for: smart_tree/model/weights/__init__.py

@@ -32,6 +32,7 @@ def __init__(
         preprocessing_cfg,
         repair_skeletons=False,
         smooth_skeletons=False,
+        smooth_kernel_size=0,
         prune_skeletons=False,
         min_skeleton_radius=0.0,
         min_skeleton_length=1000,
@@ -51,6 +52,7 @@ def __init__(
 
         self.repair_skeletons = repair_skeletons
         self.smooth_skeletons = smooth_skeletons
+        self.smooth_kernel_size = smooth_kernel_size
         self.prune_skeletons = prune_skeletons
 
         self.min_skeleton_radius = min_skeleton_radius
@@ -92,7 +94,7 @@ def process_cloud(self, path: Path):
                     skeleton.to_o3d_tube(colour=False),
                     cloud.to_o3d_cld(),
                 ],
-                line_width=10,
+                line_width=5,
             )
 
         if self.save_outputs:
@@ -111,7 +113,7 @@ def post_process(self, skeleton: DisjointTreeSkeleton):
             skeleton.repair()
 
         if self.smooth_skeletons:
-            skeleton.smooth(kernel_size=27)
+            skeleton.smooth(self.smooth_kernel_size)
 
     @staticmethod
     def from_cfg(inferer, skeletonizer, cfg):
@@ -121,6 +123,7 @@ def from_cfg(inferer, skeletonizer, cfg):
             preprocessing_cfg=cfg.preprocessing,
             repair_skeletons=cfg.repair_skeletons,
             smooth_skeletons=cfg.smooth_skeletons,
+            smooth_kernel_size=cfg.self.smooth_kernel_size,
             prune_skeletons=cfg.prune_skeletons,
             min_skeleton_radius=cfg.min_skeleton_radius,
             min_skeleton_length=cfg.min_skeleton_length,