Repository: deepchecks/deepchecks

Issue:
[BUG][CV] Confusion matrix does not support unseen classes
**Describe the bug**
For object detection:
The model might return classes that we haven't seen before from our dataset. For example, our dataset has class ids between 0 to 100. The model can return suddenly 120, and we don't have a row for it in the matrix. Right now, we just put all the "unseen" classes under the same cateogry.

**To Reproduce**
Steps to reproduce the behavior:
Limit coco dataset to only a few classes and run the check on it

**Expected behavior**
Have each unseen classes in a row of its own.



Discussion:
This can be done by either having the matrix data structure be a pandas data frame with the class as index, or better a sparse matrix.

Resolving patch:
diff --git a/deepchecks/vision/checks/performance/confusion_matrix.py b/deepchecks/vision/checks/performance/confusion_matrix.py
--- a/deepchecks/vision/checks/performance/confusion_matrix.py
+++ b/deepchecks/vision/checks/performance/confusion_matrix.py
@@ -9,15 +9,19 @@
 # ----------------------------------------------------------------------------
 #
 """Module containing confusion matrix report check."""
+import pandas as pd
 import numpy as np
+import torch
 from plotly.express import imshow
 from queue import PriorityQueue
+from collections import defaultdict
 
 from deepchecks.core import CheckResult, DatasetKind
 from deepchecks.vision import SingleDatasetCheck, Context, Batch
 from deepchecks.vision.vision_data import TaskType
 from deepchecks.vision.metrics_utils.iou_utils import jaccard_iou
 
+
 __all__ = ['ConfusionMatrixReport']
 
 
@@ -72,33 +76,14 @@ def __init__(self,
     def initialize_run(self, context: Context, dataset_kind: DatasetKind = None):
         """Initialize run by creating an empty matrix the size of the data."""
         context.assert_task_type(TaskType.CLASSIFICATION, TaskType.OBJECT_DETECTION)
-
-        if dataset_kind == DatasetKind.TRAIN:
-            dataset = context.train
-        else:
-            dataset = context.test
+        dataset = context.get_data_by_kind(dataset_kind)
         self.task_type = dataset.task_type
-
-        # In case of object detection add last category for "not found label/prediction overlap"
-        if self.task_type == TaskType.OBJECT_DETECTION:
-            # In detection, we might have non-consecutive ids. For example, we might have class ids 10, 14, 20. So we
-            # will use the class list as a map of matrix id to class id: 0: 10, 1: 14, 2: 20
-            self.classes_list = sorted([int(x) for x in dataset.n_of_samples_per_class.keys()])
-            # Adding 2 extra categories. One for predictions with unseen before classes, and second for label and
-            # prediction that have no overlapping prediction/label
-            matrix_size = len(self.classes_list)
-            self.not_found_idx = matrix_size - 1
-            self.unseen_class_idx = matrix_size - 2
-        else:
-            matrix_size = dataset.num_classes
-
-        self.matrix = np.zeros((matrix_size, matrix_size))
+        self.matrix = defaultdict(lambda: defaultdict(int))
 
     def update(self, context: Context, batch: Batch, dataset_kind: DatasetKind = DatasetKind.TRAIN):
         """Add batch to confusion matrix."""
         labels = batch.labels
         predictions = batch.predictions
-
         if self.task_type == TaskType.CLASSIFICATION:
             self.update_classification(predictions, labels)
         elif self.task_type == TaskType.OBJECT_DETECTION:
@@ -106,122 +91,132 @@ def update(self, context: Context, batch: Batch, dataset_kind: DatasetKind = Dat
 
     def compute(self, context: Context, dataset_kind: DatasetKind = None) -> CheckResult:
         """Compute and plot confusion matrix after all batches were processed."""
-        if dataset_kind == DatasetKind.TRAIN:
-            dataset = context.train
-        else:
-            dataset = context.test
-        display_confusion_matrix, categories = filter_confusion_matrix(self.matrix, self.categories_to_display)
+        assert self.matrix is not None
+
+        dataset = context.get_data_by_kind(dataset_kind)
+        matrix = pd.DataFrame(self.matrix).T
+        matrix.replace(np.nan, 0, inplace=True)
+
+        classes = sorted(
+            set(matrix.index).union(set(matrix.columns)),
+            key=lambda x: np.inf if isinstance(x, str) else x
+        )
 
-        description = []
+        matrix = pd.DataFrame(matrix, index=classes, columns=classes).to_numpy()
+
+        confusion_matrix, categories = filter_confusion_matrix(
+            matrix,
+            self.categories_to_display
+        )
+
+        description = [f'Showing {self.categories_to_display} of {dataset.num_classes} classes:']
+        classes_to_display = []
+        classes_map = dict(enumerate(classes))  # class index -> class label
 
-        display_categories = []
-        add_not_found_category = False
-        add_unseen_category = False
         for category in categories:
-            if self.not_found_idx == category:
-                add_not_found_category = True
-            elif self.unseen_class_idx == category:
-                add_unseen_category = True
+            category = classes_map[category]
+            if category == 'no-overlapping':
+                description.append(
+                    '"No overlapping" categories are labels and prediction which did not have a matching '
+                    'label/prediction. For example a predictions that did not have a sufficiently overlapping '
+                    'label bounding box will appear under the "No overlapping label" category'
+                )
+                classes_to_display.append('no-overlapping')
+            elif isinstance(category, int):
+                classes_to_display.append(dataset.label_id_to_name(category))
             else:
-                display_categories.append(self.matrix_id_to_class_name(category, dataset))
-
-        x = display_categories
-        y = x.copy()
-
-        if add_not_found_category:
-            description += ['"No overlapping" categories are labels and prediction which did not have a matching '
-                            'label/prediction. For example a predictions that did not have a sufficiently overlapping '
-                            'label bounding box will appear under the "No overlapping label" category']
-            x += ['No overlapping prediction']
-            y += ['No overlapping label']
-        if add_unseen_category:
-            description += ['Unseen classes are classes that did not appear in the dataset labels, but did were '
-                            'predicted to be present.']
-            x += ['Unseen classes']
-            y += ['Unseen classes']
-
-        description += [f'Showing {self.categories_to_display} of {dataset.num_classes} classes:']
-
-        fig = imshow(display_confusion_matrix,
-                     x=x,
-                     y=y,
-                     text_auto=True)
-
-        fig.update_layout(width=600, height=600)
-        fig.update_xaxes(title='Predicted Value', type='category')
-        fig.update_yaxes(title='True value', type='category')
-
+                raise RuntimeError(
+                    'Internal Error! categories list must '
+                    'contain items of type - Union[int, Literal["no-overlapping"]]'
+                )
+
+        x = []
+        y = []
+
+        for it in classes_to_display:
+            if it != 'no-overlapping':
+                x.append(it)
+                y.append(it)
+            else:
+                x.append('No overlapping prediction')
+                y.append('No overlapping label')
+
+        description.append(
+            imshow(
+                confusion_matrix,
+                x=x,
+                y=y,
+                text_auto=True)
+            .update_layout(width=600, height=600)
+            .update_xaxes(title='Predicted Value', type='category')
+            .update_yaxes(title='True value', type='category')
+        )
         return CheckResult(
-            self.matrix,
+            matrix,
             header='Confusion Matrix',
-            display=[*description, fig]
+            display=description
         )
 
-    def class_id_to_matrix_id(self, class_id):
-        """Convert a class id to its id in the matrix."""
-        try:
-            return self.classes_list.index(class_id)
-        # If the class is not in the labels list (detection returned class not in the list) then return not found
-        except ValueError:
-            return self.unseen_class_idx
-
-    def matrix_id_to_class_name(self, matrix_id, dataset):
-        """Convert matrix id to the real class id, and return its name if possible."""
-        class_id = self.classes_list[matrix_id] if self.classes_list else matrix_id
-        return dataset.label_id_to_name(class_id)
-
     def update_object_detection(self, predictions, labels):
         """Update the confusion matrix by batch for object detection task."""
+        assert self.matrix is not None
+
         for image_detections, image_labels in zip(predictions, labels):
             detections_passed_threshold = [
-                detection for detection in image_detections if detection[4] > self.confidence_threshold
+                detection for detection in image_detections
+                if detection[4] > self.confidence_threshold
             ]
+
             if len(detections_passed_threshold) == 0:
                 # detections are empty, update matrix for labels
                 for label in image_labels:
-                    gt_class = self.class_id_to_matrix_id(int(label[0].item()))
-                    self.matrix[gt_class, self.not_found_idx] += 1
+                    label_class = int(label[0].item())
+                    self.matrix[label_class]['no-overlapping'] += 1
                 continue
 
-            all_ious = np.zeros((len(image_labels), len(detections_passed_threshold)))
-
-            for label_index, label in enumerate(image_labels):
-                for detected_index, detected in enumerate(detections_passed_threshold):
-                    all_ious[label_index, detected_index] = jaccard_iou(detected, label)
-
-            want_idx = np.where(all_ious > self.iou_threshold)
-
-            all_matches = [[want_idx[0][i], want_idx[1][i], all_ious[want_idx[0][i], want_idx[1][i]]]
-                           for i in range(want_idx[0].shape[0])]
-            all_matches = np.array(all_matches)
+            list_of_ious = (
+                (label_index, detected_index, jaccard_iou(detected, label))
+                for label_index, label in enumerate(image_labels)
+                for detected_index, detected in enumerate(detections_passed_threshold)
+            )
+            matches = np.array([
+                [label_index, detected_index, ious]
+                for label_index, detected_index, ious in list_of_ious
+                if ious > self.iou_threshold
+            ])
 
             # remove duplicate matches
-            if all_matches.shape[0] > 0:
-                all_matches = all_matches[all_matches[:, 2].argsort()[::-1]]
+            if len(matches) > 0:
+                # sort by ious, in descend order
+                matches = matches[matches[:, 2].argsort()[::-1]]
+                # leave matches with unique prediction and the highest ious
+                matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
+                # sort by ious, in descend order
+                matches = matches[matches[:, 2].argsort()[::-1]]
+                # leave matches with unique label and the highest ious
+                matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
+
+            n_of_matches = len(matches)
 
-                all_matches = all_matches[np.unique(all_matches[:, 1], return_index=True)[1]]
-
-                all_matches = all_matches[all_matches[:, 2].argsort()[::-1]]
-
-                all_matches = all_matches[np.unique(all_matches[:, 0], return_index=True)[1]]
-
-            for i, label in enumerate(image_labels):
-                gt_class = self.class_id_to_matrix_id(int(label[0]))
-                if all_matches.shape[0] > 0 and all_matches[all_matches[:, 0] == i].shape[0] == 1:
-                    detection_class = int(image_detections[int(all_matches[all_matches[:, 0] == i, 1][0])][5])
-                    detection_class = self.class_id_to_matrix_id(detection_class)
-                    self.matrix[gt_class, detection_class] += 1
+            for label_index, label in enumerate(image_labels):
+                label_class = int(label[0])
+                if n_of_matches > 0 and (matches[:, 0] == label_index).any():
+                    detection_index = int(matches[matches[:, 0] == label_index, 1][0])
+                    detected_class = int(image_detections[detection_index][5])
+                    self.matrix[label_class][detected_class] += 1
                 else:
-                    self.matrix[gt_class, self.not_found_idx] += 1
+                    self.matrix[label_class]['no-overlapping'] += 1
 
-            for i, detection in enumerate(image_detections):
-                if all_matches.shape[0] and all_matches[all_matches[:, 1] == i].shape[0] == 0:
-                    detection_class = self.class_id_to_matrix_id(int(detection[5]))
-                    self.matrix[self.not_found_idx, detection_class] += 1
+            for detection_index, detection in enumerate(detections_passed_threshold):
+                if n_of_matches > 0 and not (matches[:, 1] == detection_index).any():
+                    detected_class = int(detection[5])
+                    self.matrix['no-overlapping'][detected_class] += 1
 
     def update_classification(self, predictions, labels):
         """Update the confusion matrix by batch for classification task."""
+        assert self.matrix is not None
+
         for predicted_classes, image_labels in zip(predictions, labels):
             detected_class = max(range(len(predicted_classes)), key=predicted_classes.__getitem__)
-
-            self.matrix[image_labels, detected_class] += 1
+            label_class = image_labels.item() if isinstance(image_labels, torch.Tensor) else image_labels
+            self.matrix[label_class][detected_class] += 1


Test patch:
diff --git a/tests/vision/checks/performance/confusion_matrix_test.py b/tests/vision/checks/performance/confusion_matrix_test.py
--- a/tests/vision/checks/performance/confusion_matrix_test.py
+++ b/tests/vision/checks/performance/confusion_matrix_test.py
@@ -10,11 +10,13 @@
 #
 """Test functions of the VISION confusion matrix."""
 
-from hamcrest import assert_that, has_entries, close_to, equal_to, raises, calling
-
+from hamcrest import assert_that, equal_to, less_than_or_equal_to as le
 from deepchecks.vision.checks import ConfusionMatrixReport
 
 
+# TODO: more tests
+
+
 def test_classification(mnist_dataset_train, trained_mnist, device):
     # Arrange
     check = ConfusionMatrixReport()
@@ -35,4 +37,5 @@ def test_detection(coco_train_visiondata, trained_yolov5_object_detection, devic
                        device=device)
 
     # Assert
-    assert_that(result.value.shape, equal_to((61, 61)))
+    num_of_classes = coco_train_visiondata.num_classes + 1 # plus no-overlapping
+    assert_that(result.value.shape, le((num_of_classes, num_of_classes)))