code: def non_max_suppression_multiclass(
      class_indices: np.ndarray,
      scores: np.ndarray,
      segments: np.ndarray,
      config: ml_collections.ConfigDict,
  ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
      """Performs multiclass non-maximal suppression.
  
      Args:
        class_indices: Predicted class indices in shape (N,).
        scores: Predicted class scores in shape (N,).
        segments: Predicted temporal segments in shape (N, 2).
        config: NMS configs.
  
      Returns:
        class indices, scores, and segments after NMS.
      """
  
      out_class_indices = []
      out_scores = []
      out_segments = []
  
      for cls_idx in range(config.dataset_configs.num_classes):
          mask = class_indices == cls_idx
          cur_class_indices = class_indices[mask]
          if cur_class_indices.size:
              cur_segments = segments[mask]
              cur_scores = scores[mask]
              (
                  selected_indices,
                  selected_scores,
              ) = tf.image.non_max_suppression_with_scores(
                  make_2d_boxes(cur_segments),
                  cur_scores,
                  max_output_size=config.get("max_detections", 100),
                  iou_threshold=config.get("iou_threshold", 0.5),
                  score_threshold=config.get("score_threshold", 0.001),
                  soft_nms_sigma=config.get("soft_nms_sigma", 0.3),
              )
              selected_indices = selected_indices.numpy()
              selected_scores = selected_scores.numpy()
              out_class_indices.append(
                  cur_class_indices[selected_indices]
              )
              out_scores.append(selected_scores)
              out_segments.append(cur_segments[selected_indices])
      out_class_indices = (
          np.concatenate(out_class_indices, axis=0)
          if out_class_indices
          else np.array([], dtype=np.int32)
      )
      out_scores = (
          np.concatenate(out_scores, axis=0)
          if out_scores
          else np.array([], dtype=np.float32)
      )
      out_segments = (
          np.concatenate(out_segments, axis=0)
          if out_segments
          else np.array([], dtype=np.float32)
      )
      return out_class_indices, out_scores, out_segments
  
file_path: scenic/scenic/projects/unloc/postprocessing_utils.py
imports: from typing import Any, List, Optional, Tuple, Union
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.unloc import model_utils
  import tensorflow
label_desc: Use image node_id 4 for calling a global function with example usage: non_max_suppression_multiclass(class_indices, scores, segments, config) and returns: out_class_indices,  out_scores,  out_segments