code: def get_instance_mask(instance_segmentation):
      """Obtain the instance mask from the blue channel of the segmentation file."""
      # Based on DevKit:
      # https://github.com/CSAILVision/ADE20K/blob/main/utils/utils_ade20k.py
      instance_segmentation_blue = instance_segmentation[:, :, 2]
      instance_mask = unique_with_inverse(instance_segmentation_blue)
      instance_mask = tf.expand_dims(
          tf.reshape(
              instance_mask, tf.shape(instance_segmentation_blue)
          ),
          -1,
      )
      return tf.cast(instance_mask, tf.uint16)
  
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 2 for calling a global function with example usage: get_instance_mask(instance_segmentation) and returns: tf