Improved progress bar

Author: harrydobbs

README.md

@@ -1,4 +1,6 @@
-# 💡 Smart-Tree: Neural Medial Axis Approximation of Point Clouds for 3D Tree Skeletonization 🌳
+#  <center> 💡🧠🤔 Smart-Tree: Neural Medial Axis Approximation of Point Clouds for 3D Tree Skeletonization 🌳🌲🌴 </center>
+
+## 📝 Description:
 
 This GitHub repository contains code from the the paper "Smart-Tree: Neural Medial Axis Approximation of Point Clouds for 3D Tree Skeletonization". <br>
 The code provided, is a deep-learning based skeletonization method for point clouds.
@@ -51,7 +53,7 @@ We supply two different models with weights:
 * `noble-elevator-58` contains branch / foliage segmentation. <br>
 * `peach-forest-65` is only trained on points from branching structure. <br>
 
-If you wish to run smart-tree using your own weights you will need to update the model paths in the tree-dataset.yaml. <br>
+If you wish to run smart-tree using your own weights you will need to update the model paths in the `tree-dataset.yaml`. <br>
 
 To run smart-tree use: <br>
 `run-smart-tree +path=cloud_path` <br>

smart_tree/conf/tree-dataset.yaml

@@ -67,9 +67,9 @@ model_inference:
 
 skeletonizer:
   K: 16
-  min_connection_length: 0.02 
+  min_connection_length: 0.00 #.02
   minimum_graph_vertices: 32
-  max_number_components: 1000
+  max_number_components: 30
   voxel_downsample: False
   edge_non_linear: None
 
@@ -98,8 +98,8 @@ pipeline:
   repair_skeletons : True
   smooth_skeletons : True
   prune_skeletons : True
-  min_skeleton_radius : 0 #0.005
-  min_skeleton_length : 0
+  min_skeleton_radius : 0.001 #0.005
+  min_skeleton_length : 0.02
   view_model_output : False
   view_skeletons : True
   save_outputs : False

smart_tree/data_types/cloud.py

@@ -153,8 +153,7 @@ def cat(self):
 
     def medial_voxel_down_sample(self, voxel_size):
         idx = voxel_downsample_idxs(self.medial_pts, voxel_size)
-        print(idx)
-        quit()
+
         return self.filter(idx)
 
     def to_torch(self):

smart_tree/model/train.py

@@ -110,7 +110,6 @@ def main(cfg: DictConfig):
 
                 model_output = model.forward(sparse_input)
 
-                # print(model_output)
                 loss = model.compute_loss(model_output, targets, loss_mask)
                 total_loss = loss["radius"] + loss["direction"] + loss["class"]
 

smart_tree/pipeline.py

@@ -110,7 +110,7 @@ def post_process(self, skeleton: DisjointTreeSkeleton):
             skeleton.repair()
 
         if self.smooth_skeletons:
-            print("Smoothing.")
+            print("Smoothing...")
             skeleton.smooth()
 
     @staticmethod

smart_tree/skeleton/path.py

@@ -2,6 +2,8 @@
 import numpy as np
 import open3d as o3d
 import torch
+import sys
+
 from tqdm import tqdm
 
 from ..data_types.branch import BranchSkeleton
@@ -39,7 +41,9 @@ def select_path_points(
     """
 
     point_path, dists, _ = nn(
-        points, path_verts, r=radii.max().item()
+        points,
+        path_verts,
+        r=radii.max().item(),
     )  # nearest path idx for each point
     valid = dists[point_path >= 0] < radii[point_path[point_path >= 0]].squeeze(
         1
@@ -72,6 +76,7 @@ def sample_tree(
     all_points,
     root_idx=0,
     visualize=False,
+    pbar=None,
 ):
     """
     Medial Points: NN estimated medial points
@@ -91,17 +96,18 @@ def sample_tree(
     idx_lookup = {}
     branches = {}
 
-    tubes = []
+    # tubes = []
 
     while True:
-        os.system("clear")
-        print(f"{(distances > 0).sum().item() / medial_pts.shape[0]:.4f}")
-
         farthest = distances.argmax().item()  # Get fartherest away medial point
 
         if distances[farthest] <= 0:
             break
 
+        if pbar:
+            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, first_idx = trace_route(
             preds, farthest, allocated=allocated_path_points
         )  # Gets IDXs along a path and the first IDX of that path

smart_tree/skeleton/skeletonize.py

@@ -3,7 +3,7 @@
 import open3d.visualization.rendering as rendering
 import torch
 from cugraph import sssp
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 from ..data_types.tree import TreeSkeleton, DisjointTreeSkeleton
 from ..util.mesh.geometries import (
@@ -69,36 +69,25 @@ def forward(self, cloud: LabelledCloud):
             max_components=self.max_number_components,
         )
 
-        return DisjointTreeSkeleton(
-            skeletons=[
-                self.process_subgraph(
-                    subgraph,
-                    cloud.medial_pts,
-                    cloud.xyz,
-                    cloud.radii.unsqueeze(1),
-                    self.min_connection_length,
-                    self.device,
-                    skeleton_id,
-                )
-                for skeleton_id, subgraph in enumerate(
-                    tqdm(subgraphs, desc="Processing Skeleton Fragment", leave=False)
-                )
-            ]
-        )
+        skeletons = []
+        pbar = tqdm(subgraphs, position=0, leave=False)
+
+        for skeleton_id, subgraph in enumerate(pbar):
+            pbar.set_description(f"Processing Subgraph {skeleton_id}")
+            skeletons.append(self.process_subgraph(subgraph, cloud, skeleton_id, pbar))
+
+        return DisjointTreeSkeleton(skeletons)
 
     def process_subgraph(
         self,
         subgraph,
-        medial_points,
-        points,
-        radii,
-        min_edge,
-        device=torch.device("cuda:0"),
+        cloud,
         skeleton_id=0,
+        pbar=None,
     ) -> TreeSkeleton:
         edges, edge_weights = decompose_cuda_graph(subgraph, self.device)
 
-        root_idx = edges[torch.argmin(medial_points[edges[:, 0]][:, 1])][0].item()
+        root_idx = edges[torch.argmin(cloud.medial_pts[edges[:, 0]][:, 1])][0].item()
 
         verts, preds, distance = shortest_paths(
             root_idx,
@@ -107,13 +96,20 @@ def process_subgraph(
             renumber=False,
         )
 
-        predecessor_graph = pred_graph(verts, preds, medial_points)
+        predecessor_graph = pred_graph(verts, preds, cloud.medial_pts)
 
         distances = torch.as_tensor(
             sssp(predecessor_graph, source=root_idx)["distance"]
         ).to(self.device)
 
-        branches = sample_tree(medial_points, radii, preds, distances, points)
+        branches = sample_tree(
+            cloud.medial_pts,
+            cloud.radii.unsqueeze(1),
+            preds,
+            distances,
+            cloud.xyz,
+            pbar=pbar,
+        )
 
         return TreeSkeleton(skeleton_id, branches)
 

smart_tree/util/misc.py

@@ -1,6 +1,7 @@
 import math
 from typing import List, Union
 
+import sys
 import cmapy
 import numpy as np
 import open3d as o3d