code: def compute_confusion_matrix_metrics(
      confusion_matrices: Sequence[Array],
      return_per_class_metrics: bool,
  ) -> Dict[str, float]:
      """Computes classification metrics from a confusion matrix.
  
      Computes the recall, precision and jaccard index (IoU) from the input
      confusion matrices. The confusion matrices are assumed to be of the form
      [ground_truth, predictions]. In other words, ground truth classes along the
      rows, and predicted classes along the columns.
  
      Args:
        confusion_matrices: Sequence of [n_batch, n_class, n_class] confusion
          matrices. The first two dimensions will be summed over to get an
          [n_class, n_class] matrix for further metrics.
        return_per_class_metrics: If true, return per-class metrics.
  
      Returns:
        A dictionary of metrics (recall, precision and jaccard index).
      """
  
      conf_matrix = np.sum(
          confusion_matrices, axis=0
      )  # Sum over eval batches.
      if conf_matrix.ndim != 3:
          raise ValueError(
              "Expecting confusion matrix to have shape "
              f"[batch_size, num_classes, num_classes], got {conf_matrix.shape}."
          )
      conf_matrix = np.sum(
          conf_matrix, axis=0
      )  # Sum over batch dimension.
      n_classes = conf_matrix.shape[0]
      metrics_dict = {}
  
      # We assume that the confusion matrix is [ground_truth x predictions].
      true_positives = np.diag(conf_matrix)
      sum_rows = np.sum(conf_matrix, axis=0)
      sum_cols = np.sum(conf_matrix, axis=1)
  
      recall_per_class = true_positives / sum_cols
      precision_per_class = true_positives / sum_rows
      jaccard_index_per_class = true_positives / (
          sum_rows + sum_cols - true_positives
      )
  
      metrics_dict["recall/mean"] = np.nanmean(recall_per_class)
      metrics_dict["precision/mean"] = np.nanmean(precision_per_class)
      metrics_dict["jaccard/mean"] = np.nanmean(jaccard_index_per_class)
  
      def add_per_class_results(metric: Array, name: str) -> None:
          for i in range(n_classes):
              # We set NaN values (from dividing by 0) to 0, to not cause problems with
              # logging.
              metrics_dict[f"{name}/{i}"] = np.nan_to_num(metric[i])
  
      if return_per_class_metrics:
          add_per_class_results(recall_per_class, "recall")
          add_per_class_results(precision_per_class, "precision")
          add_per_class_results(jaccard_index_per_class, "jaccard")
  
      return metrics_dict
  
file_path: scenic/scenic/projects/vivit/evaluation_lib.py
imports: import os
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
  from absl import logging
  import flax
  from flax.metrics import tensorboard
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import numpy
  from scenic.train_lib_deprecated import pretrain_utils
  from scenic.train_lib_deprecated import train_utils
  from sklearn.metrics import average_precision_score
  from tensorflow.io import gfile
label_desc: Use image node_id 3 for calling a global function with example usage: compute_confusion_matrix_metrics(confusion_matrices, return_per_class_metrics) and returns: metrics_dict