code: def build_dataset(
      data_root,
      num_images_per_epoch,
      num_images_per_input_file,
      batch_size,
      read_from_images=True,
  ):
      if read_from_images:
          ds = ray.data.read_images(
              data_root,
              # Use the same partitioning required by torch dataloader.
              # root_dir
              #   class_name1
              #     XXX.jpg
              #   class_name2
              #     YYY.jpg
              partitioning=Partitioning(
                  "dir", field_names=["label"], base_dir="~/data"
              ),
          )
  
          classes = {
              label: i for i, label in enumerate(ds.unique("label"))
          }
  
          def convert_class_to_idx(df, classes):
              df["label"] = df["label"].map(classes).astype("float32")
              return df
  
          ds = ds.map_batches(
              convert_class_to_idx,
              fn_kwargs={"classes": classes},
          )
          ds = ds.map(crop_and_flip_image)
      else:
          filenames = get_tfrecords_filenames(
              data_root, num_images_per_epoch, num_images_per_input_file
          )
          ds = ray.data.read_tfrecords(filenames)
          ds = ds.map_batches(
              decode_crop_and_flip_tf_record_batch,
              batch_size=batch_size,
              batch_format="pandas",
          )
  
      ds = ds.limit(num_images_per_epoch)
      return ds
  
imports: from collections import defaultdict
  import tensorflow
  import numpy
  import os
  import pandas
  import time
  import logging
  import csv
  import json
  import torchvision
  import torch
  import ray
  from ray.train.tensorflow import prepare_dataset_shard, TensorflowTrainer
  from ray.train import DataConfig, ScalingConfig
  from ray import train, tune
  from ray.tune import Tuner
  from ray.data.datasource.partitioning import Partitioning
  from tf_utils import DEFAULT_IMAGE_SIZE, NUM_CHANNELS, preprocess_image, build_tf_dataset
  from metric_utils import determine_if_memory_monitor_is_enabled_in_latest_session, get_ray_spilled_and_restored_mb, MaxMemoryUtilizationTracker
label_desc: Use image node_id 9 for calling a global function with example usage: build_dataset(data_root, num_images_per_epoch, num_images_per_input_file, batch_size, read_from_images) and returns: ds, df