implemented allignment of point cloud with gravity

digiforest_analysis/notebooks/tree_segmentation.ipynb

@@ -12,7 +12,9 @@
     "import open3d as o3d\n",
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
+    "\n",
     "from digiforest_analysis.utils.timing import Timer\n",
+    "from digiforest_analysis.utils.io import load, apply_header_transform\n",
     "\n",
     "timer = Timer()"
    ]
@@ -32,8 +34,9 @@
     "        colors = cmap(labels / (max_label if max_label > 0 else 1))\n",
     "        colors[labels < 0] = 0\n",
     "        cloud.colors = o3d.utility.Vector3dVector(colors[:, :3])\n",
+    "    cosy = o3d.geometry.TriangleMesh.create_coordinate_frame(size=20, origin=[0, 0, 0])\n",
     "    o3d.visualization.draw_geometries(\n",
-    "        [cloud],\n",
+    "        [cloud, cosy],\n",
     "        # zoom=0.5,\n",
     "        # front=[0.79, 0.02, 0.60],\n",
     "        # lookat=[2.61, 2.04, 1.53],\n",
@@ -52,8 +55,8 @@
    "outputs": [],
    "source": [
     "# load data\n",
-    "pcd_file = \"/home/matias/vilens_slam_data/test_data/forest_cloud.pcd\"\n",
-    "pcd = o3d.t.io.read_point_cloud(pcd_file)\n",
+    "pcd_file = \"/home/ori/logs/logs_evo_finland/exp01/2023-05-01-14-01-05-exp01/payload_clouds/cloud_1682946124_761436000.pcd\"\n",
+    "pcd, header = load(pcd_file, binary=True)\n",
     "print(pcd)\n",
     "visualize(pcd.to_legacy(), None, \"original_cloud\")"
    ]
@@ -126,10 +129,8 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "86d97afb-49e8-4b97-b687-faf106076c52",
-   "metadata": {
-    "tags": []
-   },
+   "id": "f7f2b4b1",
+   "metadata": {},
    "outputs": [],
    "source": [
     "# DBSCAN (sklearn)\n",
@@ -466,6 +467,23 @@
     "        ss = \"cloud_cluster_\" + str(j) + \".pcd\"\n",
     "        pcl.save(cloud_cluster, ss)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "31ba9e1c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Test header transforms. The third pc should look the same as the first one\n",
+    "pcd_file = \"/home/ori/logs/logs_evo_finland/exp01/2023-05-01-14-01-05-exp01/payload_clouds/cloud_1682946124_761436000.pcd\"\n",
+    "pcd, header = load(pcd_file, transform_to_world=False)\n",
+    "visualize(pcd.to_legacy(), None, \"No trafo\")\n",
+    "pcd, header = load(pcd_file, transform_to_world=True)\n",
+    "visualize(pcd.to_legacy(), None, \"Trafo\")\n",
+    "apply_header_transform(pcd, header, inverse=True)\n",
+    "visualize(pcd.to_legacy(), None, \"Back Trafo\")"
+   ]
   }
  ],
  "metadata": {
@@ -484,7 +502,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.8.18"
   }
  },
  "nbformat": 4,

digiforest_analysis/src/digiforest_analysis/pipeline.py

@@ -91,6 +91,10 @@ def setup_output(self, output_dir):
 
     def save_cloud(self, cloud, label="default"):
         filename = Path(self._output_dir).joinpath(label + self._cloud_format)
+
+        # retransform cloud
+        cloud = io.apply_header_transform(cloud, self._header, inverse=True)
+
         io.write(cloud, self._header, str(filename))
 
     def save_trees(self, trees, label="trees"):
@@ -102,6 +106,11 @@ def save_trees(self, trees, label="trees"):
             i = tree["info"]["id"]
             tree_cloud = tree["cloud"]
 
+            # retransform cloud
+            tree_cloud = io.apply_header_transform(
+                tree_cloud, self._header, inverse=True
+            )
+
             # Write cloud
             tree_cloud_filename = Path(
                 save_folder, f"tree_cloud_{i:04}{self._cloud_format}"

digiforest_analysis/src/digiforest_analysis/tasks/ground_segmentation.py

@@ -238,6 +238,10 @@ def debug_visualizations(self, ground_cloud, forest_cloud):
 
     ground_cloud, forest_cloud = app.process(cloud=cloud)
 
+    # retransform clouds
+    ground_cloud = io.apply_header_transform(ground_cloud, header, inverse=True)
+    forest_cloud = io.apply_header_transform(forest_cloud, header, inverse=True)
+
     # Write clouds
     header_fix = {"VIEWPOINT": header["VIEWPOINT"]}
     io.write(ground_cloud, header_fix, os.path.join(sys.argv[2], "ground_cloud.pcd"))

digiforest_analysis/src/digiforest_analysis/tasks/tree_analysis.py

@@ -241,6 +241,9 @@ def debug_visualizations(self, trees, filtered_trees):
         i = tree["info"]["id"]
         cloud = tree["cloud"]
 
+        # retransform cloud
+        cloud = io.apply_header_transform(cloud, header, inverse=True)
+
         tree_name = f"tree_cloud_{i:04}.pcd"
         tree_cloud_filename = os.path.join(sys.argv[2], tree_name)
         header_fix = {"VIEWPOINT": header["VIEWPOINT"]}

digiforest_analysis/src/digiforest_analysis/tasks/tree_segmentation.py

@@ -339,6 +339,9 @@ def debug_visualizations(self, cloud, clusters):
         # Tree height normalize
         tree_cloud = tree["cloud"]
 
+        # retransform cloud
+        tree_cloud = io.apply_header_transform(tree_cloud, header, inverse=True)
+
         # shift to zero (debugging)
         # z_shift = cloud.point.positions[:, 2].min()
         # cloud.point.positions[:, 2] = cloud.point.positions[:, 2] - z_shift

digiforest_analysis/src/digiforest_analysis/utils/io.py

@@ -6,18 +6,35 @@
 import numpy as np
 
 
-def load(filename: str, binary=True):
+def load(filename: str, binary=True, transform_to_world=True):
     path = Path(filename)
     file_format = path.suffix
     cloud = o3d.t.io.read_point_cloud(str(path))
 
     header = load_header(filename, file_format, binary=binary, cloud=cloud)
+    if transform_to_world:
+        cloud = apply_header_transform(cloud, header, inverse=False)
+        return cloud, header
     if "offset" in header:
         cloud = cloud.translate(-header["offset"])
-
     return cloud, header
 
 
+def apply_header_transform(cloud, header: dict, inverse: bool = False):
+    assert "VIEWPOINT" in header, "No viewpoint in header. Cannot apply transform"
+    header_data = [float(x) for x in header["VIEWPOINT"]]
+    location = np.array(header_data[:3])
+    rotation = np.array(header_data[3:])
+    R = o3d.geometry.TriangleMesh.get_rotation_matrix_from_quaternion(rotation)
+    if inverse:
+        cloud.rotate(R.T, center=[0, 0, 0])
+        cloud = cloud.translate(location)
+    else:
+        cloud.rotate(R, center=location)
+        cloud = cloud.translate(-location)
+    return cloud
+
+
 def write(cloud, header, filename):
     write_open3d(cloud, header, filename)