Merge pull request #267 from JdeRobot/dph/issue-240

Update metrics & improvements in datasets

Author: dpascualhe

detectionmetrics/cli/evaluate.py

@@ -33,8 +33,9 @@ def get_dataset(
     split,
 ):
     # Check if required data is available
-    if dataset_format == "gaia" and dataset_fname is None:
-        raise ValueError("--dataset is required for 'gaia' format")
+    if dataset_format == "gaia":
+        if dataset_fname is None:
+            raise ValueError("--dataset is required for 'gaia' format")
 
     elif dataset_format == "rellis3d":
         if dataset_dir is None:

detectionmetrics/datasets/dataset.py

@@ -28,6 +28,7 @@ def __init__(self, dataset: pd.DataFrame, dataset_dir: str, ontology: dict):
         self.dataset = dataset
         self.dataset_dir = os.path.abspath(dataset_dir)
         self.ontology = ontology
+        self.has_label_count = all("label_count" in v for v in self.ontology.values())
 
     def __len__(self):
         return len(self.dataset)
@@ -54,6 +55,16 @@ def append(self, new_dataset: Self):
             [self.dataset, new_dataset.dataset], verify_integrity=True
         )
 
+    def get_label_count(self, splits: List[str] = ["train", "val"]) -> np.ndarray:
+        """Get label count for each class in the dataset
+
+        :param splits: Dataset splits to consider, defaults to ["train", "val"]
+        :type splits: List[str], optional
+        :return: Label count for the dataset
+        :rtype: np.ndarray
+        """
+        raise NotImplementedError
+
 
 class ImageSegmentationDataset(SegmentationDataset):
     """Parent image segmentation dataset class
@@ -94,7 +105,9 @@ def export(
         outdir: str,
         new_ontology: Optional[dict] = None,
         ontology_translation: Optional[dict] = None,
-        ignored_classes: Optional[List[str]] = [],
+        ignored_classes: Optional[List[str]] = None,
+        resize: Optional[Tuple[int, int]] = None,
+        include_label_count: bool = True,
     ):
         """Export dataset dataframe and image files in SemanticKITTI format. Optionally, modify ontology before exporting.
 
@@ -106,6 +119,10 @@ def export(
         :type ontology_translation: Optional[dict], optional
         :param ignored_classes: Classes to ignore from the old ontology, defaults to []
         :type ignored_classes: Optional[List[str]], optional
+        :param resize: Resize images and labels to the given dimensions, defaults to None
+        :type resize: Optional[Tuple[int, int]], optional
+        :param include_label_count: Whether to include class weights in the dataset, defaults to True
+        :type include_label_count: bool, optional
         """
         os.makedirs(outdir, exist_ok=True)
 
@@ -117,7 +134,7 @@ def export(
         if ontology_translation is not None and new_ontology is None:
             raise ValueError("New ontology must be provided")
 
-        # Create ontology conversion lookup table
+        # Create ontology conversion lookup table if needed and get number of classes
         ontology_conversion_lut = None
         if new_ontology is not None:
             ontology_conversion_lut = uc.get_ontology_conversion_lut(
@@ -126,6 +143,16 @@ def export(
                 ontology_translation=ontology_translation,
                 ignored_classes=ignored_classes,
             )
+            n_classes = max(c["idx"] for c in new_ontology.values()) + 1
+        else:
+            n_classes = max(c["idx"] for c in self.ontology.values()) + 1
+
+        # Check if label count is missing and create empty array if needed
+        label_count_missing = include_label_count and (
+            not self.has_label_count or new_ontology is not None
+        )
+        if label_count_missing:
+            label_count = np.zeros(n_classes, dtype=np.uint64)
 
         # Export each sample
         for sample_name, row in pbar:
@@ -149,20 +176,29 @@ def export(
                     label_fname = os.path.join(self.dataset_dir, label_fname)
 
             # If image mode is not appropriate: read, convert, and rewrite image
-            if uio.get_image_mode(image_fname) != "RGB":
+            if uio.get_image_mode(image_fname) != "RGB" or resize is not None:
                 image = cv2.imread(image_fname, 1)  # convert to RGB
+
+                # Resize image if needed
+                if resize is not None:
+                    image = cv2.resize(image, resize, interpolation=cv2.INTER_CUBIC)
                 cv2.imwrite(os.path.join(outdir, rel_image_fname), image)
-            # if image mode is appropriate simply copy image to new location
+
+            # If image mode is appropriate simply copy image to new location
             else:
                 shutil.copy2(image_fname, os.path.join(outdir, rel_image_fname))
             self.dataset.at[sample_name, "image"] = rel_image_fname
 
-            # Same for labels (plus ontology conversion if needed)
+            # Same for labels (plus ontology conversion and label count if needed)
             if label_fname:
                 image_mode = uio.get_image_mode(label_fname)
-                if image_mode == "L" and ontology_conversion_lut is None:
-                    shutil.copy2(label_fname, os.path.join(outdir, rel_label_fname))
-                else:
+                if (
+                    image_mode != "L"
+                    or ontology_conversion_lut is not None
+                    or resize is not None
+                    or label_count_missing
+                ):
+                    # Read and convert label from RGB to L
                     if self.is_label_rgb:
                         label_rgb = cv2.imread(label_fname)[:, :, ::-1]
                         label = np.zeros(label_rgb.shape[:2], dtype=np.uint8)
@@ -172,16 +208,37 @@ def export(
                             label[(label_rgb == rgb).all(axis=2)] = idx
                     else:
                         label = cv2.imread(label_fname, 0)  # convert to L
-                        if ontology_conversion_lut is not None:
-                            label = ontology_conversion_lut[label]
+
+                    # Convert label to new ontology if needed
+                    if ontology_conversion_lut is not None:
+                        label = ontology_conversion_lut[label]
+
+                    # Resize label if needed
+                    if resize is not None:
+                        label = cv2.resize(
+                            label, resize, interpolation=cv2.INTER_NEAREST
+                        )
+
+                    # Update label count if needed
+                    if label_count_missing:
+                        indices, counts = np.unique(label, return_counts=True)
+                        label_count[indices] += counts.astype(np.uint64)
+
                     cv2.imwrite(os.path.join(outdir, rel_label_fname), label)
+                else:
+                    shutil.copy2(label_fname, os.path.join(outdir, rel_label_fname))
+
                 self.dataset.at[sample_name, "label"] = rel_label_fname
 
         # Update dataset directory and ontology if needed
         self.dataset_dir = outdir
         self.ontology = new_ontology if new_ontology is not None else self.ontology
 
         # Write ontology and store relative path in dataset attributes
+        if label_count_missing:
+            for class_data in self.ontology.values():
+                class_data["label_count"] = int(label_count[class_data["idx"]])
+
         ontology_fname = "ontology.json"
         self.dataset.attrs = {"ontology_fname": ontology_fname}
         uio.write_json(os.path.join(outdir, ontology_fname), self.ontology)

detectionmetrics/models/tensorflow.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 import os
 import time
 from typing import List, Optional, Tuple, Union
@@ -94,6 +95,8 @@ class ImageSegmentationTensorflowDataset:
     :type split: str, optional
     :param lut_ontology: LUT to transform label classes, defaults to None
     :type lut_ontology: dict, optional
+    :param normalization: Parameters for normalizing input images, defaults to None
+    :type normalization: dict, optional
     """
 
     def __init__(
@@ -103,8 +106,14 @@ def __init__(
         batch_size: int = 1,
         split: str = "all",
         lut_ontology: Optional[dict] = None,
+        normalization: Optional[dict] = None,
     ):
         self.image_size = image_size
+        self.normalization = None
+        if normalization is not None:
+            mean = tf.constant(normalization["mean"], dtype=tf.float32)
+            std = tf.constant(normalization["std"], dtype=tf.float32)
+            self.normalization = {"mean": mean, "std": std}
 
         # Filter split and make filenames global
         if split != "all":
@@ -155,9 +164,17 @@ def read_image(self, fname: str, label=False) -> tf.Tensor:
         # Resize (use NN to avoid interpolation when dealing with labels)
         method = "nearest" if label else "bilinear"
         image = tf_image.resize(images=image, size=self.image_size, method=method)
+
+        # If label, round values to avoid interpolation artifacts
         if label:
             image = tf.round(image)
 
+        # If normalization parameters are provided, normalize image
+        else:
+            if self.normalization is not None:
+                image = tf.cast(image, tf.float32) / 255.0
+                image = (image - self.normalization["mean"]) / self.normalization["std"]
+
         return image
 
     def load_data(
@@ -217,6 +234,11 @@ def t_in(image):
             tensor = tf.convert_to_tensor(image)
             tensor = tf_image.resize(images=tensor, size=self.model_cfg["image_size"])
             tensor = tf.expand_dims(tensor, axis=0)
+            if "normalization" in self.model_cfg:
+                mean = tf.constant(self.model_cfg["normalization"]["mean"])
+                std = tf.constant(self.model_cfg["normalization"]["std"])
+                tensor = tf.cast(tensor, tf.float32) / 255.0
+                tensor = (tensor - mean) / std
             return tensor
 
         self.t_in = t_in
@@ -275,18 +297,23 @@ def eval(
             batch_size=self.model_cfg.get("batch_size", 1),
             split=split,
             lut_ontology=lut_ontology,
+            normalization=self.model_cfg.get("normalization", None),
         )
 
+        # Retrieve ignored label indices
+        ignored_label_indices = []
+        for ignored_class in self.model_cfg.get("ignored_classes", []):
+            ignored_label_indices.append(dataset.ontology[ignored_class]["idx"])
+
         # Init metrics
         results = {}
-        iou = um.IoU(self.n_classes)
-        cm = um.ConfusionMatrix(self.n_classes)
+        metrics_factory = um.MetricsFactory(self.n_classes)
 
         # Evaluation loop
         pbar = tqdm(dataset.dataset)
         for image, label in pbar:
             if self.model_type == "native":
-                pred = self.model(image)
+                pred = self.model(image, training=False)
             elif self.model_type == "compiled":
                 pred = self.model.signatures["serving_default"](image)
             else:
@@ -295,37 +322,48 @@ def eval(
             if isinstance(pred, dict):
                 pred = list(pred.values())[0]
 
+            # Get valid points masks depending on ignored label indices
+            if ignored_label_indices:
+                valid_mask = tf.ones_like(label, dtype=tf.bool)
+                for idx in ignored_label_indices:
+                    valid_mask *= label != idx
+            else:
+                valid_mask = None
+
             label = tf.squeeze(label, axis=3)
             pred = tf.argmax(pred, axis=3)
-            cm.update(pred.numpy(), label.numpy())
-
-            pred = tf.one_hot(pred, self.n_classes)
-            pred = tf.transpose(pred, perm=[0, 3, 1, 2])
+            if valid_mask is not None:
+                valid_mask = tf.squeeze(valid_mask, axis=3)
+            metrics_factory.update(
+                pred.numpy(),
+                label.numpy(),
+                valid_mask.numpy() if valid_mask is not None else None,
+            )
 
-            label = tf.one_hot(label, self.n_classes)
-            label = tf.transpose(label, perm=[0, 3, 1, 2])
+        # Build results dataframe
+        results = defaultdict(dict)
 
-            iou.update(pred.numpy(), label.numpy())
+        # Add per class and global metrics
+        for metric in metrics_factory.get_metric_names():
+            per_class = metrics_factory.get_metric_per_name(metric, per_class=True)
 
-        # Get metrics results
-        iou_per_class, iou = iou.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]
+            for class_name, class_data in self.ontology.items():
+                results[class_name][metric] = float(per_class[class_data["idx"]])
 
-        # Build results dataframe
-        results = {}
-        for class_name, class_data in self.ontology.items():
-            results[class_name] = {
-                "iou": iou_per_class[class_data["idx"]],
-                "acc": acc_per_class[class_data["idx"]],
-            }
-        results["global"] = {"iou": iou, "acc": acc}
+            if metric not in ["tp", "fp", "fn", "tn"]:
+                for avg_method in ["macro", "micro"]:
+                    results[avg_method][metric] = metrics_factory.get_averaged_metric(
+                        metric, avg_method
+                    )
 
-        results = pd.DataFrame(results)
-        results.index.name = "metric"
+        # Add confusion matrix
+        for class_name_a, class_data_a in self.ontology.items():
+            for class_name_b, class_data_b in self.ontology.items():
+                results[class_name_a][class_name_b] = metrics_factory.confusion_matrix[
+                    class_data_a["idx"], class_data_b["idx"]
+                ]
 
-        return results
+        return pd.DataFrame(results)
 
     def get_computational_cost(self, runs: int = 30, warm_up_runs: int = 5) -> dict:
         """Get different metrics related to the computational cost of the model