code: def get_custom_dataset(split, args):
      """Creates input data pipeline from custom PNG images."""
      with tf.device("/cpu:0"):
          files = glob.glob(args.train_glob)
          if not files:
              raise RuntimeError(
                  f"No training images found with glob "
                  f"'{args.train_glob}'."
              )
          dataset = tf.data.Dataset.from_tensor_slices(files)
          dataset = dataset.shuffle(
              len(files), reshuffle_each_iteration=True
          )
          if split == "train":
              dataset = dataset.repeat()
          dataset = dataset.map(
              lambda x: crop_image(read_png(x), args.patchsize),
              num_parallel_calls=args.preprocess_threads,
          )
          dataset = dataset.batch(args.batchsize, drop_remainder=True)
      return dataset
  
imports: import argparse
  import glob
  import sys
  from absl import app
  from absl.flags import argparse_flags
  import tensorflow
  import tensorflow_compression
  import tensorflow_datasets
label_desc: Use image node_id 6 for calling a global function with example usage: get_custom_dataset(split, args) and returns: dataset