merge

Author: harrydobbs

.gitignore

@@ -131,12 +131,16 @@ dmypy.json
 
 .vscode/
 
+debug/
 outputs/
 data/
 multirun/
 wandb/
 *.out
 *.sl
+*.pcd
+*.ply
+
 speed-tree-outputs/
 smart_tree/conf/tree-dataset-test.yaml
 smart_tree/conf/tree-split-test.json

smart_tree/conf/tree-dataset-test.yaml

@@ -0,0 +1,116 @@
+training:
+
+  name: "tree"
+  wblogging: "online" 
+  fp16: False
+  num_epoch: 10
+  lr_decay: True
+  early_stop_epoch: 20
+  early_stop: True
+  use_colour: False
+
+  dataset:
+    _target_: smart_tree.dataset.dataset.TreeDataset
+    voxel_size: 0.01
+    json_path: /local/uc-vision/smart-tree/smart_tree/conf/tree-split.json
+    directory: /local/uc-vision/dataset/branches/
+    blocking: True
+    block_size: 4
+    buffer_size: 0.4
+
+    # augmentation:
+    #     _target_:  smart_tree.dataset.augmentations.Scale
+    #     min_scale: 0.9
+    #     max_scale: 1.1
+      
+    #     _target_:  smart_tree.dataset.augmentations.RandomDropout
+    #     max_drop_out: 0.1
+      
+    #     _target_:  smart_tree.dataset.augmentations.RandomColourDropout
+    #     max_drop_out: 0.1
+
+  data_loader:
+    _target_: torch.utils.data.DataLoader
+    batch_size: 16
+    drop_last: True
+    pin_memory: False
+    num_workers: 0
+    #prefetch_factor: 
+    shuffle: True
+    collate_fn:
+      _target_: smart_tree.model.sparse.batch_collate
+      _partial_: True
+
+  model:
+    _target_: smart_tree.model.model.Smart_Tree
+    input_channels: 3
+    unet_planes: [8, 16, 32, 64, 128, 256]
+    radius_fc_planes: [8, 8, 4, 1]
+    direction_fc_planes: [8, 8, 4, 3]
+    class_fc_planes: [8, 8, 4, 2] # last one is number of classes..
+    bias: False
+
+  optimizer:
+    _target_: torch.optim.Adam # SGD, Adadelta etc
+    lr: 0.1
+    #eps: 0.00001
+
+
+  scheduler:
+    _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
+  voxel_size: 0.01
+  block_size: 4
+  buffer_size: 0.4
+  num_workers : 8
+  batch_size : 4
+  use_colour: False
+
+skeletonizer:
+  K: 16
+  min_connection_length: 0.00 #.02
+  minimum_graph_vertices: 32
+  max_number_components: 30
+  voxel_downsample: False
+  edge_non_linear: None
+
+
+pipeline:
+  
+  preprocessing:
+    # Scale:
+    #     _target_:  smart_tree.dataset.augmentations.Scale
+    #     min_scale: 1
+    #     max_scale: 1
+    # Scale:
+    #   _target_:  smart_tree.dataset.augmentations.Scale
+    #   min_scale: 1
+    #   max_scale: 1
+    # VoxelDownsample:
+    #    _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 : True
+  view_skeletons : True
+  save_outputs : False
+  branch_classes: [0]
+  cmap:
+    - [1, 0, 0] # Trunk
+    - [0, 1, 0] # Foliage

smart_tree/data_types/cloud.py

@@ -75,12 +75,17 @@ def max_xyz(self):
     def min_xyz(self):
         return torch.min(self.xyz, 0)[0]
 
+    @property
+    def centre(self):
+        return self.min_xyz + (self.max_xyz - self.min_xyz) / 2
+
+    @property
+    def dimensions(self):
+        return self.max_xyz - self.min_xyz
+
     @property
     def bbox(self):
-        # defined by centre coordinate, x/2, y/2, z/2
-        dimensions = (self.max_xyz - self.min_xyz) / 2
-        centre = self.min_xyz + dimensions
-        return centre, dimensions
+        return self.centre, self.dimensions
 
     @staticmethod
     def from_numpy(xyz, rgb, device=torch.device("cpu")):
@@ -99,9 +104,21 @@ class LabelledCloud(Cloud):
     vector: torch.Tensor
     class_l: torch.Tensor
 
-    def __post_init__(self):
-        num_classes = int(torch.max(self.class_l, 0)[0].item())
-        self.cmap = torch.rand(num_classes + 1, 3)
+    @property
+    def number_classes(self):
+        return int(torch.max(self.class_l, 0)[0].item()) + 1
+
+    @property
+    def cmap(self):
+        return torch.rand(self.number_classes, 3)
+
+    def __add__(self, other):
+        xyz = torch.cat((self.xyz, other.xyz))
+        rgb = torch.cat((self.rgb, other.rgb))
+        vector = torch.cat((self.vector, other.vector))
+        class_l = torch.cat((self.class_l, other.class_l))
+
+        return LabelledCloud(xyz, rgb, vector, class_l)
 
     def filter(self, mask):
         return LabelledCloud(
@@ -198,3 +215,12 @@ def from_numpy(xyz, rgb, vector, class_l):
             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),
+        )

smart_tree/dataset/augmentations.py

@@ -1,13 +1,25 @@
+import os
 import random
 
 import numpy as np
 import torch
+import random
+import open3d as o3d
 
+from tqdm import tqdm
+from pathlib import Path
 from abc import ABC, abstractmethod
 from typing import List
 
-from smart_tree.data_types.cloud import Cloud
+from smart_tree.data_types.cloud import Cloud, LabelledCloud
 from smart_tree.util.math.maths import euler_angles_to_rotation
+from smart_tree.util.file import (
+    load_o3d_mesh,
+    save_o3d_mesh,
+    save_o3d_cloud,
+    load_o3d_cloud,
+)
+from smart_tree.util.mesh.geometries import o3d_cloud, o3d_sphere
 from hydra.utils import call, get_original_cwd, instantiate, to_absolute_path
 
 
@@ -73,6 +85,89 @@ def __call__(self, cloud):
         )
 
 
+class RandomMesh(Augmentation):
+    def __init__(
+        self,
+        mesh_directory: Path,
+        preprocessed_path: Path,
+        voxel_size: float,
+        number_meshes: int,
+        min_size: float,
+        max_size: float,
+        max_pts: int,
+    ):
+        """We want to preprocess the meshes by getting them to the right scale,
+        which is done by first converting them from mm to metres, we then
+        scale them up to the max_size and then do a point sample, based on target voxel_size
+        (ensure we have enough point density at the max size), then revert
+        the scale back normal scale in metres. During inference we randomly scale based on the
+        min_size and max_size and then translate the points and merge with the input cloud
+        """
+
+        self.voxel_size = voxel_size
+        self.number_meshes = number_meshes
+        self.min_size = min_size
+        self.max_size = max_size
+        self.class_number = 2
+        self.preprocessed_path = preprocessed_path
+
+        if not (os.path.exists(preprocessed_path)):
+            os.makedirs(preprocessed_path)
+
+        for mesh_path in tqdm(
+            Path(mesh_directory).glob("*.stl"),
+            desc="Preprocessing Meshes",
+            leave=False,
+        ):
+            if os.path.isfile(f"{preprocessed_path}/{mesh_path.stem}.pcd"):
+                continue
+            try:
+                mesh = load_o3d_mesh(str(mesh_path))
+                pcd = (
+                    mesh.scale(0.001, mesh.get_center())
+                    .translate(-mesh.get_center())
+                    .paint_uniform_color(np.random.rand(3))
+                    .scale(max_size, mesh.get_center())
+                    .sample_points_uniformly(
+                        min(
+                            max(int(mesh.get_surface_area() / (voxel_size**2)), 10),
+                            max_pts,
+                        )
+                    )
+                    .scale(1 / max_size, mesh.get_center())
+                )
+                save_o3d_cloud(f"{preprocessed_path}/{mesh_path.stem}.pcd", pcd)
+            except:
+                print(f"Cloud Generation Failed on {mesh_path}")
+
+    def __call__(self, cloud):
+        centre, dimensions = cloud.bbox
+
+        for i in range(self.number_meshes):
+            random_pcd_path = random.choice(list(self.preprocessed_path.glob("*.pcd")))
+            pcd = load_o3d_cloud(str(random_pcd_path))
+            scaled_pcd = pcd.scale(
+                random.uniform(self.min_size, self.max_size), pcd.get_center()
+            )
+
+            lc = LabelledCloud.from_o3d_cld(
+                pcd,
+                class_l=torch.ones(np.asarray(pcd.points).shape[0]) * self.class_number,
+            )
+
+            lc = lc.rotate(
+                euler_angles_to_rotation(torch.rand(3) * torch.pi * 2).to(
+                    cloud.xyz.device
+                )
+            )
+            lc = lc.translate(cloud.min_xyz - lc.centre)
+            lc = lc.translate(dimensions * torch.rand(3))
+
+            cloud += lc
+
+        return cloud
+
+
 class RandomDropout(Augmentation):
     def __init__(self, max_drop_out):
         self.max_drop_out = max_drop_out
@@ -127,13 +222,36 @@ def from_cfg(cfg):
 
 if __name__ == "__main__":
     from pathlib import Path
-    from smart_tree.util.file import load_cloud
+    from smart_tree.util.file import load_data_npz
     from smart_tree.util.visualizer.view import o3d_viewer
 
-    cld = load_cloud(
-        Path("/media/harry/harry's-data/PhD/training-data/apple/apple_1.npz")
+    mesh_adder = RandomMesh(
+        mesh_directory=Path("/local/Datasets/Thingi10K/raw_meshes/"),
+        preprocessed_path=Path(
+            "/local/uc-vision/smart-tree/data/things10K_sampled_1mm/"
+        ),
+        voxel_size=0.001,
+        number_meshes=20,
+        min_size=0.01,
+        max_size=20,
+        max_pts=50000,
+    )
+
+    cld, _ = load_data_npz(Path("/local/_smart-tree/evaluation-data/gt/apple_12.npz"))
+
+    cld = mesh_adder(cld)
+
+    cld.view()
+
+    o3d_viewer(
+        [
+            cld.to_o3d_cld(),
+            o3d_sphere(cld.min_xyz, radius=0.1, colour=(0, 1, 0)),
+            o3d_sphere(cld.max_xyz, radius=0.1),
+        ]
     )
 
+    quit()
     centre = CentreCloud()
     rotater = FixedRotate(torch.tensor([torch.pi / 2, torch.pi / 2, torch.pi * 2]))
     do = RandomDropout(0.5)

smart_tree/dataset/dataset.py

@@ -51,9 +51,6 @@ def __getitem__(self, idx):
 
         labelled_cld = labelled_cld.to_device(torch.device("cuda"))
 
-        if self.augmentation != None:
-            labelled_cld = self.augmentation(labelled_cld)
-
         if self.blocking:
             block_center_idx = torch.randint(
                 labelled_cld.xyz.shape[0], size=(1,), device=labelled_cld.xyz.device
@@ -66,6 +63,9 @@ def __getitem__(self, idx):
             )
             labelled_cld = labelled_cld.filter(block_filter)
 
+        if self.augmentation != None:
+            labelled_cld = self.augmentation(labelled_cld)
+
         xyzmin, _ = torch.min(labelled_cld.xyz, axis=0)
         xyzmax, _ = torch.max(labelled_cld.xyz, axis=0)
         make_voxel_gen = time.time()