code: def get_dataset(
      *,
      batch_size: int,
      eval_batch_size: int,
      num_shards: int,
      dtype_str: str = "float32",
      shuffle_seed: int = 0,
      rng: Optional[Tuple[int, int]] = None,
      dataset_configs: ml_collections.ConfigDict = ml_collections.ConfigDict(),
      dataset_service_address: Optional[str] = None,
  ) -> dataset_utils.Dataset:
      """Returns generators for train and validation splits of the specified dataset.
  
      Args:
        batch_size: Determines the train batch size.
        eval_batch_size: Determines the evaluation batch size.
        num_shards: Number of shards --> batch shape: [num_shards, bs, ...].
        dtype_str: Data type of the image (e.g. 'float32').
        shuffle_seed: Seed for shuffling the training data.
        rng: JAX rng key, which can be used for augmentation, shuffling, etc.
        dataset_configs: Dataset specific configurations.
        dataset_service_address: If set, will distribute the training dataset using
          the given tf.data service at the given address.
  
      Returns:
        A dataset_utils.Dataset() which includes a train_iter, a valid_iter,
          a test_iter, and a dict of meta_data.
      """
      del batch_size, shuffle_seed, dataset_service_address
  
      dtype = getattr(tf, dtype_str)
  
      if dataset_configs.name is None:
          raise ValueError("The name of the dataset must be specified")
      if dataset_configs.train_target_size is None:
          raise ValueError("Target size must be specified")
  
      denoise_configs = dataset_configs.get("denoise")
      dataset_info = datasets_info.get_info(dataset_configs.name)
      validation_split = dataset_configs.get(
          "validation_split", "validation"
      )
      num_eval_examples = dataset_utils.get_num_examples(
          dataset=dataset_info.tfds_name,
          split=validation_split,
          data_dir=dataset_info.data_dir,
      )
  
      logging.info(
          "Loading validation split of the %s dataset.",
          dataset_configs.name,
      )
      preprocess_ex_eval = functools.partial(
          preprocess_example,
          train=False,
          dtype=dtype,
          resize=dataset_configs.get(
              "eval_target_size", dataset_configs.train_target_size
          ),
          dataset_configs=dataset_configs,
          dataset_info=dataset_info,
          rng=int(rng[0]),
      )
  
      # TODO(kellybuchanan): merge cityscapes_c and ade20k_c
      eval_ds, _ = dataset_utils.load_split_from_tfds(
          dataset_info.tfds_name,
          eval_batch_size,
          split=validation_split,
          preprocess_example=preprocess_ex_eval,
          data_dir=dataset_info.data_dir,
      )
  
      maybe_pad_batches_eval = functools.partial(
          dataset_utils.maybe_pad_batch,
          train=False,
          batch_size=eval_batch_size,
          pixel_level=True,
      )
      shard_batches = functools.partial(
          dataset_utils.shard, n_devices=num_shards
      )
      # define classes to exclude
      # if ood classes are present use those to remap the classes
      class_to_exclude = (
          dataset_info.ood_classes
          if dataset_info.ood_classes
          else dataset_info.classes
      )
      exclude_classes = functools.partial(
          exclude_bad_classes,
          new_labels=get_post_exclusion_labels(class_to_exclude),
      )
  
      eval_iter = iter(eval_ds)
      eval_iter = map(dataset_utils.tf_to_numpy, eval_iter)
      eval_iter = map(maybe_pad_batches_eval, eval_iter)
      if not denoise_configs and dataset_configs.name != "pascal_voc":
          eval_iter = map(exclude_classes, eval_iter)
      eval_iter = map(shard_batches, eval_iter)
  
      input_shape = (
          (-1,) + tuple(dataset_configs.train_target_size) + (3,)
      )
  
      class_proportions = (
          get_class_proportions(
              dataset_info.classes, dataset_info.pixels_per_class
          )
          if dataset_info.pixels_per_class
          else None
      )
      meta_data = {
          "num_classes": 3
          if denoise_configs
          else len(
              [
                  c.id
                  for c in dataset_info.classes
                  if not c.ignore_in_eval
              ]
          ),
          "input_shape": input_shape,
          "num_train_examples": 0,
          "num_eval_examples": num_eval_examples,
          "input_dtype": getattr(jnp, dtype_str),
          "target_is_onehot": False,
          "class_names": get_class_names(dataset_info.classes),
          "class_colors": get_class_colors(dataset_info.classes),
          "class_proportions": class_proportions,
      }
      return dataset_utils.Dataset(None, eval_iter, None, meta_data)
  
file_path: scenic/scenic/projects/robust_segvit/datasets/segmentation_variants.py
imports: import functools
  from typing import Dict, List, Optional, Tuple
  from absl import logging
  import jax.numpy
  import ml_collections
  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
  from scenic.projects.robust_segvit.datasets.segmentation_datasets import exclude_bad_classes
  from scenic.projects.robust_segvit.datasets.segmentation_datasets import get_class_colors
  from scenic.projects.robust_segvit.datasets.segmentation_datasets import get_class_names
  from scenic.projects.robust_segvit.datasets.segmentation_datasets import get_class_proportions
  from scenic.projects.robust_segvit.datasets.segmentation_datasets import get_post_exclusion_labels
  import tensorflow
label_desc: Use image node_id 4 for calling a global function with example usage: get_dataset() and returns: dataset_utils