Documentation

@dataclass
class Bitmask:
    """
    Represents a binary mask with an associated semantic label.

    Parameters
    ----------
    mask : NDArray[np.bool_]
        A NumPy array of boolean values representing the mask.
    label : str
        The semantic label associated with the mask.

    Examples
    --------
    >>> import numpy as np
    >>> mask = np.array([[True, False], [False, True]], dtype=np.bool_)
    >>> bitmask = Bitmask(mask=mask, label='ocean')
    """

    mask: NDArray[np.bool_]
    label: str

    def __post_init__(self):
        if self.mask.dtype != np.bool_:
            raise ValueError(
                f"Bitmask recieved mask with dtype '{self.mask.dtype}'."
            )

@dataclass
class Segmentation:
    """
    Segmentation data structure holding ground truth and prediction bitmasks for semantic segmentation tasks.

    Parameters
    ----------
    uid : str
        Unique identifier for the image or sample.
    groundtruths : List[Bitmask]
        List of ground truth bitmasks.
    predictions : List[Bitmask]
        List of predicted bitmasks.
    shape : tuple of int, optional
        The shape of the segmentation masks. This is set automatically after initialization.
    size : int, optional
        The total number of pixels in the masks. This is set automatically after initialization.

    Examples
    --------
    >>> import numpy as np
    >>> mask1 = np.array([[True, False], [False, True]], dtype=np.bool_)
    >>> groundtruth = Bitmask(mask=mask1, label='object')
    >>> mask2 = np.array([[False, True], [True, False]], dtype=np.bool_)
    >>> prediction = Bitmask(mask=mask2, label='object')
    >>> segmentation = Segmentation(
    ...     uid='123',
    ...     groundtruths=[groundtruth],
    ...     predictions=[prediction]
    ... )
    """

    uid: str
    groundtruths: list[Bitmask]
    predictions: list[Bitmask]
    shape: tuple[int, ...] = field(default_factory=lambda: (0, 0))
    size: int = field(default=0)

    def __post_init__(self):

        groundtruth_shape = {
            groundtruth.mask.shape for groundtruth in self.groundtruths
        }
        prediction_shape = {
            prediction.mask.shape for prediction in self.predictions
        }
        if len(groundtruth_shape) == 0:
            raise ValueError("The segmenation is missing ground truths.")
        elif len(prediction_shape) == 0:
            raise ValueError("The segmenation is missing predictions.")
        elif (
            len(groundtruth_shape) != 1
            or len(prediction_shape) != 1
            or groundtruth_shape != prediction_shape
        ):
            raise ValueError(
                "A shape mismatch exists within the segmentation."
            )

        self.shape = groundtruth_shape.pop()
        self.size = int(np.prod(np.array(self.shape)))

class DataLoader:
    """
    Segmentation DataLoader.
    """

    def __init__(self):
        self._evaluator = Evaluator()
        self.groundtruth_count = defaultdict(defaultdict_int)
        self.prediction_count = defaultdict(defaultdict_int)
        self.matrices = list()
        self.pixel_count = list()

    def _add_datum(self, uid: str) -> int:
        """
        Helper function for adding a datum to the cache.

        Parameters
        ----------
        uid : str
            The datum uid.

        Returns
        -------
        int
            The datum index.
        """
        if uid in self._evaluator.uid_to_index:
            raise ValueError(f"Datum with uid `{uid}` has already been added.")
        index = len(self._evaluator.uid_to_index)
        self._evaluator.uid_to_index[uid] = index
        self._evaluator.index_to_uid[index] = uid
        return index

    def _add_label(self, label: str) -> int:
        """
        Helper function for adding a label to the cache.

        Parameters
        ----------
        label : str
            A string label.

        Returns
        -------
        int
            The label's index.
        """
        if label not in self._evaluator.label_to_index:
            label_id = len(self._evaluator.index_to_label)
            self._evaluator.label_to_index[label] = label_id
            self._evaluator.index_to_label[label_id] = label
        return self._evaluator.label_to_index[label]

    def add_data(
        self,
        segmentations: list[Segmentation],
        show_progress: bool = False,
    ):
        """
        Adds segmentations to the cache.

        Parameters
        ----------
        segmentations : list[Segmentation]
            A list of Segmentation objects.
        show_progress : bool, default=False
            Toggle for tqdm progress bar.
        """

        disable_tqdm = not show_progress
        for segmentation in tqdm(segmentations, disable=disable_tqdm):

            # update metadata
            self._evaluator.n_datums += 1
            self._evaluator.n_groundtruths += len(segmentation.groundtruths)
            self._evaluator.n_predictions += len(segmentation.predictions)
            self._evaluator.n_pixels += segmentation.size
            self._evaluator.n_groundtruth_pixels += segmentation.size * len(
                segmentation.groundtruths
            )
            self._evaluator.n_prediction_pixels += segmentation.size * len(
                segmentation.predictions
            )

            # update datum cache
            uid_index = self._add_datum(segmentation.uid)

            groundtruth_labels = np.full(
                len(segmentation.groundtruths), fill_value=-1
            )
            for idx, groundtruth in enumerate(segmentation.groundtruths):
                label_idx = self._add_label(groundtruth.label)
                groundtruth_labels[idx] = label_idx
                self.groundtruth_count[label_idx][
                    uid_index
                ] += groundtruth.mask.sum()

            prediction_labels = np.full(
                len(segmentation.predictions), fill_value=-1
            )
            for idx, prediction in enumerate(segmentation.predictions):
                label_idx = self._add_label(prediction.label)
                prediction_labels[idx] = label_idx
                self.prediction_count[label_idx][
                    uid_index
                ] += prediction.mask.sum()

            combined_groundtruths = np.stack(
                [
                    groundtruth.mask.flatten()
                    for groundtruth in segmentation.groundtruths
                ],
                axis=0,
            )
            combined_predictions = np.stack(
                [
                    prediction.mask.flatten()
                    for prediction in segmentation.predictions
                ],
                axis=0,
            )

            self.matrices.append(
                compute_intermediate_confusion_matrices(
                    groundtruths=combined_groundtruths,
                    predictions=combined_predictions,
                    groundtruth_labels=groundtruth_labels,
                    prediction_labels=prediction_labels,
                    n_labels=len(self._evaluator.index_to_label),
                )
            )
            self.pixel_count.append(segmentation.size)

    def finalize(self) -> Evaluator:
        """
        Performs data finalization and some preprocessing steps.

        Returns
        -------
        Evaluator
            A ready-to-use evaluator object.
        """

        if len(self.matrices) == 0:
            raise ValueError("No data available to create evaluator.")

        n_datums = self._evaluator.n_datums
        n_labels = len(self._evaluator.index_to_label)

        self._evaluator.n_labels = n_labels

        self._evaluator._label_metadata_per_datum = np.zeros(
            (2, n_datums, n_labels), dtype=np.int32
        )
        for datum_idx in range(n_datums):
            for label_idx in range(n_labels):
                gt_count = (
                    self.groundtruth_count[label_idx].get(datum_idx, 0)
                    if label_idx in self.groundtruth_count
                    else 0
                )
                pd_count = (
                    self.prediction_count[label_idx].get(datum_idx, 0)
                    if label_idx in self.prediction_count
                    else 0
                )
                self._evaluator._label_metadata_per_datum[
                    :, datum_idx, label_idx
                ] = np.array([gt_count, pd_count])

        self._evaluator._label_metadata = np.array(
            [
                [
                    np.sum(
                        self._evaluator._label_metadata_per_datum[
                            0, :, label_idx
                        ]
                    ),
                    np.sum(
                        self._evaluator._label_metadata_per_datum[
                            1, :, label_idx
                        ]
                    ),
                ]
                for label_idx in range(n_labels)
            ],
            dtype=np.int32,
        )

        self._evaluator._n_pixels_per_datum = np.array(
            self.pixel_count, dtype=np.int32
        )

        self._evaluator._confusion_matrices = np.zeros(
            (n_datums, n_labels + 1, n_labels + 1), dtype=np.int32
        )
        for idx, matrix in enumerate(self.matrices):
            h, w = matrix.shape
            self._evaluator._confusion_matrices[idx, :h, :w] = matrix

        return self._evaluator

class Evaluator:
    """
    Segmentation Evaluator
    """

    def __init__(self):

        # metadata
        self.n_datums = 0
        self.n_groundtruths = 0
        self.n_predictions = 0
        self.n_pixels = 0
        self.n_groundtruth_pixels = 0
        self.n_prediction_pixels = 0
        self.n_labels = 0

        # datum reference
        self.uid_to_index: dict[str, int] = dict()
        self.index_to_uid: dict[int, str] = dict()

        # label reference
        self.label_to_index: dict[str, int] = dict()
        self.index_to_label: dict[int, str] = dict()

        # computation caches
        self._confusion_matrices = np.array([])
        self._label_metadata = np.array([], dtype=np.int32)
        self._label_metadata_per_datum = np.array([], dtype=np.int32)
        self._n_pixels_per_datum = np.array([], dtype=np.int32)

    @property
    def ignored_prediction_labels(self) -> list[str]:
        """
        Prediction labels that are not present in the ground truth set.
        """
        glabels = set(np.where(self._label_metadata[:, 0] > 0)[0])
        plabels = set(np.where(self._label_metadata[:, 1] > 0)[0])
        return [
            self.index_to_label[label_id] for label_id in (plabels - glabels)
        ]

    @property
    def missing_prediction_labels(self) -> list[str]:
        """
        Ground truth labels that are not present in the prediction set.
        """
        glabels = set(np.where(self._label_metadata[:, 0] > 0)[0])
        plabels = set(np.where(self._label_metadata[:, 1] > 0)[0])
        return [
            self.index_to_label[label_id] for label_id in (glabels - plabels)
        ]

    @property
    def metadata(self) -> dict:
        """
        Evaluation metadata.
        """
        return {
            "number_of_datums": self.n_datums,
            "number_of_groundtruths": self.n_groundtruths,
            "number_of_predictions": self.n_predictions,
            "number_of_groundtruth_pixels": self.n_groundtruth_pixels,
            "number_of_prediction_pixels": self.n_prediction_pixels,
            "number_of_labels": self.n_labels,
            "ignored_prediction_labels": self.ignored_prediction_labels,
            "missing_prediction_labels": self.missing_prediction_labels,
        }

    def create_filter(
        self,
        datum_uids: list[str] | NDArray[np.int32] | None = None,
        labels: list[str] | NDArray[np.int32] | None = None,
    ) -> Filter:
        """
        Creates a boolean mask that can be passed to an evaluation.

        Parameters
        ----------
        datum_uids : list[str] | NDArray[np.int32], optional
            An optional list of string uids or a numpy array of uid indices.
        labels : list[tuple[str, str]] | NDArray[np.int32], optional
            An optional list of labels or a numpy array of label indices.

        Returns
        -------
        Filter
            A filter object that can be passed to the `evaluate` method.
        """
        n_datums = self._label_metadata_per_datum.shape[1]
        n_labels = self._label_metadata_per_datum.shape[2]

        mask_datums = np.ones(n_datums, dtype=np.bool_)
        mask_labels = np.ones(n_labels, dtype=np.bool_)

        if datum_uids is not None:
            if isinstance(datum_uids, list):
                datum_uids = np.array(
                    [self.uid_to_index[uid] for uid in datum_uids],
                    dtype=np.int32,
                )
            if datum_uids.size == 0:
                mask_datums[mask_datums] = False
            else:
                mask = (
                    np.arange(n_datums).reshape(-1, 1)
                    == datum_uids.reshape(1, -1)
                ).any(axis=1)
                mask_datums[~mask] = False

        if labels is not None:
            if isinstance(labels, list):
                labels = np.array(
                    [self.label_to_index[label] for label in labels],
                    dtype=np.int32,
                )
            if labels.size == 0:
                mask_labels[mask_labels] = False
            else:
                mask = (
                    np.arange(n_labels).reshape(-1, 1) == labels.reshape(1, -1)
                ).any(axis=1)
                mask_labels[~mask] = False

        mask = mask_datums[:, np.newaxis] & mask_labels[np.newaxis, :]
        label_metadata_per_datum = self._label_metadata_per_datum.copy()
        label_metadata_per_datum[:, ~mask] = 0

        label_metadata = np.zeros_like(self._label_metadata, dtype=np.int32)
        label_metadata = np.transpose(
            np.sum(
                label_metadata_per_datum,
                axis=1,
            )
        )
        n_datums = int(np.sum(label_metadata[:, 0]))

        return Filter(
            indices=np.where(mask_datums)[0],
            label_metadata=label_metadata,
            n_pixels=self._n_pixels_per_datum[mask_datums].sum(),
        )

    def compute_precision_recall_iou(
        self,
        filter_: Filter | None = None,
    ) -> dict[MetricType, list]:
        """
        Performs an evaluation and returns metrics.

        Parameters
        ----------
        filter_ : Filter, optional
            An optional filter object.

        Returns
        -------
        dict[MetricType, list]
            A dictionary mapping MetricType enumerations to lists of computed metrics.
        """

        # apply filters
        data = self._confusion_matrices
        label_metadata = self._label_metadata
        n_pixels = self.n_pixels
        if filter_ is not None:
            data = data[filter_.indices]
            label_metadata = filter_.label_metadata
            n_pixels = filter_.n_pixels

        results = compute_metrics(
            data=data,
            label_metadata=label_metadata,
            n_pixels=n_pixels,
        )

        return unpack_precision_recall_iou_into_metric_lists(
            results=results,
            label_metadata=label_metadata,
            index_to_label=self.index_to_label,
        )

    def evaluate(
        self,
        filter_: Filter | None = None,
    ) -> dict[MetricType, list[Metric]]:
        """
        Computes all available metrics.

        Parameters
        ----------
        filter_ : Filter, optional
            An optional filter object.

        Returns
        -------
        dict[MetricType, list[Metric]]
            Lists of metrics organized by metric type.
        """
        return self.compute_precision_recall_iou(filter_=filter_)