Merge pull request #253 from JdeRobot/dph/confusion_matrix

Add confusion matrix

Author: dpascualhe

detectionmetrics/cli/batch.py

@@ -80,27 +80,32 @@ def batch(command, jobs_cfg):
         pbar.set_description(f"Processing {job_id}")
 
         ctx = click.get_current_context()
-        result = ctx.invoke(
-            cli_registry[command],
-            task=jobs_cfg["task"],
-            input_type=jobs_cfg["input_type"],
-            model_format=model_cfg["format"],
-            model=model_cfg["path"],
-            model_ontology=model_cfg["ontology"],
-            model_cfg=model_cfg["cfg"],
-            dataset_format=dataset_cfg["format"],
-            dataset_fname=dataset_cfg.get("fname", None),
-            dataset_dir=dataset_cfg.get("dir", None),
-            train_dataset_dir=dataset_cfg.get("train_dir", None),
-            val_dataset_dir=dataset_cfg.get("val_dir", None),
-            test_dataset_dir=dataset_cfg.get("test_dir", None),
-            data_suffix=dataset_cfg.get("data_suffix", None),
-            label_suffix=dataset_cfg.get("label_suffix", None),
-            dataset_ontology=dataset_cfg.get("ontology", None),
-            split=dataset_cfg["split"],
-            ontology_translation=jobs_cfg.get("ontology_translation", None),
-            out_fname=job_out_fname,
-        )
+        try:
+            result = ctx.invoke(
+                cli_registry[command],
+                task=jobs_cfg["task"],
+                input_type=jobs_cfg["input_type"],
+                model_format=model_cfg["format"],
+                model=model_cfg["path"],
+                model_ontology=model_cfg["ontology"],
+                model_cfg=model_cfg["cfg"],
+                dataset_format=dataset_cfg["format"],
+                dataset_fname=dataset_cfg.get("fname", None),
+                dataset_dir=dataset_cfg.get("dir", None),
+                split_dir=dataset_cfg.get("split_dir", None),
+                train_dataset_dir=dataset_cfg.get("train_dir", None),
+                val_dataset_dir=dataset_cfg.get("val_dir", None),
+                test_dataset_dir=dataset_cfg.get("test_dir", None),
+                data_suffix=dataset_cfg.get("data_suffix", None),
+                label_suffix=dataset_cfg.get("label_suffix", None),
+                dataset_ontology=dataset_cfg.get("ontology", None),
+                split=dataset_cfg["split"],
+                ontology_translation=jobs_cfg.get("ontology_translation", None),
+                out_fname=job_out_fname,
+            )
+        except Exception as e:
+            print(f"Error processing job {job_id}: {e}")
+            continue
 
         # We assume that the command returns the results as a pandas DataFrame
         result["job_id"] = job_id

detectionmetrics/cli/evaluate.py

@@ -21,6 +21,7 @@ def get_dataset(
     dataset_format,
     dataset_fname,
     dataset_dir,
+    split_dir,
     train_dataset_dir,
     val_dataset_dir,
     test_dataset_dir,
@@ -33,8 +34,13 @@ def get_dataset(
     if dataset_format == "gaia" and dataset_fname is None:
         raise ValueError("--dataset is required for 'gaia' format")
 
-    elif dataset_format == "rellis3d" and dataset_dir is None:
-        raise ValueError("--dataset_dir is required for 'rellis3d' format")
+    elif dataset_format == "rellis3d":
+        if dataset_dir is None:
+            raise ValueError("--dataset_dir is required for 'rellis3d' format")
+        if split_dir is None:
+            raise ValueError("--split_dir is required for 'rellis3d' format")
+        if ontology is None:
+            raise ValueError("--dataset_ontology is required for 'rellis3d' format")
 
     elif dataset_format in ["goose", "generic"]:
         if split == "train" and train_dataset_dir is None:
@@ -62,7 +68,11 @@ def get_dataset(
     if dataset_format == "gaia":
         dataset_args = {"dataset_fname": dataset_fname}
     elif dataset_format == "rellis3d":
-        dataset_args = {"dataset_dir": dataset_dir}
+        dataset_args = {
+            "dataset_dir": dataset_dir,
+            "split_dir": split_dir,
+            "ontology_fname": ontology,
+        }
     elif dataset_format == "goose":
         dataset_args = {
             "train_dataset_dir": train_dataset_dir,
@@ -99,7 +109,9 @@ def get_dataset(
 # model
 @click.option(
     "--model_format",
-    type=click.Choice(["torch", "tensorflow"], case_sensitive=False),
+    type=click.Choice(
+        ["torch", "tensorflow", "tensorflow_explicit"], case_sensitive=False
+    ),
     show_default=True,
     default="torch",
     help="Trained model format",
@@ -140,6 +152,11 @@ def get_dataset(
     type=click.Path(exists=True, file_okay=False, dir_okay=True),
     help="Dataset directory (used for 'Rellis3D' format)",
 )
+@click.option(
+    "--split_dir",
+    type=click.Path(exists=True, file_okay=False, dir_okay=True),
+    help="Directory containing .lst split files (used for 'Rellis3D' format)",
+)
 @click.option(
     "--train_dataset_dir",
     type=click.Path(exists=True, file_okay=False, dir_okay=True),
@@ -168,7 +185,7 @@ def get_dataset(
 @click.option(
     "--dataset_ontology",
     type=click.Path(exists=True, dir_okay=False),
-    help="JSON file containing dataset ontology (used for 'Generic' format)",
+    help="JSON containing dataset ontology (used for 'Generic' and 'Rellis3D' formats)",
 )
 @click.option(
     "--split",
@@ -199,6 +216,7 @@ def evaluate(
     dataset_format,
     dataset_fname,
     dataset_dir,
+    split_dir,
     train_dataset_dir,
     val_dataset_dir,
     test_dataset_dir,
@@ -218,6 +236,7 @@ def evaluate(
         dataset_format,
         dataset_fname,
         dataset_dir,
+        split_dir,
         train_dataset_dir,
         val_dataset_dir,
         test_dataset_dir,

detectionmetrics/datasets/__init__.py

@@ -23,6 +23,6 @@
     "generic_lidar_segmentation": GenericLiDARSegmentationDataset,
     "goose_image_segmentation": GOOSEImageSegmentationDataset,
     "goose_lidar_segmentation": GOOSELiDARSegmentationDataset,
-    "rellis_image_segmentation": Rellis3DImageSegmentationDataset,
-    "rellis_lidar_segmentation": Rellis3DLiDARSegmentationDataset,
+    "rellis3d_image_segmentation": Rellis3DImageSegmentationDataset,
+    "rellis3d_lidar_segmentation": Rellis3DLiDARSegmentationDataset,
 }

detectionmetrics/models/__init__.py

@@ -1,14 +1,22 @@
-from detectionmetrics.models.torch import (
-    TorchImageSegmentationModel,
-    TorchLiDARSegmentationModel,
-)
-from detectionmetrics.models.tensorflow import (
-    TensorflowImageSegmentationModel,
-)
-
-
-REGISTRY = {
-    "torch_image_segmentation": TorchImageSegmentationModel,
-    "torch_lidar_segmentation": TorchLiDARSegmentationModel,
-    "tensorflow_image_segmentation": TensorflowImageSegmentationModel,
-}
+REGISTRY = {}
+
+try:
+    from detectionmetrics.models.torch import (
+        TorchImageSegmentationModel,
+        TorchLiDARSegmentationModel,
+    )
+
+    REGISTRY["torch_image_segmentation"] = TorchImageSegmentationModel
+    REGISTRY["torch_lidar_segmentation"] = TorchLiDARSegmentationModel
+except ImportError:
+    print("Torch not available")
+
+try:
+    from detectionmetrics.models.tensorflow import TensorflowImageSegmentationModel
+
+    REGISTRY["tensorflow_image_segmentation"] = TensorflowImageSegmentationModel
+except ImportError:
+    print("Tensorflow not available")
+
+if not REGISTRY:
+    raise Exception("No valid deep learning framework found")

detectionmetrics/models/tensorflow.py

@@ -157,7 +157,7 @@ def inference(self, image: Image.Image) -> Image.Image:
         # TODO: check if this is consistent across different models
         result = self.model.signatures["serving_default"](tensor)
         if isinstance(result, dict):
-            result = result["output_0"]
+            result = list(result.values())[0]
 
         return self.t_out(result)
 
@@ -195,18 +195,19 @@ def eval(
         # Init metrics
         results = {}
         iou = um.IoU(self.n_classes)
-        acc = um.Accuracy(self.n_classes)
+        cm = um.ConfusionMatrix(self.n_classes)
 
         # Evaluation loop
         pbar = tqdm(dataset.dataset)
         for image, label in pbar:
+            # TODO: check if this is consistent across different models
             pred = self.model.signatures["serving_default"](image)
             if isinstance(pred, dict):
-                pred = pred["output_0"]
+                pred = list(pred.values())[0]
 
             label = tf.squeeze(label, axis=3)
             pred = tf.argmax(pred, axis=3)
-            acc.update(pred.numpy(), label.numpy())
+            cm.update(pred.numpy(), label.numpy())
 
             pred = tf.one_hot(pred, self.n_classes)
             pred = tf.transpose(pred, perm=[0, 3, 1, 2])
@@ -218,7 +219,7 @@ def eval(
 
         # Get metrics results
         iou_per_class, iou = iou.compute()
-        acc_per_class, acc = acc.compute()
+        acc_per_class, acc = cm.get_accuracy()
         iou_per_class = [float(n) for n in iou_per_class]
         acc_per_class = [float(n) for n in acc_per_class]
 

detectionmetrics/models/torch.py

@@ -12,8 +12,6 @@
 from detectionmetrics.datasets import dataset as dm_dataset
 from detectionmetrics.models import model as dm_model
 from detectionmetrics.models import torch_model_utils as tmu
-import detectionmetrics.utils.conversion as uc
-import detectionmetrics.utils.io as uio
 import detectionmetrics.utils.lidar as ul
 import detectionmetrics.utils.metrics as um
 
@@ -224,11 +222,24 @@ def __init__(self, model_fname: str, model_cfg: str, ontology_fname: str):
                 )
             ]
 
-        self.transform_input += [
-            transforms.ToImage(),
-            transforms.ToDtype(torch.float32, scale=True),
-        ]
-        self.transform_label += [transforms.ToImage(), transforms.ToDtype(torch.int64)]
+        try:
+            self.transform_input += [
+                transforms.ToImage(),
+                transforms.ToDtype(torch.float32, scale=True),
+            ]
+            self.transform_label += [
+                transforms.ToImage(),
+                transforms.ToDtype(torch.int64),
+            ]
+        except AttributeError:  # adapt for older versions of torchvision transforms v2
+            self.transform_input += [
+                transforms.ToImageTensor(),
+                transforms.ConvertDtype(torch.float32),
+            ]
+            self.transform_label += [
+                transforms.ToImageTensor(),
+                transforms.ToDtype(torch.int64),
+            ]
 
         if "normalization" in self.model_cfg:
             self.transform_input += [
@@ -311,7 +322,7 @@ def eval(
         # Init metrics
         results = {}
         iou = um.IoU(self.n_classes)
-        acc = um.Accuracy(self.n_classes)
+        cm = um.ConfusionMatrix(self.n_classes)
 
         # Evaluation loop
         with torch.no_grad():
@@ -335,13 +346,13 @@ def eval(
                 if lut_ontology is not None:
                     label = lut_ontology[label]
 
-                # Prepare data and update accuracy
+                # Prepare data and update confusion matrix
                 label = label.squeeze(dim=1).cpu()
                 pred = torch.argmax(pred, axis=1).cpu()
                 if valid_mask is not None:
                     valid_mask = valid_mask.squeeze(dim=1).cpu()
 
-                acc.update(
+                cm.update(
                     pred.numpy(),
                     label.numpy(),
                     valid_mask.numpy() if valid_mask is not None else None,
@@ -363,7 +374,7 @@ def eval(
 
         # Get metrics results
         iou_per_class, iou = iou.compute()
-        acc_per_class, acc = acc.compute()
+        acc_per_class, acc = cm.get_accuracy()
         iou_per_class = [float(n) for n in iou_per_class]
         acc_per_class = [float(n) for n in acc_per_class]
 
@@ -526,7 +537,7 @@ def eval(
 
         # Init metrics
         iou = um.IoU(self.n_classes)
-        acc = um.Accuracy(self.n_classes)
+        cm = um.ConfusionMatrix(self.n_classes)
 
         # Evaluation loop
         results = {}
@@ -590,13 +601,13 @@ def eval(
                 if lut_ontology is not None:
                     label = lut_ontology[label]
 
-                # Prepare data and update accuracy
+                # Prepare data and update confusion matrix
                 label = label.cpu().unsqueeze(0)
                 pred = self.transform_output(pred).cpu().unsqueeze(0).to(torch.int64)
                 if valid_mask is not None:
                     valid_mask = valid_mask.cpu().unsqueeze(0)
 
-                acc.update(
+                cm.update(
                     pred.numpy(),
                     label.numpy(),
                     valid_mask.numpy() if valid_mask is not None else None,
@@ -618,7 +629,7 @@ def eval(
 
         # Get metrics results
         iou_per_class, iou = iou.compute()
-        acc_per_class, acc = acc.compute()
+        acc_per_class, acc = cm.get_accuracy()
         iou_per_class = [float(n) for n in iou_per_class]
         acc_per_class = [float(n) for n in acc_per_class]