scale-nucleus 0.12b1__py3-none-any.whl → 0.14.14b0__py3-none-any.whl

nucleus/metrics/segmentation_loader.py

@@ -0,0 +1,30 @@
+import abc
+from typing import Dict
+
+import numpy as np
+
+
+class SegmentationMaskLoader(abc.ABC):
+    @abc.abstractmethod
+    def fetch(self, url: str) -> np.ndarray:
+        pass
+
+
+class DummyLoader(SegmentationMaskLoader):
+    def fetch(self, url: str) -> np.ndarray:
+        raise NotImplementedError(
+            "This dummy loader has to be replaced with an actual implementation of an image loader"
+        )
+
+
+class InMemoryLoader(SegmentationMaskLoader):
+    """We use this loader in the tests, this allows us to serve images from memory instead of fetching
+    from a filesystem.
+    """
+
+    def __init__(self, url_to_array: Dict[str, np.ndarray]):
+        self.url_to_array = url_to_array
+
+    def fetch(self, url: str):
+        array = self.url_to_array[url]
+        return array

nucleus/metrics/segmentation_metrics.py

@@ -1,42 +1,26 @@
 import abc
-from typing import List, Optional, Union
+from typing import List, Optional, Set, Tuple, Union
 
-import fsspec
 import numpy as np
-from PIL import Image
-from s3fs import S3FileSystem
 
-from nucleus.annotation import AnnotationList, SegmentationAnnotation
+from nucleus.annotation import AnnotationList, Segment, SegmentationAnnotation
 from nucleus.metrics.base import MetricResult
 from nucleus.metrics.filtering import ListOfAndFilters, ListOfOrAndFilters
 from nucleus.prediction import PredictionList, SegmentationPrediction
 
 from .base import Metric, ScalarResult
-from .metric_utils import compute_average_precision
+from .segmentation_loader import DummyLoader, SegmentationMaskLoader
+from .segmentation_utils import (
+    FALSE_POSITIVES,
+    convert_to_instance_seg_confusion,
+    fast_confusion_matrix,
+    non_max_suppress_confusion,
+    setup_iou_thresholds,
+)
 
 # pylint: disable=useless-super-delegation
 
 
-def _fast_hist(label_true, label_pred, n_class):
-    """Calculates confusion matrix - fast!"""
-    mask = (label_true >= 0) & (label_true < n_class)
-    hist = np.bincount(
-        n_class * label_true[mask].astype(int) + label_pred[mask],
-        minlength=n_class ** 2,
-    ).reshape(n_class, n_class)
-    return hist
-
-
-class SegmentationMaskLoader:
-    def __init__(self, fs: fsspec):
-        self.fs = fs
-
-    def fetch(self, url: str):
-        with self.fs.open(url) as fh:
-            img = Image.open(fh)
-        return img
-
-
 class SegmentationMaskMetric(Metric):
     def __init__(
         self,
@@ -46,6 +30,7 @@ class SegmentationMaskMetric(Metric):
         prediction_filters: Optional[
             Union[ListOfOrAndFilters, ListOfAndFilters]
         ] = None,
+        iou_threshold: float = 0.5,
     ):
         """Initializes PolygonMetric abstract object.
 
@@ -71,13 +56,12 @@ class SegmentationMaskMetric(Metric):
         """
         # TODO -> add custom filtering to Segmentation(Annotation|Prediction).annotations.(metadata|label)
         super().__init__(annotation_filters, prediction_filters)
-        self.loader = SegmentationMaskLoader(S3FileSystem(anon=False))
-        # NOTE: We store histogram for re-use in subsequently calculated metrics
-        self.confusion: Optional[np.ndarray] = None
+        self.loader: SegmentationMaskLoader = DummyLoader()
+        self.iou_threshold = iou_threshold
 
     def call_metric(
         self, annotations: AnnotationList, predictions: PredictionList
-    ) -> MetricResult:
+    ) -> ScalarResult:
         assert (
             len(annotations.segmentation_annotations) <= 1
         ), f"Expected only one segmentation mask, got {annotations.segmentation_annotations}"
@@ -94,29 +78,71 @@ class SegmentationMaskMetric(Metric):
             if predictions.segmentation_predictions
             else None
         )
-        annotation_img = self.loader.fetch(annotation.mask_url)
-        pred_img = self.loader.fetch(prediction.mask_url)
-        return self._metric_impl(
-            np.asarray(annotation_img, dtype=np.int32),
-            np.asarray(pred_img, dtype=np.int32),
-            annotation,
-            prediction,
-        )
+        if (
+            annotation
+            and prediction
+            and annotation.annotations
+            and prediction.annotations
+        ):
+            annotation_img = self.get_mask_channel(annotation)
+            pred_img = self.get_mask_channel(prediction)
+            return self._metric_impl(
+                np.asarray(annotation_img, dtype=np.int32),
+                np.asarray(pred_img, dtype=np.int32),
+                annotation,
+                prediction,
+            )
+        else:
+            return ScalarResult(0, weight=0)
+
+    def get_mask_channel(self, ann_or_pred):
+        """Some annotations are stored as RGB instead of L (single-channel).
+        We expect the image to be faux-single-channel with all the channels repeating so we choose the first one.
+        """
+        img = self.loader.fetch(ann_or_pred.mask_url)
+        if len(img.shape) > 2:
+            # TODO: Do we have to do anything more advanced? Currently expect all channels to have same data
+            min_dim = np.argmin(img.shape)
+            if min_dim == 0:
+                img = img[0, :, :]
+            elif min_dim == 1:
+                img = img[:, 0, :]
+            else:
+                img = img[:, :, 0]
+        return img
 
     @abc.abstractmethod
     def _metric_impl(
         self,
         annotation_img: np.ndarray,
         prediction_img: np.ndarray,
-        annotation: Optional[SegmentationAnnotation],
-        prediction: Optional[SegmentationPrediction],
+        annotation: SegmentationAnnotation,
+        prediction: SegmentationPrediction,
     ):
         pass
 
     def _calculate_confusion_matrix(
-        self, annotation, annotation_img, prediction, prediction_img
-    ):
-        # NOTE: This creates a max(class_index) * max(class_index) MAT. If we have np.int16 this could become
+        self,
+        annotation,
+        annotation_img,
+        prediction,
+        prediction_img,
+        iou_threshold,
+    ) -> Tuple[np.ndarray, Set[int]]:
+        """This calculates a confusion matrix with ground_truth_index X predicted_index summary
+
+        Notes:
+            If filtering has been applied we filter out missing segments from the confusion matrix.
+
+        Returns:
+            Class-based confusion matrix and a set of indexes that are not considered a part of the taxonomy (and are
+            only considered for FPs not as a part of mean calculations)
+
+
+        TODO(gunnar): Allow pre-seeding confusion matrix (all of the metrics calculate the same confusion matrix ->
+            we can calculate it once and then use it for all other metrics in the chain)
+        """
+        # NOTE: This creates a max(class_index) * max(class_index) MAT. If we have np.int32 this could become
         #  huge. We could probably use a sparse matrix instead or change the logic to only create count(index) ** 2
         #  matrix (we only need to keep track of available indexes)
         num_classes = (
@@ -126,11 +152,80 @@ class SegmentationMaskMetric(Metric):
             )
             + 1  # to include 0
         )
-        confusion = (
-            _fast_hist(annotation_img, prediction_img, num_classes)
-            if self.confusion is None
-            else self.confusion
+        confusion = fast_confusion_matrix(
+            annotation_img, prediction_img, num_classes
+        )
+        confusion = self._filter_confusion_matrix(
+            confusion, annotation, prediction
+        )
+        confusion = non_max_suppress_confusion(confusion, iou_threshold)
+        false_positive = Segment(FALSE_POSITIVES, index=confusion.shape[0] - 1)
+        if annotation.annotations[-1].label != FALSE_POSITIVES:
+            annotation.annotations.append(false_positive)
+            if annotation.annotations is not prediction.annotations:
+                # Probably likely that this structure is re-used -> check if same list instance and only append once
+                # TODO(gunnar): Should this uniqueness be handled by the base class?
+                prediction.annotations.append(false_positive)
+
+        # TODO(gunnar): Detect non_taxonomy classes for segmentation as well as instance segmentation
+        non_taxonomy_classes = set()
+        if self._is_instance_segmentation(annotation, prediction):
+            (
+                confusion,
+                _,
+                non_taxonomy_classes,
+            ) = convert_to_instance_seg_confusion(
+                confusion, annotation, prediction
+            )
+        else:
+            ann_labels = list(
+                dict.fromkeys(s.label for s in annotation.annotations)
+            )
+            pred_labels = list(
+                dict.fromkeys(s.label for s in prediction.annotations)
+            )
+            missing_or_filtered_labels = set(ann_labels) - set(pred_labels)
+            non_taxonomy_classes = {
+                segment.index
+                for segment in annotation.annotations
+                if segment.label in missing_or_filtered_labels
+            }
+
+        return confusion, non_taxonomy_classes
+
+    def _is_instance_segmentation(self, annotation, prediction):
+        """Guesses that we're dealing with instance segmentation if we have multiple segments with same label.
+        Degenerate case is same as semseg so fine to misclassify in that case."""
+        # This is a trick to get ordered sets
+        ann_labels = list(
+            dict.fromkeys(s.label for s in annotation.annotations)
+        )
+        pred_labels = list(
+            dict.fromkeys(s.label for s in prediction.annotations)
         )
+        # NOTE: We assume instance segmentation if labels are duplicated in annotations or predictions
+        is_instance_segmentation = len(ann_labels) != len(
+            annotation.annotations
+        ) or len(pred_labels) != len(prediction.annotations)
+        return is_instance_segmentation
+
+    def _filter_confusion_matrix(self, confusion, annotation, prediction):
+        if self.annotation_filters or self.prediction_filters:
+            new_confusion = np.zeros_like(confusion)
+            # we mask the confusion matrix instead of the images
+            if self.annotation_filters:
+                annotation_indexes = {
+                    segment.index for segment in annotation.annotations
+                }
+                for row in annotation_indexes:
+                    new_confusion[row, :] = confusion[row, :]
+            if self.prediction_filters:
+                prediction_indexes = {
+                    segment.index for segment in prediction.annotations
+                }
+                for col in prediction_indexes:
+                    new_confusion[:, col] = confusion[:, col]
+            confusion = new_confusion
         return confusion
 
 
@@ -143,6 +238,7 @@ class SegmentationIOU(SegmentationMaskMetric):
         prediction_filters: Optional[
             Union[ListOfOrAndFilters, ListOfAndFilters]
         ] = None,
+        iou_threshold: float = 0.5,
     ):
         """Initializes PolygonIOU object.
 
@@ -169,30 +265,36 @@ class SegmentationIOU(SegmentationMaskMetric):
         super().__init__(
             annotation_filters,
             prediction_filters,
+            iou_threshold,
         )
 
     def _metric_impl(
         self,
         annotation_img: np.ndarray,
         prediction_img: np.ndarray,
-        annotation: Optional[SegmentationAnnotation],
-        prediction: Optional[SegmentationPrediction],
+        annotation: SegmentationAnnotation,
+        prediction: SegmentationPrediction,
     ) -> ScalarResult:
-        if annotation is None or prediction is None:
-            # TODO: Throw error when we wrap each item in try catch
-            return ScalarResult(0, weight=0)
-
-        self.confusion = self._calculate_confusion_matrix(
-            annotation, annotation_img, prediction, prediction_img
+        confusion, non_taxonomy_classes = self._calculate_confusion_matrix(
+            annotation,
+            annotation_img,
+            prediction,
+            prediction_img,
+            self.iou_threshold,
         )
 
         with np.errstate(divide="ignore", invalid="ignore"):
-            iou = np.diag(self.confusion) / (
-                self.confusion.sum(axis=1)
-                + self.confusion.sum(axis=0)
-                - np.diag(self.confusion)
+            tp = confusion[:-1, :-1]
+            fp = confusion[:, -1]
+            iou = np.diag(tp) / (
+                tp.sum(axis=1) + tp.sum(axis=0) + fp.sum() - np.diag(tp)
             )
-        return ScalarResult(value=np.nanmean(iou), weight=annotation_img.size)  # type: ignore
+            non_taxonomy_classes = non_taxonomy_classes - {
+                confusion.shape[1] - 1
+            }
+            iou.put(list(non_taxonomy_classes), np.nan)
+            mean_iou = np.nanmean(iou)
+            return ScalarResult(value=mean_iou, weight=annotation_img.size)  # type: ignore
 
     def aggregate_score(self, results: List[MetricResult]) -> ScalarResult:
         return ScalarResult.aggregate(results)  # type: ignore
@@ -207,6 +309,7 @@ class SegmentationPrecision(SegmentationMaskMetric):
         prediction_filters: Optional[
             Union[ListOfOrAndFilters, ListOfAndFilters]
         ] = None,
+        iou_threshold: float = 0.5,
     ):
         """Calculates mean per-class precision
 
@@ -233,93 +336,37 @@ class SegmentationPrecision(SegmentationMaskMetric):
         super().__init__(
             annotation_filters,
             prediction_filters,
+            iou_threshold,
         )
 
     def _metric_impl(
         self,
         annotation_img: np.ndarray,
         prediction_img: np.ndarray,
-        annotation: Optional[SegmentationAnnotation],
-        prediction: Optional[SegmentationPrediction],
-    ) -> ScalarResult:
-        if annotation is None or prediction is None:
-            # TODO: Throw error when we wrap each item in try catch
-            return ScalarResult(0, weight=0)
-
-        self.confusion = self._calculate_confusion_matrix(
-            annotation, annotation_img, prediction, prediction_img
-        )
-
-        with np.errstate(divide="ignore", invalid="ignore"):
-            true_pos = np.diag(self.confusion)
-            precision = true_pos / np.sum(self.confusion, axis=0)
-            mean_precision = np.nanmean(precision)
-        return ScalarResult(value=mean_precision, weight=1)  # type: ignore
-
-    def aggregate_score(self, results: List[MetricResult]) -> ScalarResult:
-        return ScalarResult.aggregate(results)  # type: ignore
-
-
-class SegmentationAveragePrecision(SegmentationMaskMetric):
-    def __init__(
-        self,
-        annotation_filters: Optional[
-            Union[ListOfOrAndFilters, ListOfAndFilters]
-        ] = None,
-        prediction_filters: Optional[
-            Union[ListOfOrAndFilters, ListOfAndFilters]
-        ] = None,
-    ):
-        """Initializes SegmentationAveragePrecision object.
-
-        Args:
-            annotation_filters: Filter predicates. Allowed formats are:
-                ListOfAndFilters where each Filter forms a chain of AND predicates.
-                    or
-                ListOfOrAndFilters where Filters are expressed in disjunctive normal form (DNF), like
-                [[MetadataFilter("short_haired", "==", True), FieldFilter("label", "in", ["cat", "dog"]), ...].
-                DNF allows arbitrary boolean logical combinations of single field predicates. The innermost structures
-                each describe a single column predicate. The list of inner predicates is interpreted as a conjunction
-                (AND), forming a more selective `and` multiple field predicate.
-                Finally, the most outer list combines these filters as a disjunction (OR).
-            prediction_filters: Filter predicates. Allowed formats are:
-                ListOfAndFilters where each Filter forms a chain of AND predicates.
-                    or
-                ListOfOrAndFilters where Filters are expressed in disjunctive normal form (DNF), like
-                [[MetadataFilter("short_haired", "==", True), FieldFilter("label", "in", ["cat", "dog"]), ...].
-                DNF allows arbitrary boolean logical combinations of single field predicates. The innermost structures
-                each describe a single column predicate. The list of inner predicates is interpreted as a conjunction
-                (AND), forming a more selective `and` multiple field predicate.
-                Finally, the most outer list combines these filters as a disjunction (OR).
-        """
-        super().__init__(
-            annotation_filters,
-            prediction_filters,
-        )
-
-    def _metric_impl(
-        self,
-        annotation_img: np.ndarray,
-        prediction_img: np.ndarray,
-        annotation: Optional[SegmentationAnnotation],
-        prediction: Optional[SegmentationPrediction],
+        annotation: SegmentationAnnotation,
+        prediction: SegmentationPrediction,
     ) -> ScalarResult:
-        if annotation is None or prediction is None:
-            # TODO: Throw error when we wrap each item in try catch
-            return ScalarResult(0, weight=0)
-
-        self.confusion = self._calculate_confusion_matrix(
-            annotation, annotation_img, prediction, prediction_img
+        confusion, non_taxonomy_classes = self._calculate_confusion_matrix(
+            annotation,
+            annotation_img,
+            prediction,
+            prediction_img,
+            self.iou_threshold,
         )
 
         with np.errstate(divide="ignore", invalid="ignore"):
-            true_pos = np.diag(self.confusion)
-            precision = true_pos / np.sum(self.confusion, axis=0)
-            recall = true_pos / np.sum(self.confusion, axis=1)
-            average_precision = compute_average_precision(
-                np.nan_to_num(recall), np.nan_to_num(precision)
-            )
-        return ScalarResult(value=average_precision, weight=1)
+            # TODO(gunnar): Logic can be simplified
+            confused = confusion[:-1, :-1]
+            tp = confused.diagonal()
+            fp = confusion[:, -1][:-1] + confused.sum(axis=0) - tp
+            tp_and_fp = tp + fp
+            precision = tp / tp_and_fp
+            non_taxonomy_classes = non_taxonomy_classes - {
+                confusion.shape[1] - 1
+            }
+            precision.put(list(non_taxonomy_classes), np.nan)
+            avg_precision = np.nanmean(precision)
+        return ScalarResult(value=np.nan_to_num(avg_precision), weight=confusion.sum())  # type: ignore
 
     def aggregate_score(self, results: List[MetricResult]) -> ScalarResult:
         return ScalarResult.aggregate(results)  # type: ignore
@@ -337,6 +384,7 @@ class SegmentationRecall(SegmentationMaskMetric):
         prediction_filters: Optional[
             Union[ListOfOrAndFilters, ListOfAndFilters]
         ] = None,
+        iou_threshold: float = 0.5,
     ):
         """Initializes PolygonRecall object.
 
@@ -361,29 +409,33 @@ class SegmentationRecall(SegmentationMaskMetric):
                 Finally, the most outer list combines these filters as a disjunction (OR).
         """
         super().__init__(
-            annotation_filters=annotation_filters,
-            prediction_filters=prediction_filters,
+            annotation_filters,
+            prediction_filters,
+            iou_threshold,
         )
 
     def _metric_impl(
         self,
         annotation_img: np.ndarray,
         prediction_img: np.ndarray,
-        annotation: Optional[SegmentationAnnotation],
-        prediction: Optional[SegmentationPrediction],
+        annotation: SegmentationAnnotation,
+        prediction: SegmentationPrediction,
     ) -> ScalarResult:
-        if annotation is None or prediction is None:
-            # TODO: Throw error when we wrap each item in try catch
-            return ScalarResult(0, weight=0)
-
-        self.confusion = self._calculate_confusion_matrix(
-            annotation, annotation_img, prediction, prediction_img
+        confusion, non_taxonomy_classes = self._calculate_confusion_matrix(
+            annotation,
+            annotation_img,
+            prediction,
+            prediction_img,
+            self.iou_threshold,
         )
 
         with np.errstate(divide="ignore", invalid="ignore"):
-            true_pos = np.diag(self.confusion)
-            recall = np.nanmean(true_pos / np.sum(self.confusion, axis=1))
-        return ScalarResult(value=recall, weight=annotation_img.size)  # type: ignore
+            recall = confusion.diagonal() / confusion.sum(axis=1)
+            recall.put(
+                list(non_taxonomy_classes), np.nan
+            )  # We don't consider non taxonomy classes, i.e. FPs and background
+            mean_recall = np.nanmean(recall)
+        return ScalarResult(value=np.nan_to_num(mean_recall), weight=annotation_img.size)  # type: ignore
 
     def aggregate_score(self, results: List[MetricResult]) -> ScalarResult:
         return ScalarResult.aggregate(results)  # type: ignore
@@ -424,6 +476,8 @@ class SegmentationMAP(SegmentationMaskMetric):
         metric(annotations, predictions)
     """
 
+    iou_setups = {"coco"}
+
     # TODO: Remove defaults once these are surfaced more cleanly to users.
     def __init__(
         self,
@@ -433,6 +487,7 @@ class SegmentationMAP(SegmentationMaskMetric):
         prediction_filters: Optional[
             Union[ListOfOrAndFilters, ListOfAndFilters]
         ] = None,
+        iou_thresholds: Union[List[float], str] = "coco",
     ):
         """Initializes PolygonRecall object.
 
@@ -455,44 +510,42 @@ class SegmentationMAP(SegmentationMaskMetric):
                 each describe a single column predicate. The list of inner predicates is interpreted as a conjunction
                 (AND), forming a more selective `and` multiple field predicate.
                 Finally, the most outer list combines these filters as a disjunction (OR).
+            map_thresholds: Provide a list of threshold to compute over or literal "coco"
         """
         super().__init__(
-            annotation_filters=annotation_filters,
-            prediction_filters=prediction_filters,
+            annotation_filters,
+            prediction_filters,
         )
+        self.iou_thresholds = setup_iou_thresholds(iou_thresholds)
 
     def _metric_impl(
         self,
         annotation_img: np.ndarray,
         prediction_img: np.ndarray,
-        annotation: Optional[SegmentationAnnotation],
-        prediction: Optional[SegmentationPrediction],
+        annotation: SegmentationAnnotation,
+        prediction: SegmentationPrediction,
     ) -> ScalarResult:
-        if annotation is None or prediction is None:
-            # TODO: Throw error when we wrap each item in try catch
-            return ScalarResult(0, weight=0)
 
-        self.confusion = self._calculate_confusion_matrix(
-            annotation, annotation_img, prediction, prediction_img
-        )
-        label_to_index = {a.label: a.index for a in annotation.annotations}
-        num_classes = len(label_to_index.keys())
-        ap_per_class = np.ndarray(num_classes)  # type: ignore
-        with np.errstate(divide="ignore", invalid="ignore"):
-            for class_idx, (_, index) in enumerate(label_to_index.items()):
-                true_pos = self.confusion[index, index]
-                false_pos = self.confusion[:, index].sum()
-                samples = true_pos + false_pos
-                if samples:
-                    ap_per_class[class_idx] = true_pos / samples
-                else:
-                    ap_per_class[class_idx] = np.nan
-
-        if num_classes > 0:
-            m_ap = np.nanmean(ap_per_class)
-            return ScalarResult(m_ap, weight=1)  # type: ignore
-        else:
-            return ScalarResult(0, weight=0)
+        ap_per_threshold = []
+        weight = 0
+        for iou_threshold in self.iou_thresholds:
+            ap = SegmentationPrecision(
+                self.annotation_filters, self.prediction_filters, iou_threshold
+            )
+            ap.loader = self.loader
+            ap_result = ap(
+                AnnotationList(segmentation_annotations=[annotation]),
+                PredictionList(segmentation_predictions=[prediction]),
+            )
+            ap_per_threshold.append(ap_result.value)  # type: ignore
+            weight += ap_result.weight  # type: ignore
+
+        thresholds = np.concatenate([[0], self.iou_thresholds, [1]])
+        steps = np.diff(thresholds)
+        mean_ap = (
+            np.array(ap_per_threshold + [ap_per_threshold[-1]]) * steps
+        ).sum()
+        return ScalarResult(mean_ap, weight=weight)
 
     def aggregate_score(self, results: List[MetricResult]) -> ScalarResult:
         return ScalarResult.aggregate(results)  # type: ignore
@@ -542,6 +595,7 @@ class SegmentationFWAVACC(SegmentationMaskMetric):
         prediction_filters: Optional[
             Union[ListOfOrAndFilters, ListOfAndFilters]
         ] = None,
+        iou_threshold: float = 0.5,
     ):
         """Initializes SegmentationFWAVACC object.
 
@@ -566,33 +620,40 @@ class SegmentationFWAVACC(SegmentationMaskMetric):
                 Finally, the most outer list combines these filters as a disjunction (OR).
         """
         super().__init__(
-            annotation_filters=annotation_filters,
-            prediction_filters=prediction_filters,
+            annotation_filters,
+            prediction_filters,
+            iou_threshold,
         )
 
     def _metric_impl(
         self,
         annotation_img: np.ndarray,
         prediction_img: np.ndarray,
-        annotation: Optional[SegmentationAnnotation],
-        prediction: Optional[SegmentationPrediction],
+        annotation: SegmentationAnnotation,
+        prediction: SegmentationPrediction,
     ) -> ScalarResult:
-        if annotation is None or prediction is None:
-            # TODO: Throw error when we wrap each item in try catch
-            return ScalarResult(0, weight=0)
-
-        self.confusion = self._calculate_confusion_matrix(
-            annotation, annotation_img, prediction, prediction_img
+        confusion, non_taxonomy_classes = self._calculate_confusion_matrix(
+            annotation,
+            annotation_img,
+            prediction,
+            prediction_img,
+            self.iou_threshold,
         )
         with np.errstate(divide="ignore", invalid="ignore"):
-            iu = np.diag(self.confusion) / (
-                self.confusion.sum(axis=1)
-                + self.confusion.sum(axis=0)
-                - np.diag(self.confusion)
+            iu = np.diag(confusion) / (
+                confusion.sum(axis=1)
+                + confusion.sum(axis=0)
+                - np.diag(confusion)
             )
-            freq = self.confusion.sum(axis=1) / self.confusion.sum()
-            fwavacc = (freq[freq > 0] * iu[freq > 0]).sum()
-        return ScalarResult(value=np.nanmean(fwavacc), weight=1)  # type: ignore
+            predicted_counts = confusion.sum(axis=0).astype(np.float_)
+            predicted_counts.put(list(non_taxonomy_classes), np.nan)
+            freq = predicted_counts / np.nansum(predicted_counts)
+            iu.put(list(non_taxonomy_classes), np.nan)
+            fwavacc = (
+                np.nan_to_num(freq[freq > 0]) * np.nan_to_num(iu[freq > 0])
+            ).sum()
+            mean_fwavacc = np.nanmean(fwavacc)
+        return ScalarResult(value=np.nan_to_num(mean_fwavacc), weight=confusion.sum())  # type: ignore
 
     def aggregate_score(self, results: List[MetricResult]) -> ScalarResult:
         return ScalarResult.aggregate(results)  # type: ignore