Add HeLiMOS (#25)

* Add helimos dataloader

* Update readme

* Update readme

* Update readme

* Take fix from 4DMOS

* Switch to polyscope

* Fix name

* Fix voxelization

* Add belief visualization

* Revert "Fix voxelization"

This reverts commit 3b2198f.

* Fix wrong voxel size

* Fix voxelization

* Fix voxelization

* Fix transformations in visualizer

* Readme

* Add link to paper

* Not needed

README.md

@@ -71,7 +71,7 @@ mapmos_pipeline --visualize /path/to/weights.ckpt /path/to/data
 <details>
 <summary>Want to evaluate with ground truth labels?</summary>
 
-Because these lables come in all shapes, you need to specify a dataloader. This is currently available for SemanticKITTI and NuScenes as well as our post-processed KITTI Tracking sequence 19 and Apollo sequences (see [Downloads](#downloads)).
+Because these labels come in all shapes, you need to specify a dataloader. This is currently available for SemanticKITTI, NuScenes, HeLiMOS, and our labeled KITTI Tracking sequence 19 and Apollo sequences (see [Downloads](#downloads)).
 
 </details>
 
@@ -108,6 +108,21 @@ The training log and checkpoints will be saved by default to the current working
 
 </details>
 
+## HeLiMOS
+We provide additional training and evaluation data for different sensor types in our [HeLiMOS paper](https://www.ipb.uni-bonn.de/pdfs/lim2024iros.pdf). To train on the HeLiMOS data, use the following commands:
+
+```shell
+python3 scripts/precache.py /path/to/HeLiMOS helimos /path/to/cache --config config/helimos/*_training.yaml
+python3 scripts/train.py /path/to/HeLiMOS helimos /path/to/cache --config config/helimos/*_training.yaml
+```
+
+by replacing the paths and the config file names. To evaluate for example on the Velodyne test data, run
+
+```shell
+mapmos_pipeline /path/to/weights.ckpt /path/to/HeLiMOS --dataloader helimos -s Velodyne/test.txt
+```
+
+Note that our sequence `-s` encodes both the sensor type `Velodyne` and split `test.txt`, just replace these with `Ouster`, `Aeva`, or `Avia` and/or `train.txt` or `val.txt` to run MapMOS on different sensors and/or splits. 
 
 ## Downloads
 You can download the post-processed and labeled [Apollo dataset](https://www.ipb.uni-bonn.de/html/projects/apollo_dataset/LiDAR-MOS.zip) and [KITTI Tracking sequence 19](https://www.ipb.uni-bonn.de/html/projects/kitti-tracking/post-processed/kitti-tracking.zip) from our website.

config/example.yaml

@@ -13,11 +13,11 @@ odometry:
 
 mos:
   voxel_size_mos: 0.1
+  delay_mos: 10
   max_range_mos: 50.0
   min_range_mos: 0.0
   voxel_size_belief: 0.25
   max_range_belief: 150
-  delay_belief: 10
 
 training:
   id: "experiment_id"

config/helimos/all_training.yaml

@@ -0,0 +1,12 @@
+training:
+  id: "helimos_all"
+  train:
+    - "Avia/train.txt"
+    - "Aeva/train.txt"
+    - "Velodyne/train.txt"
+    - "Ouster/train.txt"
+  val:
+    - "Avia/val.txt"
+    - "Aeva/val.txt"
+    - "Velodyne/val.txt"
+    - "Ouster/val.txt"

config/helimos/inference.yaml

@@ -0,0 +1,3 @@
+mos:
+  max_range_mos: -1.0
+  max_range_belief: 50

config/helimos/omni_training.yaml

@@ -0,0 +1,8 @@
+training:
+  id: "helimos_omni"
+  train:
+    - "Velodyne/train.txt"
+    - "Ouster/train.txt"
+  val:
+    - "Velodyne/val.txt"
+    - "Ouster/val.txt"

config/helimos/solid_training.yaml

@@ -0,0 +1,8 @@
+training:
+  id: "helimos_solid"
+  train:
+    - "Avia/train.txt"
+    - "Aeva/train.txt"
+  val:
+    - "Avia/val.txt"
+    - "Aeva/val.txt"

scripts/cache_to_ply.py

@@ -82,7 +82,7 @@ def cache_to_ply(
             dataloader=dataloader,
             data_dir=data,
             config=cfg,
-            sequences=sequence,
+            sequences=[sequence],
             cache_dir=cache_dir,
         ),
         batch_size=1,
@@ -103,36 +103,37 @@ def cache_to_ply(
     for idx, batch in enumerate(
         tqdm(data_iterable, desc="Writing data to ply", unit=" items", dynamic_ncols=True)
     ):
-        mask_scan = batch[:, 4] == idx
-        scan_points = batch[mask_scan, 1:4]
-        scan_labels = batch[mask_scan, 6]
-
-        map_points = batch[~mask_scan, 1:4]
-        map_timestamps = batch[~mask_scan, 5]
-        map_labels = batch[~mask_scan, 6]
-
-        min_time = torch.min(batch[:, 5])
-        max_time = torch.max(batch[:, 5])
-
-        pcd_scan = o3d.geometry.PointCloud(
-            o3d.utility.Vector3dVector(scan_points.numpy())
-        ).paint_uniform_color([0, 0, 1])
-        scan_colors = np.array(pcd_scan.colors)
-        scan_colors[scan_labels == 1] = [1, 0, 0]
-        pcd_scan.colors = o3d.utility.Vector3dVector(scan_colors)
-
-        pcd_map = o3d.geometry.PointCloud(
-            o3d.utility.Vector3dVector(map_points.numpy())
-        ).paint_uniform_color([0, 0, 0])
-        map_colors = np.array(pcd_map.colors)
-        map_timestamps_norm = (map_timestamps - min_time) / (max_time - min_time)
-        for i in range(len(map_colors)):
-            t = map_timestamps_norm[i]
-            map_colors[i, :] = [t, t, t]
-        map_colors[map_labels == 1] = [1, 0, 0]
-        pcd_map.colors = o3d.utility.Vector3dVector(map_colors)
-
-        o3d.io.write_point_cloud(os.path.join(path, f"{idx:06}.ply"), pcd_scan + pcd_map)
+        if len(batch) > 0:
+            mask_scan = batch[:, 4] == idx
+            scan_points = batch[mask_scan, 1:4]
+            scan_labels = batch[mask_scan, 6]
+
+            map_points = batch[~mask_scan, 1:4]
+            map_timestamps = batch[~mask_scan, 5]
+            map_labels = batch[~mask_scan, 6]
+
+            min_time = torch.min(batch[:, 5])
+            max_time = torch.max(batch[:, 5])
+
+            pcd_scan = o3d.geometry.PointCloud(
+                o3d.utility.Vector3dVector(scan_points.numpy())
+            ).paint_uniform_color([0, 0, 1])
+            scan_colors = np.array(pcd_scan.colors)
+            scan_colors[scan_labels == 1] = [1, 0, 0]
+            pcd_scan.colors = o3d.utility.Vector3dVector(scan_colors)
+
+            pcd_map = o3d.geometry.PointCloud(
+                o3d.utility.Vector3dVector(map_points.numpy())
+            ).paint_uniform_color([0, 0, 0])
+            map_colors = np.array(pcd_map.colors)
+            map_timestamps_norm = (map_timestamps - min_time) / (max_time - min_time)
+            for i in range(len(map_colors)):
+                t = map_timestamps_norm[i]
+                map_colors[i, :] = [t, t, t]
+            map_colors[map_labels == 1] = [1, 0, 0]
+            pcd_map.colors = o3d.utility.Vector3dVector(map_colors)
+
+            o3d.io.write_point_cloud(os.path.join(path, f"{idx:06}.ply"), pcd_scan + pcd_map)
 
 
 if __name__ == "__main__":

src/mapmos/config/config.py

@@ -39,11 +39,11 @@ class OdometryConfig(BaseModel):
 
 class MOSConfig(BaseModel):
     voxel_size_mos: float = 0.1
+    delay_mos: int = 10
     max_range_mos: float = 50.0
     min_range_mos: float = 0.0
     voxel_size_belief: float = 0.25
     max_range_belief: float = 150
-    delay_belief: int = 10
 
 
 class TrainingConfig(BaseModel):

src/mapmos/datasets/__init__.py

@@ -44,6 +44,7 @@ def sequence_dataloaders():
         "kitti_tracking",
         "nuscenes",
         "apollo",
+        "helimos",
     ]
 
 

src/mapmos/datasets/helimos.py

@@ -0,0 +1,137 @@
+# MIT License
+#
+# Copyright (c) 2023 Benedikt Mersch, Tiziano Guadagnino, Ignacio Vizzo, Cyrill Stachniss
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import glob
+import os
+import numpy as np
+
+
+class HeliMOSDataset:
+    def __init__(self, data_dir, sequence: str, *_, **__):
+        self.sequence_id = sequence.split("/")[0]
+        split_file = sequence.split("/")[1]
+        self.sequence_dir = os.path.join(data_dir, self.sequence_id)
+        self.scan_dir = os.path.join(self.sequence_dir, "velodyne/")
+
+        self.scan_files = sorted(glob.glob(self.scan_dir + "*.bin"))
+        self.calibration = self.read_calib_file(os.path.join(self.sequence_dir, "calib.txt"))
+
+        # Load GT Poses (if available)
+        self.poses_fn = os.path.join(self.sequence_dir, "poses.txt")
+        if os.path.exists(self.poses_fn):
+            self.gt_poses = self.load_poses(self.poses_fn)
+
+        # No correction
+        self.correct_kitti_scan = lambda frame: frame
+
+        # Load labels
+        self.label_dir = os.path.join(self.sequence_dir, "labels/")
+        label_files = sorted(glob.glob(self.label_dir + "*.label"))
+
+        # Get labels for train/val split if desired
+        label_indices = np.loadtxt(os.path.join(data_dir, split_file), dtype=int).tolist()
+
+        # Filter based on split if desired
+        getIndex = lambda filename: int(os.path.basename(filename).split(".label")[0])
+        self.dict_label_files = {
+            getIndex(filename): filename
+            for filename in label_files
+            if getIndex(filename) in label_indices
+        }
+
+    def __getitem__(self, idx):
+        points = self.scans(idx)
+        timestamps = np.zeros(len(points))
+        labels = (
+            self.read_labels(self.dict_label_files[idx])
+            if idx in self.dict_label_files.keys()
+            else np.full(len(points), -1, dtype=np.int32)
+        )
+        return points, timestamps, labels
+
+    def __len__(self):
+        return len(self.scan_files)
+
+    def scans(self, idx):
+        return self.read_point_cloud(self.scan_files[idx])
+
+    def apply_calibration(self, poses: np.ndarray) -> np.ndarray:
+        """Converts from Velodyne to Camera Frame"""
+        Tr = np.eye(4, dtype=np.float64)
+        Tr[:3, :4] = self.calibration["Tr"].reshape(3, 4)
+        return Tr @ poses @ np.linalg.inv(Tr)
+
+    def read_point_cloud(self, scan_file: str):
+        points = np.fromfile(scan_file, dtype=np.float32).reshape((-1, 4))[:, :3].astype(np.float64)
+        return points
+
+    def load_poses(self, poses_file):
+        def _lidar_pose_gt(poses_gt):
+            _tr = self.calibration["Tr"].reshape(3, 4)
+            tr = np.eye(4, dtype=np.float64)
+            tr[:3, :4] = _tr
+            left = np.einsum("...ij,...jk->...ik", np.linalg.inv(tr), poses_gt)
+            right = np.einsum("...ij,...jk->...ik", left, tr)
+            return right
+
+        poses = np.loadtxt(poses_file, delimiter=" ")
+        n = poses.shape[0]
+        poses = np.concatenate(
+            (poses, np.zeros((n, 3), dtype=np.float32), np.ones((n, 1), dtype=np.float32)), axis=1
+        )
+        poses = poses.reshape((n, 4, 4))  # [N, 4, 4]
+
+        # Ensure rotations are SO3
+        rotations = poses[:, :3, :3]
+        U, _, Vh = np.linalg.svd(rotations)
+        poses[:, :3, :3] = U @ Vh
+
+        return _lidar_pose_gt(poses)
+
+    @staticmethod
+    def read_calib_file(file_path: str) -> dict:
+        calib_dict = {}
+        with open(file_path, "r") as calib_file:
+            for line in calib_file.readlines():
+                tokens = line.split(" ")
+                if tokens[0] == "calib_time:":
+                    continue
+                # Only read with float data
+                if len(tokens) > 0:
+                    values = [float(token) for token in tokens[1:]]
+                    values = np.array(values, dtype=np.float32)
+
+                    # The format in KITTI's file is <key>: <f1> <f2> <f3> ...\n -> Remove the ':'
+                    key = tokens[0][:-1]
+                    calib_dict[key] = values
+        return calib_dict
+
+    def read_labels(self, filename):
+        """Load moving object labels from .label file"""
+        orig_labels = np.fromfile(filename, dtype=np.int32).reshape((-1))
+        orig_labels = orig_labels & 0xFFFF  # Mask semantics in lower half
+
+        labels = np.zeros_like(orig_labels)
+        labels[orig_labels <= 1] = -1  # Unlabeled (0), outlier (1)
+        labels[orig_labels > 250] = 1  # Moving
+        labels = labels.astype(dtype=np.int32).reshape(-1)
+        return labels

src/mapmos/datasets/mapmos_dataset.py

@@ -105,6 +105,7 @@ def setup(self, stage=None):
             shuffle=self.shuffle,
             num_workers=self.config.training.num_workers,
             pin_memory=True,
+            persistent_workers=True,
             drop_last=False,
             timeout=0,
         )
@@ -118,6 +119,7 @@ def setup(self, stage=None):
             shuffle=False,
             num_workers=self.config.training.num_workers,
             pin_memory=True,
+            persistent_workers=True,
             drop_last=False,
             timeout=0,
         )
@@ -148,6 +150,7 @@ def __init__(
     ):
         self.config = config
         self.sequences = sequences
+        self._print = False
 
         # Cache
         if cache_dir is not None:
@@ -208,14 +211,18 @@ def get_scan_and_map(
         local frame to allow for efficient cropping (sample point does not change). We use the
         VoxelHashMap to keep track of the GT labels for map points.
         """
+        if not self._print:
+            print("*****Caching now*****")
+            self._print = True
+
         scan_points, timestamps, scan_labels = self.datasets[sequence][scan_index]
 
         # Only consider valid points
         valid_mask = scan_labels != -1
         scan_points = scan_points[valid_mask]
         scan_labels = scan_labels[valid_mask]
 
-        if self.sequence != sequence:
+        if self.sequence != sequence or len(scan_points) == 0:
             data_config = DataConfig().model_validate(data_config_dict)
             odometry_config = OdometryConfig().model_validate(odometry_config_dict)
 

src/mapmos/mapmos_net.py

@@ -80,7 +80,8 @@ def forward(self, coordinates: torch.Tensor, indices: torch.Tensor):
         logits = out.features.reshape(-1)
         return logits
 
-    def to_label(self, logits):
+    @staticmethod
+    def to_label(logits):
         labels = copy.deepcopy(logits)
         mask = logits > 0
         labels[mask] = 1.0

src/mapmos/pipeline.py

@@ -78,7 +78,7 @@ def __init__(
             voxel_size=self.config.mos.voxel_size_belief,
             max_distance=self.config.mos.max_range_belief,
         )
-        self.buffer = deque(maxlen=self.config.mos.delay_belief)
+        self.buffer = deque(maxlen=self.config.mos.delay_mos)
 
         # Results
         self.results = MOSPipelineResults()