43 panoptic quality metrics

alodataset/coco_panoptic_dataset.py

@@ -9,6 +9,7 @@
 import numpy as np
 import torch
 from PIL import Image
+from typing import Union
 
 from alodataset.utils.panoptic_utils import rgb2id
 from alodataset.utils.panoptic_utils import masks_to_boxes
@@ -70,7 +71,7 @@ def __init__(
         split=Split.TRAIN,
         return_masks: bool = True,
         classes: list = None,
-        fix_classes_len: int = 250,  # Match with pre-trained weights
+        fix_classes_len: int = None,  # Match with pre-trained weights
         **kwargs,
     ):
         super(CocoPanopticDataset, self).__init__(name=name, split=split, **kwargs)
@@ -82,7 +83,7 @@ def __init__(
         self.ann_folder = os.path.join(self.dataset_dir, self.get_split_ann_folder())
         self.ann_file = os.path.join(self.dataset_dir, self.get_split_ann_file())
         self.return_masks = return_masks
-        self.label_names = None
+        self.label_names, self.label_types, self.label_types_names = None, None, None
         self.items = self._get_sequences()
 
         # Fix classes if it is desired
@@ -110,6 +111,15 @@ def __init__(
                     items.append(self.items[i])
             self.items = items
 
+            # Fix label_types: If `classes` is desired, remove types that not include this classes and fix indices
+            if self.label_types is not None:
+                for ltype, vtype in self.label_types.items():
+                    vtype = [x for b, x in enumerate(vtype) if self._ids_renamed[b] != -1]
+                    ltn = list(sorted(set([self.label_types_names[ltype][vt] for vt in vtype])))
+                    index = {b: ltn.index(p) for b, p in enumerate(self.label_types_names[ltype]) if p in ltn}
+                    self.label_types[ltype] = [index[idx] for idx in vtype]
+                    self.label_types_names[ltype] = ltn
+
         # Fix number of label names if desired
         if fix_classes_len is not None:
             if fix_classes_len > len(self.label_names):
@@ -146,9 +156,26 @@ def _get_sequences(self):
         if "categories" in coco:
             nb_category = max(cat["id"] for cat in coco["categories"])
             self.label_names = ["N/A"] * (nb_category + 1)
+
+            # Get types names
+            self.label_types_names = {
+                k: list(sorted(set([cat[k] for cat in coco["categories"]]))) + ["N/A"]
+                for k in coco["categories"][0].keys()
+                if k not in ["id", "name"]
+            }
+
+            # Make index between type category id and label id
+            self.label_types = {
+                k: [len(self.label_types_names[k]) - 1] * (nb_category + 1) for k in self.label_types_names
+            }
+            if "isthing" in self.label_types_names:
+                self.label_types_names["isthing"] = ["stuff", "thing", "N/A"]
             for cat in coco["categories"]:
                 self.label_names[cat["id"]] = cat["name"]
-        print("Done")
+                for k in self.label_types:
+                    self.label_types[k][cat["id"]] = (
+                        cat[k] if k == "isthing" else self.label_types_names[k].index(cat[k])
+                    )
         return items
 
     def get_split_ann_folder(self):
@@ -173,6 +200,18 @@ def get_split_ann_file(self):
         assert self.split in self.SPLIT_ANN_FILES
         return self.SPLIT_ANN_FILES[self.split]
 
+    def _append_type_labels(self, element: Union[BoundingBoxes2D, Mask], labels):
+        if self.label_types is not None:
+            for ktype in self.label_types:
+                label_types = torch.as_tensor(self.label_types[ktype])[labels]
+                label_types = Labels(
+                    label_types.to(torch.float32),
+                    labels_names=self.label_types_names[ktype],
+                    names=("N"),
+                    encoding="id",
+                )
+                element.append_labels(label_types, name=ktype)
+
     def getitem(self, idx):
         """Get the :mod:`Frame <aloscene.frame>` corresponds to *idx* index
 
@@ -213,28 +252,32 @@ def getitem(self, idx):
 
         # Make aloscene.frame
         frame = Frame(img_path)
-
         labels_2d = Labels(labels.to(torch.float32), labels_names=self.label_names, names=("N"), encoding="id")
         boxes_2d = BoundingBoxes2D(
-            masks_to_boxes(masks),
-            boxes_format="xyxy",
-            absolute=True,
-            frame_size=frame.HW,
-            names=("N", None),
-            labels=labels_2d,
+            masks_to_boxes(masks), boxes_format="xyxy", absolute=True, frame_size=frame.HW, names=("N", None),
         )
+        boxes_2d.append_labels(labels_2d, name="category")
+        self._append_type_labels(boxes_2d, labels)
         frame.append_boxes2d(boxes_2d)
 
         if self.return_masks:
-            masks_2d = Mask(masks, names=("N", "H", "W"), labels=labels_2d)
+            masks_2d = Mask(masks, names=("N", "H", "W"))
+            masks_2d.append_labels(labels_2d, name="category")
+            self._append_type_labels(masks_2d, labels)
             frame.append_segmentation(masks_2d)
         return frame
 
 
 if __name__ == "__main__":
-    coco_seg = CocoPanopticDataset(sample=True)
+    coco_seg = CocoPanopticDataset(sample=False)
     for f, frames in enumerate(coco_seg.train_loader(batch_size=2)):
         frames = Frame.batch_list(frames)
-        frames.get_view().render()
+        labels_set = "category" if isinstance(frames.boxes2d[0].labels, dict) else None
+        views = [fr.boxes2d.get_view(fr, labels_set=labels_set) for fr in frames]
+        if hasattr(frames, "segmentation"):
+            views += [fr.segmentation.get_view(fr, labels_set=labels_set) for fr in frames]
+        frames.get_view(views).render()
+        # frames.get_view(labels_set=labels_set).render()
+
         if f > 1:
             break

alodataset/utils/panoptic_utils.py

@@ -1,11 +1,17 @@
+from typing import Union
 import numpy as np
 import torch
 
+# from alonet.metrics.compute_pq import VOID
+
+VOID_CLASS_ID = -1
 GLOBAL_COLOR_SET = np.random.uniform(0, 1, (300, 3))
+GLOBAL_COLOR_SET[VOID_CLASS_ID] = [0, 0, 0]
+OFFSET = 256 * 256 * 256
 
 
 # Function get from PanopticAPI: https://github.com/cocodataset/panopticapi/blob/master/panopticapi/utils.py
-def rgb2id(color):
+def rgb2id(color: Union[list, np.ndarray]):
     if isinstance(color, np.ndarray) and len(color.shape) == 3:
         if color.dtype == np.uint8:
             color = color.astype(np.int32)
@@ -14,20 +20,20 @@ def rgb2id(color):
 
 
 # Function get from PanopticAPI: https://github.com/cocodataset/panopticapi/blob/master/panopticapi/utils.py
-def id2rgb(id_map, random_color=True):
+def id2rgb(id_map: np.ndarray, random_color: bool = True):
     if random_color:
-        return (256 * GLOBAL_COLOR_SET[id_map]).astype(np.uint8)
+        return GLOBAL_COLOR_SET[id_map]
     if isinstance(id_map, np.ndarray):
         id_map_copy = id_map.copy()
         rgb_shape = tuple(list(id_map.shape) + [3])
         rgb_map = np.zeros(rgb_shape, dtype=np.uint8)
         for i in range(3):
             rgb_map[..., i] = id_map_copy % 256
             id_map_copy //= 256
-        return rgb_map
+        return rgb_map / 255.0
     color = []
     for _ in range(3):
-        color.append(id_map % 256)
+        color.append((id_map % 256) / 255.0)
         id_map //= 256
     return color
 

alonet/callbacks/__init__.py

@@ -1,3 +1,5 @@
 from .object_detector_callback import ObjectDetectorCallback
 from .metrics_callback import MetricsCallback
+from .base_metrics_callback import InstancesBaseMetricsCallback
 from .map_metrics_callback import ApMetricsCallback
+from .pq_metrics_callback import PQMetricsCallback

alonet/callbacks/base_metrics_callback.py

@@ -0,0 +1,154 @@
+import pytorch_lightning as pl
+import aloscene
+from alonet import metrics
+from pytorch_lightning.utilities import rank_zero_only
+
+# import wandb
+
+
+class InstancesBaseMetricsCallback(pl.Callback):
+    def __init__(self, base_metric: metrics, *args, **kwargs):
+        self.metrics = []
+        self.base_metric = base_metric
+        super().__init__(*args, **kwargs)
+
+    def inference(self, pl_module: pl.LightningModule, m_outputs: dict, **kwargs):
+        """This method will call the `infernece` method of the module's model and will expect to receive the
+        predicted boxes2D and/or Masks.
+
+        Parameters
+        ----------
+        pl_module : pl.LightningModule
+            Pytorch lighting module with inference function
+        m_outputs : dict
+            Forward outputs
+
+        Returns
+        -------
+        :mod:`~aloscene.bounding_boxes_2d`, :mod:`~aloscene.Mask`
+            Boxes and masks predicted from inference function
+
+        Notes
+        -----
+            If `m_outputs` does not contain "pred_masks" attribute, a [None]*B attribute will be returned by default
+        """
+        b_pred_masks = None
+        if "pred_masks" in m_outputs:
+            b_pred_boxes, b_pred_masks = pl_module.inference(m_outputs, **kwargs)
+        else:
+            b_pred_boxes = pl_module.inference(m_outputs, **kwargs)
+        if not isinstance(m_outputs, list):
+            b_pred_boxes = [b_pred_boxes]
+            b_pred_masks = [b_pred_masks]
+        elif b_pred_masks is None:
+            b_pred_masks = [None] * len(b_pred_boxes)
+        return b_pred_boxes, b_pred_masks
+
+    @rank_zero_only
+    def on_validation_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: dict,
+        batch: list,
+        batch_idx: int,
+        dataloader_idx: int,
+    ):
+        """Method call after each validation batch. This class is a pytorch lightning callback, therefore
+        this method will by automaticly call by pl.
+
+        This method will call the `infernece` method of the module's model and will expect to receive the
+        predicted boxes2D and/or Masks. Theses elements will be aggregate to compute the different metrics in the
+        `on_validation_end` method.
+        The infernece method will be call using the `m_outputs` key from the outputs dict. If `m_outputs` is a list,
+        then the list will be consider as an temporal list. Therefore, this callback will aggregate the prediction
+        for each element of the sequence and will log the final results with the timestep prefix val/t/ instead of
+        simply /val/
+
+        Parameters
+        ----------
+        trainer: pl.Trainer
+            Pytorch lightning trainer
+        pl_module: pl.LightningModule
+            Pytorch lightning module. The "m_outputs" key is expected for this this callback to work properly.
+        outputs:
+            Training/Validation step outputs of the pl.LightningModule class.
+        batch: list
+            Batch comming from the dataloader. Usually, a list of frame.
+        batch_idx: int
+            Id the batch
+        dataloader_idx: int
+            Dataloader batch ID.
+        """
+        if isinstance(batch, list):  # Resize frames for mask procedure
+            batch = batch[0].batch_list(batch)
+
+        b_pred_boxes, b_pred_masks = self.inference(pl_module, outputs["m_outputs"])
+        is_temporal = isinstance(outputs["m_outputs"], list)
+        for b, (t_pred_boxes, t_pred_masks) in enumerate(zip(b_pred_boxes, b_pred_masks)):
+
+            # Retrieve the matching GT boxes at the same time step
+            t_gt_boxes = batch[b].boxes2d
+            t_gt_masks = batch[b].segmentation
+
+            if not is_temporal:
+                t_gt_boxes = [t_gt_boxes]
+                t_gt_masks = [t_gt_masks]
+
+            if t_pred_masks is None:
+                t_pred_masks = [None] * len(t_gt_masks)
+
+            # Add the samples to metrics for each batch of the current sequence
+            for t, (gt_boxes, pred_boxes, gt_masks, pred_masks) in enumerate(
+                zip(t_gt_boxes, t_pred_boxes, t_gt_masks, t_pred_masks)
+            ):
+                if t + 1 > len(self.metrics):
+                    self.metrics.append(self.base_metric())
+                self.add_sample(self.metrics[t], pred_boxes, gt_boxes, pred_masks, gt_masks)
+
+    @rank_zero_only
+    def add_sample(
+        self,
+        base_metric: metrics,
+        pred_boxes: aloscene.BoundingBoxes2D,
+        gt_boxes: aloscene.BoundingBoxes2D,
+        pred_masks: aloscene.Mask = None,
+        gt_masks: aloscene.Mask = None,
+    ):
+        """Add a smaple to some `alonet.metrics` class. One might want to inhert this method
+        to edit the `pred_boxes` and `gt_boxes` boxes before to add them to the ApMetrics class.
+
+        Parameters
+        ----------
+        ap_metrics: Union[:mod:`~alonet.metrics.ApMetrics`, :mod:`~alonet.metrics.PQMetrics`
+            ApMetrics intance.
+        pred_boxes: :mod:`~aloscene.BoundingBoxes2D`
+            Predicted boxes2D.
+        gt_boxes: :mod:`~aloscene.BoundingBoxes2D`
+            GT boxes2d.
+        pred_masks: :mod:`~aloscene.Mask`
+            Predicted Masks for segmentation task
+        gt_masks: :mod:`~aloscene.Mask`
+            GT masks in segmentation task.
+        """
+        base_metric.add_sample(p_bbox=pred_boxes, t_bbox=gt_boxes, p_mask=pred_masks, t_mask=gt_masks)
+
+    @rank_zero_only
+    def on_validation_end(self, trainer, pl_module):
+        """Method call at the end of each validation epoch. The method will use all the aggregate
+        data over the epoch to log the final metrics on wandb.
+        This class is a pytorch lightning callback, therefore this method will by automaticly call by pl.
+
+        This method is currently a WIP since some metrics are not logged due to some wandb error when loading
+        Table.
+
+        Parameters
+        ----------
+        trainer: pl.Trainer
+            Pytorch lightning trainer
+        pl_module: pl.LightningModule
+            Pytorch lightning module
+        """
+        if trainer.logger is None:
+            return
+        raise Exception("To inhert in a child class")