Upgrade LiDAR segmentation

.gitignore

@@ -7,4 +7,6 @@ __pycache__
 dist
 poetry.lock
 
-local/
+local/
+
+.DS_Store

README.md

@@ -41,8 +41,8 @@ Now, we're excited to introduce ***DetectionMetrics v2***! While retaining the f
   </tr>
   <tr>
     <td>LiDAR</td>
-    <td>Rellis3D, GOOSE, custom GAIA format</td>
-    <td>PyTorch (tested with RandLA-Net and KPConv from <a href="https://github.com/isl-org/Open3D-ML">Open3D-ML</a>)</td>  </tr>
+    <td>Rellis3D, GOOSE, WildScenes, custom GAIA format</td>
+    <td>PyTorch (tested with <a href="https://github.com/isl-org/Open3D-ML">Open3D-ML</a>, <a href="https://github.com/open-mmlab/mmdetection3d">mmdetection3d</a>, <a href="https://github.com/dvlab-research/SphereFormer">SphereFormer</a>, and <a href="https://github.com/FengZicai/LSK3DNet">LSK3DNet</a> models)</td>  </tr>
   <tr>
     <td>Object detection</td>
     <td>Image</td>

detectionmetrics/cli/batch.py

@@ -1,4 +1,4 @@
-from itertools import product
+from itertools import product, chain
 from glob import glob
 import os
 
@@ -30,9 +30,19 @@ def batch(command, jobs_cfg):
     for model_cfg in jobs_cfg["model"]:
 
         model_path = model_cfg["path"]
-        model_paths = glob(model_path) if model_cfg["path_is_pattern"] else [model_path]
-        assert model_paths, f"No files found for pattern {model_cfg['path']}"
+        is_pattern = model_cfg.get("path_is_pattern", False)
+        if isinstance(model_path, list):
+            if is_pattern:
+                model_paths = list(chain.from_iterable(glob(p) for p in model_path))
+            else:
+                model_paths = model_path
+        else:
+            model_paths = glob(model_path) if is_pattern else [model_path]
+
+        if not model_paths:
+            raise FileNotFoundError(f"No files found for path/pattern: {model_path}")
 
+        print(f"Found {len(model_paths)} model(s) for pattern: {model_path}")
         for new_path in model_paths:
             assert os.path.exists(new_path), f"File or directory {new_path} not found"
 
@@ -41,7 +51,8 @@ def batch(command, jobs_cfg):
             if os.path.isfile(new_path):
                 new_model_id, _ = os.path.splitext(new_model_id)
 
-            new_model_cfg = model_cfg | {
+            new_model_cfg = {
+                **model_cfg,
                 "path": new_path,
                 "id": f"{model_cfg['id']}-{new_model_id.replace('-', '_')}",
             }
@@ -102,9 +113,20 @@ def batch(command, jobs_cfg):
                     "model": model_cfg["path"],
                     "model_ontology": model_cfg["ontology"],
                     "model_cfg": model_cfg["cfg"],
-                    # "image_size": model_cfg.get("image_size", None),
                 }
             )
+
+            if command == "computational_cost":
+                if jobs_cfg["input_type"] == "image":
+                    params["image_size"] = model_cfg.get("image_size", [512, 512])
+                elif jobs_cfg["input_type"] == "lidar":
+                    params["point_cloud_range"] = model_cfg.get(
+                        "point_cloud_range", [-50, -50, -5, 50, 50, 5]
+                    )
+                    params["num_points"] = model_cfg.get("num_points", 100000)
+                else:
+                    raise ValueError(f"Unknown input type: {jobs_cfg['input_type']}")
+
             if has_dataset:
                 dataset_cfg = job_components[1]
                 params.update(

detectionmetrics/cli/computational_cost.py

@@ -1,7 +1,6 @@
 import click
 
 from detectionmetrics import cli
-from detectionmetrics.utils.io import read_json
 
 
 @click.command(name="computational_cost", help="Estimate model computational cost")
@@ -12,9 +11,7 @@
 # model
 @click.option(
     "--model_format",
-    type=click.Choice(
-        ["torch", "tensorflow", "tensorflow_explicit"], case_sensitive=False
-    ),
+    type=click.Choice(["torch", "tensorflow"], case_sensitive=False),
     show_default=True,
     default="torch",
     help="Trained model format",
@@ -39,14 +36,35 @@
 )
 @click.option(
     "--image_size",
-    type=(int, int),
+    nargs=2,
+    type=int,
     required=False,
-    help="Dummy image size used for computational cost estimation",
+    help="Dummy image size. Should be provided as two integers: width height",
+)
+@click.option(
+    "--point_cloud_range",
+    nargs=6,
+    type=int,
+    required=False,
+    help="Dummy point cloud range (meters). Should be provided as six integers: x_min y_min z_min x_max y_max z_max",
+)
+@click.option(
+    "--num_points",
+    type=int,
+    required=False,
+    help="Number of points for the dummy point cloud (uniformly sampled)",
+)
+@click.option(
+    "--has_intensity",
+    is_flag=True,
+    default=False,
+    help="Whether the dummy point cloud has intensity values",
 )
 # output
 @click.option(
     "--out_fname",
     type=click.Path(writable=True),
+    required=True,
     help="CSV file where the computational cost estimation results will be stored",
 )
 def computational_cost(
@@ -57,23 +75,46 @@ def computational_cost(
     model_ontology,
     model_cfg,
     image_size,
+    point_cloud_range,
+    num_points,
+    has_intensity,
     out_fname,
 ):
     """Estimate model computational cost"""
-
-    if image_size is None:
-        parsed_model_cfg = read_json(model_cfg)
-        if "image_size" in parsed_model_cfg:
-            image_size = parsed_model_cfg["image_size"]
-        else:
+    if input_type == "image":
+        if image_size is None:
+            raise ValueError("Image size must be provided for image models")
+        if point_cloud_range is not None or num_points is not None:
+            raise ValueError(
+                "Point cloud range and number of points cannot be provided for image models"
+            )
+        if has_intensity:
+            raise ValueError("Intensity flag cannot be set for image models")
+        params = {"image_size": image_size}
+    elif input_type == "lidar":
+        if point_cloud_range is None or num_points is None:
             raise ValueError(
-                "Image size must be provided either as an argument or in the model configuration file"
+                "Point cloud range and number of points must be provided for lidar models"
             )
+        if image_size is not None:
+            raise ValueError("Image size cannot be provided for lidar models")
+
+        params = {
+            "point_cloud_range": point_cloud_range,
+            "num_points": num_points,
+            "has_intensity": has_intensity,
+        }
+    else:
+        raise ValueError(f"Unknown input type: {input_type}")
 
     model = cli.get_model(
         task, input_type, model_format, model, model_ontology, model_cfg
     )
-    results = model.get_computational_cost(image_size)
+    results = model.get_computational_cost(**params)
     results.to_csv(out_fname)
 
     return results
+
+
+if __name__ == "__main__":
+    computational_cost()

detectionmetrics/cli/evaluate.py

@@ -25,7 +25,7 @@ def parse_split(ctx, param, value):
 @click.option(
     "--model_format",
     type=click.Choice(
-        ["torch", "tensorflow", "tensorflow_explicit"], case_sensitive=False
+        ["torch", "tensorflow"], case_sensitive=False
     ),
     show_default=True,
     default="torch",
@@ -197,3 +197,7 @@ def evaluate(
     results.to_csv(out_fname)
 
     return results
+
+
+if __name__ == "__main__":
+    evaluate()

detectionmetrics/datasets/gaia.py

@@ -23,8 +23,15 @@ def build_dataset(dataset_fname: str) -> Tuple[pd.DataFrame, str, dict]:
     dataset_dir = os.path.dirname(dataset_fname)
 
     # Read ontology file
-    ontology_fname = dataset.attrs["ontology_fname"]
-    ontology = uio.read_json(os.path.join(dataset_dir, ontology_fname))
+    try:
+        ontology_fname = dataset.attrs["ontology_fname"]
+    except KeyError:
+        ontology_fname =  "ontology.json"
+
+    ontology_fname = os.path.join(dataset_dir, ontology_fname)
+    assert os.path.isfile(ontology_fname), "Ontology file not found"
+
+    ontology = uio.read_json(ontology_fname)
     for name, data in ontology.items():
         ontology[name]["rgb"] = tuple(data["rgb"])
 

detectionmetrics/datasets/goose.py

@@ -16,6 +16,7 @@ def build_dataset(
     train_dataset_dir: Optional[str] = None,
     val_dataset_dir: Optional[str] = None,
     test_dataset_dir: Optional[str] = None,
+    is_goose_ex: bool = False,
 ) -> Tuple[dict, dict]:
     """Build dataset and ontology dictionaries from GOOSE dataset structure
 
@@ -31,6 +32,8 @@ def build_dataset(
     :type val_dataset_dir: str, optional
     :param test_dataset_dir: Directory containing test data, defaults to None
     :type test_dataset_dir: str, optional
+    :param is_goose_ex: Whether the dataset is GOOSE Ex or GOOSE, defaults to False
+    :type is_goose_ex: bool, optional
     :return: Dataset and onotology
     :rtype: Tuple[dict, dict]
     """
@@ -66,13 +69,23 @@ def build_dataset(
         train_data = os.path.join(dataset_dir, f"{data_type}/{split}/*/*_{data_suffix}")
         for data_fname in glob(train_data):
             sample_dir, sample_base_name = os.path.split(data_fname)
-            sample_base_name = sample_base_name.split("__")[-1]
+
+            # GOOSE Ex uses a different label file naming convention
+            if is_goose_ex:
+                sample_base_name = "sequence" + sample_base_name.split("_sequence")[-1]
+            else:
+                sample_base_name = sample_base_name.split("__")[-1]
+
             sample_base_name = sample_base_name.split("_" + data_suffix)[0]
 
             scene = os.path.split(sample_dir)[-1]
             sample_name = f"{scene}-{sample_base_name}"
 
-            label_base_name = f"{scene}__{sample_base_name}_{label_suffix}"
+            if is_goose_ex:
+                label_base_name = f"{scene}_{sample_base_name}_{label_suffix}"
+            else:
+                label_base_name = f"{scene}__{sample_base_name}_{label_suffix}"
+
             label_fname = os.path.join(
                 dataset_dir, "labels", split, scene, label_base_name
             )
@@ -131,31 +144,35 @@ def __init__(
 class GOOSELiDARSegmentationDataset(dm_segmentation_dataset.LiDARSegmentationDataset):
     """Specific class for GOOSE-styled LiDAR segmentation datasets. All data can be
     downloaded from the official webpage (https://goose-dataset.de):
-    train -> https://goose-dataset.de/storage/goose_3d_train.zip
-    val   -> https://goose-dataset.de/storage/goose_3d_val.zip
-    test  -> https://goose-dataset.de/storage/goose_3d_test.zip
+    train -> https://goose-dataset.de/storage/gooseEx_3d_train.zip
+    val   -> https://goose-dataset.de/storage/gooseEx_3d_val.zip
+    test  -> https://goose-dataset.de/storage/gooseEx_3d_test.zip
 
     :param train_dataset_dir: Directory containing training data
     :type train_dataset_dir: str
     :param val_dataset_dir: Directory containing validation data, defaults to None
     :type val_dataset_dir: str, optional
     :param test_dataset_dir: Directory containing test data, defaults to None
     :type test_dataset_dir: str, optional
+    :param is_goose_ex: Whether the dataset is GOOSE Ex or GOOSE, defaults to False
+    :type is_goose_ex: bool, optional
     """
 
     def __init__(
         self,
         train_dataset_dir: Optional[str] = None,
         val_dataset_dir: Optional[str] = None,
         test_dataset_dir: Optional[str] = None,
+        is_goose_ex: bool = False,
     ):
         dataset, ontology = build_dataset(
             "lidar",
-            "vls128.bin",
+            "pcl.bin" if is_goose_ex else "vls128.bin",
             "goose.label",
             train_dataset_dir,
             val_dataset_dir,
             test_dataset_dir,
+            is_goose_ex=is_goose_ex,
         )
 
         # Convert to Pandas