code: def filter_confusion_matrix(
      confusion_matrix: pd.DataFrame, number_of_categories: int
  ) -> t.Tuple[np.ndarray, int]:
      pq = PriorityQueue()
      for row, values in enumerate(confusion_matrix):
          for col, value in enumerate(values):
              if row != col:
                  pq.put((-value, (row, col)))
      categories = set()
      while not pq.empty():
          if len(categories) >= number_of_categories:
              break
          _, (row, col) = pq.get()
          categories.add(row)
          categories.add(col)
      categories = sorted(categories)
      return (
          confusion_matrix[np.ix_(categories, categories)],
          categories,
      )
  
imports: import typing
  from collections import defaultdict
  from queue import PriorityQueue
  import numpy
  import pandas
  from deepchecks.core import CheckResult, DatasetKind
  from deepchecks.utils.abstracts.confusion_matrix_abstract import create_confusion_matrix_figure
  from deepchecks.vision._shared_docs import docstrings
  from deepchecks.vision.base_checks import SingleDatasetCheck
  from deepchecks.vision.context import Context
  from deepchecks.vision.metrics_utils.iou_utils import jaccard_iou
  from deepchecks.vision.vision_data import TaskType
  from deepchecks.vision.vision_data.batch_wrapper import BatchWrapper
label_desc: Use image node_id 1 for calling a global function with example usage: filter_confusion_matrix(confusion_matrix, number_of_categories) and returns: confusion_matrix,  categories