code: def augment_example(
      example, dtype=tf.float32, resize=None, **inception_crop_kws
  ):
      """Augments the given train image.
  
      Args:
        example: dict; Example coming from TFDS.
        dtype: Tensorflow data type; Data type of the image.
        resize: sequence; [H, W] to which image and labels should be resized.
        **inception_crop_kws: Keyword arguments passed on to
          inception_crop_with_mask.
  
      Returns:
        An example dict as required by the model.
      """
      image = example["inputs"]
      mask = example["label"][..., tf.newaxis]
  
      # Random crop and resize ("Inception crop"):
      image, mask = dataset_utils.inception_crop_with_mask(
          image,
          mask,
          resize_size=image.shape[-3:-1] if resize is None else resize,
          **inception_crop_kws,
      )
  
      # Random LR flip:
      seed = tf.random.uniform(
          shape=[2], maxval=2**31 - 1, dtype=tf.int32
      )
      image = tf.image.stateless_random_flip_left_right(image, seed)
      mask = tf.image.stateless_random_flip_left_right(mask, seed)
  
      image = tf.cast(image, dtype)
      mask = tf.cast(mask, dtype)
      mask = tf.squeeze(mask, axis=2)
      return {"inputs": image, "label": mask}
  
file_path: scenic/scenic/dataset_lib/cityscapes_dataset.py
imports: import collections
  import functools
  from typing import Optional
  from absl import logging
  import jax.numpy
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: augment_example(example, dtype, resize) and returns: dict