Added colour training / colour dropout augmentation

Author: harrydobbs

smart_tree/conf/tree-dataset.yaml

@@ -7,6 +7,7 @@ training:
   lr_decay: True
   early_stop_epoch: 20
   early_stop: True
+  use_colour: False
 
   dataset:
     _target_: smart_tree.dataset.dataset.TreeDataset
@@ -17,14 +18,16 @@ training:
     block_size: 4
     buffer_size: 0.4
 
-    augmentation:
-    # Scale:
-    #    _target_:  smart_tree.dataset.augmentations.Scale
-    #    min_scale: 0.1
-    #    max_scale: 0.1
-    # Dropout:
+    # 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
@@ -64,6 +67,7 @@ model_inference:
   buffer_size: 0.4
   num_workers : 8
   batch_size : 4
+  use_colour: False
 
 skeletonizer:
   K: 16
@@ -88,19 +92,20 @@ pipeline:
     # VoxelDownsample:
     #    _target_:  smart_tree.dataset.augmentations.VoxelDownsample
     #    voxel_size : 0.01
-    #FixedRotate:
-    #   _target_:  smart_tree.dataset.augmentations.FixedRotate
-    #   xyz: [0, 0, 90]
+
+    # FixedRotate:
+    #    _target_:  smart_tree.dataset.augmentations.FixedRotate
+    #    xyz: [0, 0, 90]
     # CentreCloud:
-    #   _target_:  smart_tree.dataset.augmentations.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]

smart_tree/dataset/augmentations.py

@@ -33,7 +33,7 @@ def __init__(self, xyz):
 
     def __call__(self, cloud):
         self.rot_mat = euler_angles_to_rotation(
-            torch.tensor(self.xyz), device=cloud.device
+            torch.tensor(self.xyz), device=cloud.xyz.device
         ).float()
         return cloud.rotate(self.rot_mat)
 
@@ -60,6 +60,19 @@ def __call__(self, cloud):
         return cloud.translate(self.xyz)
 
 
+class RandomTranslate(Augmentation):
+    def __init__(self, std):
+        self.std = torch.tensor(std)
+
+    def __call__(self, cloud):
+        return cloud.translate(
+            torch.normal(
+                torch.zeros(3, device=cloud.xyz.device),
+                std=self.std.to(cloud.xyz.device),
+            )
+        )
+
+
 class RandomDropout(Augmentation):
     def __init__(self, max_drop_out):
         self.max_drop_out = max_drop_out
@@ -76,6 +89,25 @@ def __call__(self, cloud):
         return cloud.filter(indices)
 
 
+class RandomColourDropout(Augmentation):
+    def __init__(self, max_drop_out):
+        self.max_drop_out = max_drop_out
+
+    def __call__(self, cloud):
+        num_indices = int(
+            (1.0 - (self.max_drop_out * torch.rand(1, device=cloud.rgb.device)))
+            * cloud.xyz.shape[0]
+        )
+
+        indices = torch.randint(
+            high=cloud.rgb.shape[0], size=(num_indices, 1), device=cloud.rgb.device
+        ).squeeze(1)
+
+        cloud.rgb[indices] = torch.ones_like(cloud.rgb[indices])
+
+        return cloud
+
+
 class AugmentationPipeline:
     def __init__(self, augmentation_fns: List[Augmentation]):
         # config is a dict

smart_tree/model/model.py

@@ -22,10 +22,14 @@ def __init__(
         direction_fc_planes,
         class_fc_planes,
         bias=False,
+        branch_classes=[0],
         algo=spconv.ConvAlgo.Native,
+        device=torch.device("cuda"),
     ):
         super().__init__()
 
+        self.branch_classes = torch.tensor(branch_classes, device=device)
+
         norm_fn = functools.partial(
             nn.BatchNorm1d, eps=1e-4, momentum=0.1
         )  # , momentum=0.99)
@@ -51,7 +55,6 @@ def __init__(
             algo=algo,
         )
 
-        # Three Heads...
         self.radius_head = SparseFC(
             radius_fc_planes,
             norm_fn,
@@ -108,16 +111,18 @@ def compute_loss(self, outputs, targets, mask=None):
         class_target = targets[:, [3]]
         direction_target, radius_target = torch_normalized(targets[:, :3])
 
-        vector_mask = (
-            class_target == 0
-        )  # only compute vector loss on points that are meant to be branches
-        vector_mask = vector_mask.reshape(-1)
+        mask = torch.isin(
+            class_target,
+            self.branch_classes,
+        )
+
+        mask = mask.reshape(-1)
 
         losses["radius"] = self.compute_radius_loss(
-            radius_pred[vector_mask], radius_target[vector_mask]
+            radius_pred[mask], radius_target[mask]
         )
         losses["direction"] = self.compute_direction_loss(
-            direction_pred[vector_mask], direction_target[vector_mask]
+            direction_pred[mask], direction_target[mask]
         )
         losses["class"] = self.compute_class_loss(class_pred, class_target)
 

smart_tree/model/model_inference.py

@@ -42,6 +42,7 @@ def __init__(
         buffer_size: float,
         num_workers=8,
         batch_size=4,
+        use_colour=False,
         device=torch.device("cuda:0"),
     ):
         print("Initalizing Model Inference")
@@ -55,6 +56,8 @@ def __init__(
         self.num_workers = num_workers
         self.batch_size = batch_size
 
+        self.num_input_feats = 6 if use_colour else 3
+
     def forward(self, cloud: Cloud, return_masked=True):
         outputs, inputs, masks = [], [], []
 
@@ -71,7 +74,7 @@ def forward(self, cloud: Cloud, return_masked=True):
             dataloader, desc="Inferring", leave=False
         ):
             sparse_input = sparse_from_batch(
-                features[:, :3],
+                features[:, : self.num_input_feats],
                 coordinates,
                 device=self.device,
             )
@@ -111,6 +114,7 @@ def from_cfg(cfg):
             buffer_size=cfg.buffer_size,
             num_workers=cfg.num_workers,
             batch_size=cfg.batch_size,
+            use_colour=cfg.use_colour,
         )
 
 

smart_tree/model/train.py

@@ -83,6 +83,7 @@ def main(cfg: DictConfig):
 
     epochs_no_improve = 0
     best_val_loss = torch.inf
+    num_input_feats = 6 if cfg.use_colour else 3
 
     # Training Epochs
     for epoch in tqdm(range(0, cfg.num_epoch), leave=False):
@@ -96,7 +97,7 @@ def main(cfg: DictConfig):
         ):
             with amp_ctx:
                 sparse_input = sparse_from_batch(
-                    features[:, :3],
+                    features[:, :num_input_feats],
                     coordinates,
                     device=device,
                 )
@@ -143,7 +144,7 @@ def main(cfg: DictConfig):
                 desc="Validating",
             ):
                 sparse_input = sparse_from_batch(
-                    features[:, :3],
+                    features[:, :num_input_feats],
                     coordinates,
                     device=device,
                 )

smart_tree/util/math/maths.py

@@ -16,31 +16,34 @@ def torch_normalized(v):
     return F.normalize(v), v.pow(2).sum(1).sqrt().unsqueeze(1)
 
 
-def euler_angles_to_rotation(xyz: List) -> torch.tensor:
-    x, y, z = xyz
+def euler_angles_to_rotation(xyz: List, device=torch.device("cuda")) -> torch.tensor:
+    x, y, z = xyz.to(device=device)
 
     R_X = torch.tensor(
         [
             [1.0, 0.0, 0.0],
             [0.0, torch.cos(x), -torch.sin(x)],
             [0.0, torch.sin(x), torch.cos(x)],
-        ]
+        ],
+        device=device,
     )
 
     R_Y = torch.tensor(
         [
             [torch.cos(y), 0.0, torch.sin(y)],
             [0.0, 1.0, 0.0],
             [-torch.sin(y), 0.0, torch.cos(y)],
-        ]
+        ],
+        device=device,
     )
 
     R_Z = torch.tensor(
         [
             [torch.cos(z), -torch.sin(z), 0.0],
             [torch.sin(z), torch.cos(z), 0.0],
             [0.0, 0.0, 1.0],
-        ]
+        ],
+        device=device,
     )
 
     return torch.mm(R_Z, torch.mm(R_Y, R_X))