added random object generation augmentation

Author: harrydobbs

.gitignore

@@ -138,6 +138,9 @@ multirun/
 wandb/
 *.out
 *.sl
+*.pcd
+*.ply
+
 speed-tree-outputs/
 smart_tree/conf/tree-split-test.json
 smart_tree/conf/apple-trellis-split-test.json

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")):

smart_tree/dataset/augmentations.py

@@ -1,17 +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, LabelledCloud
 from smart_tree.util.math.maths import euler_angles_to_rotation
-from smart_tree.util.file import load_o3d_mesh
+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
 
 
@@ -81,45 +89,81 @@ 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,
     ):
-        self.mesh_paths = list(mesh_directory.glob("*"))
+        """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 = 3
+        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):
-            mesh = load_o3d_mesh(
-                str(self.mesh_paths[random.randint(0, len(self.mesh_paths))])
+            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()
             )
 
-            mesh = mesh.scale(0.01, center=mesh.get_center())
-
-            mesh = mesh.translate(-mesh.get_center())
-
-            mesh_pts = mesh.sample_points_uniformly(
-                int(1000 * mesh.get_surface_area() / self.voxel_size),
-            ).paint_uniform_color(np.random.rand(3))
-
             lc = LabelledCloud.from_o3d_cld(
-                mesh_pts,
-                class_l=torch.ones(np.asarray(mesh_pts.points).shape[0])
-                * self.class_number,
+                pcd,
+                class_l=torch.ones(np.asarray(pcd.points).shape[0]) * self.class_number,
             )
 
-            cloud += lc
+            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))
 
-        # load random mesh
-        # voxelize mesh
-        # create labelled cloud
-        # randomly translate / rotate it
-        # return new cloud
+            cloud += lc
 
         return cloud
 
@@ -183,10 +227,14 @@ def from_cfg(cfg):
 
     mesh_adder = RandomMesh(
         mesh_directory=Path("/local/Datasets/Thingi10K/raw_meshes/"),
-        voxel_size=0.01,
-        number_meshes=5,
+        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=0.5,
+        max_size=20,
+        max_pts=50000,
     )
 
     cld, _ = load_data_npz(Path("/local/_smart-tree/evaluation-data/gt/apple_12.npz"))
@@ -195,7 +243,13 @@ def from_cfg(cfg):
 
     cld.view()
 
-    # o3d_viewer([cld])
+    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()

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()

smart_tree/scripts/sample-meshes.py

@@ -0,0 +1,69 @@
+""" Script to sample meshes into point clouds.
+
+  usage: 
+
+  python smart_tree/scripts/sample-meshes.py --read_directory=/local/Datasets/Thingi10K/raw_meshes/ --save_directory=/local/Datasets/Thingi10K/point_clouds/ --scale=0.001 --voxel_size=0.01
+
+
+"""
+
+import json
+import os
+import random
+from pathlib import Path
+
+import click
+import numpy as np
+
+from tqdm import tqdm
+
+from smart_tree.util.file import load_o3d_mesh, load_data_npz
+from smart_tree.util.visualizer.view import o3d_viewer
+
+
+@click.command()
+@click.option(
+    "--read_directory",
+    type=click.Path(exists=True),
+    prompt="read directory?",
+)
+@click.option(
+    "--save_directory",
+    type=click.Path(exists=True),
+    prompt="save directory?",
+)
+@click.option(
+    "--scale",
+    type=float,
+    prompt="scale?",
+)
+@click.option(
+    "--voxel_size",
+    type=float,
+    prompt="voxel_size?",
+)
+def main(read_directory, save_directory, scale, voxel_size):
+    files_paths = [str(path) for path in Path(read_directory).glob("*.stl")]
+
+    pcd, skeleton = load_data_npz("/local/synthetic-trees/dataset/apple/apple_1.npz")
+
+    for file_path in tqdm(files_paths):
+        try:
+            mesh = load_o3d_mesh(file_path)
+            mesh = (
+                mesh.scale(scale, mesh.get_center())
+                .translate(-mesh.get_center())
+                .paint_uniform_color(np.random.rand(3))
+                .scale(100, mesh.get_center())
+                .sample_points_uniformly(
+                    int(mesh.get_surface_area() / (voxel_size**2))
+                )
+            )
+        except:
+            pass
+
+        # o3d_viewer([mesh, pcd.to_o3d_cld()])
+
+
+if __name__ == "__main__":
+    main()

smart_tree/scripts/split-data copy.py

@@ -0,0 +1,82 @@
+""" Script to split dataset into train and test sets 
+
+  usage: 
+
+  python smart_tree/scripts/split-data.py --read_directory=/speed-tree/speed-tree-outputs/processed_vines/ --json_save_path=/smart-tree/smart_tree/conf/vine-split.json --sample_type=random
+
+
+"""
+
+import json
+import os
+import random
+from pathlib import Path
+
+import click
+import numpy as np
+
+
+def flatten_list(l):
+    return [item for sublist in l for item in sublist]
+
+
+def random_sample(read_dir, json_save_path):
+    items = [str(path.name) for path in Path(read_dir).rglob("*.npz")]
+    random.shuffle(items)
+
+    data = {}
+
+    data["train"] = sorted(items[: int(0.8 * len(items))])
+    data["test"] = sorted(items[int(0.8 * len(items)) : int(0.9 * len(items))])
+    data["validation"] = sorted(items[int(0.9 * len(items)) :])
+
+    with open(json_save_path, "w") as outfile:
+        json.dump(data, outfile, indent=4, sort_keys=False)
+
+
+def strattified_sample(read_dir, json_save_path):
+    dirs = os.listdir(read_dir)
+
+    train_paths = []
+    test_paths = []
+    val_paths = []
+
+    for directory in dirs:
+        items = [
+            str(path.resolve())
+            for path in Path(f"{read_dir}/{directory}").rglob("*.npz")
+        ]
+        random.shuffle(items)
+
+        train_paths.append(items[: int(0.8 * len(items))])
+        test_paths.append(
+            items[int(0.8 * len(items)) : int(0.8 * len(items) + int(0.1 * len(items)))]
+        )
+        val_paths.append(items[int(0.8 * len(items)) + int(0.1 * len(items)) :])
+
+    data = {}
+
+    data["train"] = sorted(flatten_list(train_paths))
+    data["test"] = sorted(flatten_list(test_paths))
+    data["validation"] = sorted(flatten_list(val_paths))
+
+    with open(json_save_path, "w") as outfile:
+        json.dump(data, outfile, indent=4, sort_keys=False)
+
+
+@click.command()
+@click.option(
+    "--read_directory", type=click.Path(exists=True), prompt="read directory?"
+)
+@click.option("--json_save_path", prompt="json path?")
+@click.option("--sample_type", type=str, default=False, required=False)
+def main(read_directory, json_save_path, sample_type):
+    if sample_type == "random":
+        random_sample(read_directory, json_save_path)
+
+    if sample_type == "strattified":
+        strattified_sample(read_directory, json_save_path)
+
+
+if __name__ == "__main__":
+    main()

smart_tree/tests/dataloader.py

@@ -17,15 +17,18 @@
     config_name="test-dataloader",
 )
 def main(cfg: DictConfig):
-    train_dataloader = instantiate(
-        cfg.data_loader, dataset=instantiate(cfg.dataset, mode="train")
-    )
+    dataset = instantiate(cfg.dataset, mode="train")
+    train_dataloader = instantiate(cfg.data_loader, dataset)
 
     start_time = time.time()
     for data in tqdm(train_dataloader):
         pass
 
-    print(time.time() - start_time)
+        # times = [
+        #     s.elapsed_time(e) for s, e in zip(dataset.start_events, dataset.end_events)
+        # ]
+
+        # print(times)
 
 
 if __name__ == "__main__":