code: def exclude_bad_classes(batch, new_labels):
      """Adjusts masks and batch_masks to exclude void and rare classes.
  
      This must be applied after dataset_utils.maybe_pad_batch() because we also
      update the batch_mask. Note that the data is already converted to Numpy by
      then.
  
      Args:
        batch: dict; Batch of data examples.
        new_labels: nd-array; array of length num_old_classes, containing new
          labels.
  
      Returns:
        Updated batch dict.
      """
      # Convert old labels to new labels:
      batch["label"] = new_labels[batch["label"].astype(np.int32)]
  
      # Set batch_mask to 0 at pixels that have an excluded label:
      mask_dtype = batch["batch_mask"].dtype
      batch["batch_mask"] = batch["batch_mask"].astype(np.bool_) & (
          batch["label"] != -1
      )
      batch["batch_mask"] = batch["batch_mask"].astype(mask_dtype)
  
      return batch
  
file_path: scenic/scenic/projects/robust_segvit/datasets/segmentation_datasets.py
imports: import functools
  from typing import Dict, List, Optional, Tuple
  from absl import logging
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.projects.robust_segvit.datasets import datasets_info
  from scenic.projects.robust_segvit.datasets import denoise_utils
  import tensorflow
label_desc: Use image node_id 7 for calling a global function with example usage: exclude_bad_classes(batch, new_labels) and returns: batch