merge

Author: harrydobbs

pyproject.toml

@@ -3,36 +3,28 @@ requires = ["setuptools", "setuptools-scm"]
 build-backend = "setuptools.build_meta"
 
 [project]
-name = "Smart-Tree"
+name = "smart-tree"
+version = "1.0.0"
 authors = [
     {name = "Harry Dobbs", email = "harrydobbs87@gmail.com"},
 ]
 description = "Neural Network Point Cloud Tree Skeletonization"
-readme = "README.rst"
+readme = "README.md"
 requires-python = ">=3.8"
 license = {text = "BSD-3"}
 dependencies = [
     'numpy',
     'open3d',
     'hydra-core>=1.2.0',
     'click',
-#   'opencv-python',
-#   'pytorch3d@https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu113_pyt1110/pytorch3d-0.7.2-cp38-cp38-linux_x86_64.whl',
-#    'jaxtyping',
-#    'scikit-image',
-#    'scikit-learn',
-#    'scipy',
     'oauthlib',
-#   'type_enforced',
     'spconv-cu117',
     'wandb',
     'cmapy',
     'pykeops',
- #   'seaborn',
     'plyfile',
     'py_structs'
 ]
-dynamic = ["version"]
 
 [tool.setuptools.packages]
 find = {}  # Scan the project directory with the default parameters

smart_tree/__init__.py

@@ -0,0 +1 @@
+__version__ = "1.0.0"

smart_tree/cli.py

@@ -13,7 +13,7 @@
 
 @hydra.main(
     version_base=None,
-    config_path="conf",
+    config_path=".conf",
     config_name="tree-dataset",
 )
 def main(cfg: DictConfig):
@@ -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

@@ -67,9 +67,9 @@ model_inference:
 
 skeletonizer:
   K: 16
-  min_connection_length: 0.00 #.02
+  min_connection_length: 0.02 
   minimum_graph_vertices: 32
-  max_number_components: 30
+  max_number_components: 8
   voxel_downsample: False
   edge_non_linear: None
 
@@ -78,32 +78,37 @@ pipeline:
 
   preprocessing:
     # Scale:
-    #     _target_:  smart_tree.dataset.augmentations.Scale
-    #     min_scale: 1
-    #     max_scale: 1
+    #       _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
+    #     _target_:  smart_tree.dataset.augmentations.VoxelDownsample
+    #     voxel_size : 0.01
     #FixedRotate:
     #   _target_:  smart_tree.dataset.augmentations.FixedRotate
     #   xyz: [0, 0, 90]
     # CentreCloud:
     #   _target_:  smart_tree.dataset.augmentations.CentreCloud
 
 
-  repair_skeletons : True
-  smooth_skeletons : True
-  prune_skeletons : True
-  min_skeleton_radius : 0.001 #0.005
-  min_skeleton_length : 0.02
-  view_model_output : False
+  view_model_output : True
   view_skeletons : True
   save_outputs : False
   branch_classes: [0]
   cmap:
-    - [1, 0, 0] # Trunk
-    - [0, 1, 0] # Foliage
+    - [0.325, 0.207, 0.039] # Trunk
+    - [0.290, 0.703, 0.254] # Foliage
+
+
+
+
+  repair_skeletons : True
+  smooth_skeletons : True
+  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)
@@ -111,6 +111,7 @@ def number_classes(self):
     @property
     def cmap(self):
         return torch.rand(self.number_classes, 3)
+<<<<<<< HEAD
 
     def __add__(self, other):
         xyz = torch.cat((self.xyz, other.xyz))
@@ -119,6 +120,8 @@ def __add__(self, other):
         class_l = torch.cat((self.class_l, other.class_l))
 
         return LabelledCloud(xyz, rgb, vector, class_l)
+=======
+>>>>>>> revert
 
     def filter(self, mask):
         return LabelledCloud(
@@ -137,11 +140,7 @@ def filter_by_class(self, classes: List):
         return self.filter(mask)
 
     def view(self, cmap=[]):
-        if len(cmap) != 0:
-            cmap = cmap
-        else:
-            cmap = self.cmap
-
+        cmap = cmap if cmap != [] else self.cmap
         cpu_cld = self.to_device("cpu")
         input_cld = cpu_cld.to_o3d_cld()
         segmented_cld = o3d_cloud(cpu_cld.xyz, colours=cmap[cpu_cld.class_l])
@@ -210,17 +209,8 @@ def medial_pts(self):
     @staticmethod
     def from_numpy(xyz, rgb, vector, class_l):
         return LabelledCloud(
-            torch.from_numpy(xyz).float(),  # -> these data types are stupid...
-            torch.from_numpy(rgb).float(),  # float64
-            torch.from_numpy(vector).float(),  # float32
-            torch.from_numpy(class_l).int(),  # int64
-        )
-
-    @staticmethod
-    def from_o3d_cld(cld, class_l):
-        return LabelledCloud.from_numpy(
-            xyz=np.asarray(cld.points),
-            rgb=np.asarray(cld.colors),
-            vector=np.asarray(cld.points) * 0 - 1,
-            class_l=np.asarray(class_l),
+            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

@@ -93,7 +93,9 @@ def prune(self, min_radius, min_length, root_id=0):
                 continue
 
             if branch.parent_id in branches_to_keep:
-                if branch.length > min_length and branch.radii[0] > min_radius:
+                if branch.length > min_length and (
+                    (branch.radii[0] > min_radius) or branch.radii[-1] > min_radius
+                ):
                     branches_to_keep[branch_id] = branch
 
         self.branches = branches_to_keep
@@ -107,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
@@ -123,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/dataset/dataset.py

@@ -155,8 +155,6 @@ def compute_blocks(self):
 
         self.block_centres = self.block_centres.to(torch.device("cpu"))
 
-        print("Blocks Computed")
-
     def __getitem__(self, idx):
         block_centre = self.block_centres[idx]
         cloud: Cloud = self.clouds[idx]

smart_tree/model/model_inference.py

@@ -46,19 +46,21 @@ def __init__(
         batch_size=4,
         use_colour=False,
         device=torch.device("cuda:0"),
+        verbose=False,
     ):
-        print("Initalizing Model Inference")
         self.device = device
-        self.model = load_model(model_path, weights_path, self.device)
-        print("Model Loaded")
+        self.verbose = verbose
         self.voxel_size = voxel_size
         self.block_size = block_size
         self.buffer_size = buffer_size
 
         self.num_workers = num_workers
         self.batch_size = batch_size
 
-        self.num_input_feats = 6 if use_colour else 3
+        self.model = load_model(model_path, weights_path, self.device)
+
+        if self.verbose:
+            print("Model Loaded Succesfully")
 
     def forward(self, cloud: Cloud, return_masked=True):
         outputs, inputs, masks = [], [], []

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
@@ -41,7 +43,6 @@ def __init__(
         cmap=[[1, 0, 0], [0, 1, 0]],
         device=torch.device("cuda:0"),
     ):
-        print("Setting up pipeline...")
         self.inferer = inferer
         self.skeletonizer = skeletonizer
 
@@ -67,19 +68,21 @@ 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(cmap=self.cmap)
+            lc.view(self.cmap)
 
         # Filter only the branch points for skeletonizaiton
         branch_cloud: LabelledCloud = lc.filter_by_class(self.branch_classes)
 
         # 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
@@ -88,16 +91,17 @@ def process_cloud(self, path: Path):
                 [
                     skeleton.to_o3d_tube(),
                     skeleton.to_o3d_lineset(),
-                    cloud.to_o3d_cld(),
                     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_cloud("mesh.ply", cloud.to_o3d_cld())
+            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):
         if self.prune_skeletons:
@@ -110,8 +114,7 @@ def post_process(self, skeleton: DisjointTreeSkeleton):
             skeleton.repair()
 
         if self.smooth_skeletons:
-            print("Smoothing...")
-            skeleton.smooth()
+            skeleton.smooth(kernel_size=30)
 
     @staticmethod
     def from_cfg(inferer, skeletonizer, cfg):