code: def load_synthetic_nerf_dataset(
      dataset_dir: str = DATASETDIR,
      dataset_name: str = "lego",
      split: str = "train",
      scale: float = 1.0,
      batch_size: int = 10,
      shuffle: bool = True,
  ):
      """Load the synthetic data used in the NeRF paper.
  
      Args:
          dataset_dir: Directory of the synthetic dataset.
          dataset_name: Name of the dataset, eg 'lego'.
          split: Training/validation/testing split.
          scale: Scale factor for the image size.
          batch_size: Batch size of the TF Dataset.
          shuffle: Shuffle the images.
  
      Returns:
          A tf.data.Dataset containing the images and camera parameters.
      """
      anno_file = os.path.join(
          dataset_dir, dataset_name, "transforms_{}.json".format(split)
      )
      with tf.io.gfile.GFile(anno_file, "r") as fp:
          anno = json.load(fp)
  
      init_width = 800
      height, width = int(800 * scale), int(800 * scale)
  
      images = np.zeros(
          (len(anno["frames"]), height, width, 4), dtype=np.float32
      )
      focals = np.zeros((len(anno["frames"]), 2), dtype=np.float32)
      principal_points = np.zeros(
          (len(anno["frames"]), 2), dtype=np.float32
      )
      transform_matrices = np.zeros(
          (len(anno["frames"]), 4, 4), dtype=np.float32
      )
  
      for i, frame in enumerate(anno["frames"]):
          image_name = os.path.join(
              dataset_dir, dataset_name, frame["file_path"] + ".png"
          )
          with tf.io.gfile.GFile(image_name, "rb") as f:
              im = Image.open(f)
              if scale != 1.0:
                  im = im.resize((width, height))
              im = (np.array(im) / 255.0).astype(np.float32)
              images[i, :, :, :] = im
  
          camera_angle_x = float(anno["camera_angle_x"])
          focal = (
              0.5 * init_width / np.tan(0.5 * camera_angle_x)
          ) * scale
          focals[i, :] = [focal, focal]
          principal_points[i, :] = [width / 2.0, height / 2.0]
  
          transform_matrix = np.array(frame["transform_matrix"])
          transform_matrices[i, :, :] = transform_matrix
  
      dataset = tf.data.Dataset.from_tensor_slices(
          (images, focals, principal_points, transform_matrices)
      )
      dataset = dataset.batch(batch_size)
      if shuffle:
          dataset = dataset.shuffle(100, reshuffle_each_iteration=True)
  
      return dataset, height, width
  
file_path: graphics/tensorflow_graphics/projects/radiance_fields/data_loaders.py
imports: import json
  import os
  from typing import Optional
  import numpy
  from PIL import Image
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: load_synthetic_nerf_dataset(dataset_dir, dataset_name, split, scale, batch_size, shuffle) and returns: dataset,  height,  width