code: def preprocess_example(
      example: Dict[str, tf.Tensor],
      train: bool,
      dataset_configs: ml_collections.ConfigDict,
      dataset_info: datasets_info.DatasetInfo,
      dtype: tf.DType = tf.float32,
      resize: Optional[List[int]] = None,
      rng: int = 0,
  ):
      """Preprocesses the given image.
  
      Args:
        example: Example coming from TFDS.
        train: Whether to apply training-specific preprocessing or not.
        dataset_configs: Dataset configurations.
        dataset_info: Dataset specific information.
        dtype: Data type of the image.
        resize: Height and width to which image and labels should be resized.
        rng: Seed for sampling the noise for denoising.
  
      Returns:
        An example dict as required by the model.
      """
      image = dataset_utils.normalize(
          example[dataset_info.image_key], dtype
      )
      mask = example[dataset_info.label_key]
  
      # Resize test images (train images are cropped/resized during augmentation):
      if not train:
          if resize is not None:
              image = tf.image.resize(image, resize, "bilinear")
              mask = tf.image.resize(mask, resize, "nearest")
  
      # adding noise for training images is applied during augmentation
      image = tf.cast(image, dtype)
      if dataset_configs.denoise and not train:
          (
              noised_image,
              noise,
              timestep,
              gamma,
              patch,
          ) = denoise_utils.add_noise(
              image, rng, dataset_configs.denoise, dtype
          )
          example = {
              "inputs": noised_image,
              "label": noise,
              "image": image,
              "timestep": timestep,
              "gamma": gamma,
              "patch": patch,
          }
      else:
          mask = tf.cast(mask, dtype)
          mask = tf.squeeze(mask, axis=2)
          timestep = tf.constant(dataset_configs.use_timestep or 0)
          example = {
              "inputs": image,
              "label": mask,
              "timestep": timestep,
          }
  
      return example
  
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 1 for calling a global function with example usage: preprocess_example(example, train, dataset_configs, dataset_info, dtype, resize, rng) and returns: example