code: def get_dataset(split, shape_soft_labels, shape_pointclouds=None):
      """Get dataset."""
      if shape_pointclouds:
          print(shape_pointclouds)
      tfrecord_path = os.path.join(FLAGS.tfrecords_dir, split)
      buffer_size, shuffle, cycles = 10000, True, 10000
      if FLAGS.debug:
          buffer_size, shuffle, cycles = 1, False, 1
      if FLAGS.val:
          buffer_size, shuffle, cycles = 100, False, 1
  
      tfrecords_pattern = io.expand_rio_pattern(tfrecord_path)
      dataset = tf.data.Dataset.list_files(
          tfrecords_pattern, shuffle=shuffle
      )
      dataset = dataset.interleave(
          tf.data.TFRecordDataset, cycle_length=cycles
      )
  
      if shuffle:
          dataset = dataset.shuffle(buffer_size=buffer_size)
  
      if "scannet" in tfrecord_path:
          dataset = dataset.map(
              extract_protos.decode_bytes_multiple_scannet
          )
          dataset = dataset.filter(
              lambda sample: sample["num_boxes"] == 3
          )
      else:
          dataset = dataset.map(extract_protos.decode_bytes_multiple)
          dataset = dataset.filter(
              lambda sample: tf.reduce_min(sample["shapes"]) > -1
          )
  
      def augment(sample):
          image = sample["image"]
          if tf.random.uniform([1], 0, 1.0) < 0.8:
              image = tf.image.random_saturation(image, 1.0, 10.0)
              image = tf.image.random_contrast(image, 0.05, 5.0)
              image = tf.image.random_hue(image, 0.5)
              image = tf.image.random_brightness(image, 0.8)
              sample["image"] = image
          if tf.random.uniform([1], 0, 1.0) < 0.5:
              sample["image"] = tf.image.flip_left_right(
                  sample["image"]
              )
              sample["translations_3d"] *= [[-1.0, 1.0, 1.0]]
              sample["rotations_3d"] = tf.reshape(
                  sample["rotations_3d"], [-1, 3, 3]
              )
              sample["rotations_3d"] = tf.transpose(
                  sample["rotations_3d"], perm=[0, 2, 1]
              )
              sample["rotations_3d"] = tf.reshape(
                  sample["rotations_3d"], [-1, 9]
              )
              bbox = sample["groundtruth_boxes"]
              bbox = tf.stack(
                  [
                      bbox[:, 0],
                      1 - bbox[:, 3],
                      bbox[:, 2],
                      1 - bbox[:, 1],
                  ],
                  axis=-1,
              )
              sample["groundtruth_boxes"] = bbox
          if FLAGS.gaussian_augmentation:
              if tf.random.uniform([1], 0, 1.0) < 0.15:
                  sample["image"] = tf_utils.gaussian_blur(
                      sample["image"], sigma=1
                  )
              elif tf.random.uniform([1], 0, 1.0) < 0.30:
                  sample["image"] = tf_utils.gaussian_blur(
                      sample["image"], sigma=2
                  )
              elif tf.random.uniform([1], 0, 1.0) < 0.45:
                  sample["image"] = tf_utils.gaussian_blur(
                      sample["image"], sigma=3
                  )
          return sample
  
      if FLAGS.train:  # and not FLAGS.debug:
          dataset = dataset.map(
              augment, num_parallel_calls=FLAGS.num_workers
          )
  
      def add_soft_shape_labels(sample):
          sample["shapes_soft"] = tf.map_fn(
              fn=lambda t: tf.cast(shape_soft_labels[t], tf.float32),
              elems=tf.cast(sample["shapes"], tf.int32),
              fn_output_signature=tf.float32,
          )
          return sample
  
      if FLAGS.soft_shape_labels:
          dataset = dataset.map(
              add_soft_shape_labels,
              num_parallel_calls=FLAGS.num_workers,
          )
  
      # Create dataset for overfitting when debugging
      if FLAGS.debug:
          t = FLAGS.num_overfitting_samples
          dataset = dataset.take(t)
          mult = 1 if FLAGS.val else 1000
          dataset = dataset.repeat(mult)
          if t > 1:
              dataset = dataset.shuffle(buffer_size)
  
      return dataset
  
imports: import os
  import pickle
  import time
  from absl import app
  from absl import flags
  from absl import logging
  import matplotlib.pyplot
  import numpy
  import tensorflow
  import tensorflow_graphics.projects.points_to_3Dobjects.data_preparation.extract_protos
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import collision_loss
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import cross_entropy_loss
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import focal_loss
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import focal_loss_sparse
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import learning_rate_schedule
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import regression_huber_loss
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import regression_l1_loss
  import tensorflow_graphics.projects.points_to_3Dobjects.models.centernet_vid
  from tensorflow_graphics.projects.points_to_3Dobjects.transforms import transforms_factory
  from tensorflow_graphics.projects.points_to_3Dobjects.utils import evaluator
  import tensorflow_graphics.projects.points_to_3Dobjects.utils.io
  import tensorflow_graphics.projects.points_to_3Dobjects.utils.logger
  import tensorflow_graphics.projects.points_to_3Dobjects.utils.plot
  import tensorflow_graphics.projects.points_to_3Dobjects.utils.tf_utils
  from google3.pyglib import gfile
label_desc: Use image node_id 4 for calling a global function with example usage: get_dataset(split, shape_soft_labels, shape_pointclouds) and returns: dataset, sample, sample