code: def main(_):
      dataset, height, width = data_loaders.load_synthetic_nerf_dataset(
          dataset_dir=FLAGS.dataset_dir,
          dataset_name=FLAGS.dataset_name,
          split=FLAGS.split,
          scale=FLAGS.dataset_scale,
          batch_size=FLAGS.batch_size,
      )
  
      model = model_lib.NeRF(
          ray_samples_coarse=FLAGS.ray_samples_coarse,
          ray_samples_fine=FLAGS.ray_samples_fine,
          near=FLAGS.near,
          far=FLAGS.far,
          n_freq_posenc_xyz=FLAGS.n_freq_posenc_xyz,
          scene_bbox=tuple(
              [float(x) for x in FLAGS.scene_bbox.split(",")]
          ),
          n_freq_posenc_dir=FLAGS.n_freq_posenc_dir,
          n_filters=FLAGS.n_filters,
          white_background=FLAGS.white_background,
      )
      model.init_coarse_and_fine_models()
      model.init_optimizer(
          learning_rate=FLAGS.learning_rate,
          decay_steps=FLAGS.decay_steps,
      )
      model.init_checkpoint(checkpoint_dir=FLAGS.checkpoint_dir)
  
      for epoch in range(
          int(model.latest_epoch.numpy()), FLAGS.num_epochs
      ):
          epoch_loss = 0.0
          for (
              image,
              focal,
              principal_point,
              transform_matrix,
          ) in dataset:
              random_rays, random_pixels_xy = perspective.random_rays(
                  focal, principal_point, height, width, FLAGS.n_rays
              )
              random_rays = utils.change_coordinate_system(
                  random_rays, (0.0, 0.0, 0.0), (1.0, -1.0, -1.0)
              )
              (
                  rays_org,
                  rays_dir,
              ) = utils.camera_rays_from_transformation_matrix(
                  random_rays, transform_matrix
              )
              random_pixels_yx = tf.reverse(random_pixels_xy, axis=[-1])
              pixels = tf.gather_nd(
                  image, random_pixels_yx, batch_dims=1
              )
              pixels_rgb, pixels_a = tf.split(pixels, [3, 1], axis=-1)
              pixels_rgb = pixels_rgb * pixels_a + 1 - pixels_a
  
              dist_loss = model.train_step(
                  rays_org, rays_dir, pixels_rgb
              )
              epoch_loss += dist_loss.numpy()
              model.global_step.assign_add(1)
  
          with model.summary_writer.as_default():
              tf.summary.scalar("epoch_loss", epoch_loss, step=epoch)
          if epoch % 20 == 0:
              model.manager.save()
          logging.info("Epoch %d: %.3f.", epoch, epoch_loss)
          model.latest_epoch.assign(epoch + 1)
      model.manager.save()
  
file_path: graphics/tensorflow_graphics/projects/radiance_fields/nerf/train.py
imports: from absl import app
  from absl import flags
  from absl import logging
  import tensorflow
  import tensorflow_graphics.projects.radiance_fields.data_loaders
  import tensorflow_graphics.projects.radiance_fields.nerf.model
  import tensorflow_graphics.projects.radiance_fields.utils
  import tensorflow_graphics.rendering.camera.perspective
label_desc: Use image node_id 1 for calling a global function with example usage: main(_) without return types