code: def main(_):
      # ============================================================================
      # Load all colored voxels and place them in a volume table
      # ============================================================================
      color_voxels_tfrecord_dir = FLAGS.voxels_dir
      mode = "test"
      batch_size = 1
      tfrecord_pattern = os.path.join(
          color_voxels_tfrecord_dir, "{0}@100.tfrecord".format(mode)
      )
      mapping_function = voxel_proto_get
      dataset = tfrecords_to_dataset(
          tfrecord_pattern, mapping_function, batch_size, buffer_size=1
      )
      volume_table = {}
      for filename, voxels, _, _ in dataset:
          key = str(filename[0].numpy(), "utf-8")
          volume_table[key] = voxels.numpy()
  
      # ============================================================================
      # Load blender dataset
      # ============================================================================
      blender_tfrecord_dir = FLAGS.images_dir  #
      blender_tfrecord = os.path.join(
          blender_tfrecord_dir,
          "default_chairs_{0}.tfrecord".format(mode),
      )
      batch_size = 1
      mapping_function = blender_sample_proto_get
      blender_dataset = tfrecords_to_dataset(
          blender_tfrecord, mapping_function, batch_size, buffer_size=1
      )
  
      # The camera parameters are fixed for the dataset.
      focal = np.array([284.44446, 284.44446], dtype=np.float32)
      principal_point = np.array([128, 128.0], dtype=np.float32)
  
      # --------------------------------------------------------------------------
      # Where to save the estimated tf records.
      # WARNING: This will take up to 350GB!
      n_tfrecords = 100
      writer = [None] * n_tfrecords
      for tfrecord_id in range(n_tfrecords):
          writer_name = "{0}-{1:05d}-of-00100.tfrecord".format(
              mode, tfrecord_id
          )
          writer[tfrecord_id] = tf.io.TFRecordWriter(
              os.path.join(FLAGS.output_dir, writer_name)
          )
  
      for sample_id, sample in blender_dataset.enumerate():
          (
              filename_tf,
              image_data_tf,
              obj_name_tf,
              object_rotation_tf,
              object_translation_tf,
              floor_color_tf,
              light_position_tf,
              object_elevation_tf,
              camera_rotation_matrix_tf,
              camera_translation_vector_tf,
              _,
          ) = sample
  
          filename_np = str(filename_tf.numpy()[0], "utf-8")
          image_data_np = image_data_tf.numpy()[0]
          obj_name = str(obj_name_tf.numpy()[0], "utf-8")
          object_voxels = volume_table[obj_name]
          camera_rotation_matrix = camera_rotation_matrix_tf.numpy()
          camera_translation_vector = (
              camera_translation_vector_tf.numpy()
          )
          light_position = light_position_tf.numpy()[0]
          object_rotation = object_rotation_tf.numpy()
          object_translation = object_translation_tf.numpy()
          object_elevation = object_elevation_tf.numpy()
          ground_color = floor_color_tf.numpy() / 255.0
  
          # --------------------------------------------------------------------------
          # Place voxels in the scene
          # --------------------------------------------------------------------------
          object_rotation_v = object_rotation
          object_translation_v = (
              object_translation[:, 0, [1, 0, 2]] * BLENDER_SCALE
          )
          object_elevation_v = object_elevation
  
          ground_occupancy = np.zeros(
              (VOXEL_SIZE, VOXEL_SIZE, VOXEL_SIZE, 1), dtype=np.float32
          )
          ground_occupancy[-2, 1:-2, 1:-2, 0] = 1
          ground_voxel_color = (
              np.ones(
                  (VOXEL_SIZE, VOXEL_SIZE, VOXEL_SIZE, 3),
                  dtype=np.float32,
              )
              * ground_color
          )
          #  np.array(GROUND_COLOR, dtype=np.float32)
          ground_voxel_color = np.concatenate(
              [ground_voxel_color, ground_occupancy], axis=-1
          )
  
          scene_voxels = (
              object_voxels * (1 - ground_occupancy)
              + ground_voxel_color * ground_occupancy
          )
  
          euler_angles_x = np.deg2rad(
              180 - object_rotation_v
          ) * np.array([1, 0, 0], dtype=np.float32)
          euler_angles_y = np.deg2rad(
              90 - object_elevation_v
          ) * np.array([0, 1, 0], dtype=np.float32)
          translation_vector = object_translation_v / (DIAMETER * 0.5)
  
          interpolated_voxels = helpers.object_to_world(
              scene_voxels,
              euler_angles_x,
              euler_angles_y,
              translation_vector,
          )
  
          # --------------------------------------------------------------------------
          # Estimate the ground image
          # --------------------------------------------------------------------------
          ground_image, ground_alpha = helpers.generate_ground_image(
              IMAGE_SIZE,
              IMAGE_SIZE,
              focal,
              principal_point,
              camera_rotation_matrix,
              camera_translation_vector[:, :, 0],
              ground_color,
          )
  
          # --------------------------------------------------------------------------
          # Render the voxels in the image plane
          # --------------------------------------------------------------------------
          object_rotation_dvr = np.array(
              np.deg2rad(object_rotation), dtype=np.float32
          )
          object_translation_dvr = np.array(
              object_translation[..., [0, 2, 1]], dtype=np.float32
          )
          object_translation_dvr -= np.array(
              [0, 0, helpers.OBJECT_BOTTOM], dtype=np.float32
          )
  
          rerendering = helpers.render_voxels_from_blender_camera(
              object_voxels,
              object_rotation_dvr,
              object_translation_dvr,
              256,
              256,
              focal,
              principal_point,
              camera_rotation_matrix,
              camera_translation_vector,
              absorption_factor=1.0,
              cell_size=1.0,
              depth_min=3.0,
              depth_max=5.0,
              frustum_size=(128, 128, 128),
          )
          rerendering_image, rerendering_alpha = tf.split(
              rerendering, [3, 1], axis=-1
          )
  
          rerendering_image = tf.image.resize(
              rerendering_image, (256, 256)
          )
          rerendering_alpha = tf.image.resize(
              rerendering_alpha, (256, 256), method="nearest"
          )
  
          final_composite = (
              BACKGROUND_COLOR
              * (1 - rerendering_alpha)
              * (1 - ground_alpha)
              + ground_image * (1 - rerendering_alpha) * ground_alpha
              + rerendering_image * rerendering_alpha
          )
          final_composite = tf.cast(
              tf.squeeze(final_composite) * 255, tf.uint8
          )
          final_composite_bytes = tf.image.encode_jpeg(
              final_composite, quality=99
          )
  
          # --------------------------------------------------------------------------
          # Save functionality
          interpolated_voxels_np = (
              interpolated_voxels.numpy()[0] * 255
          ).astype(np.uint8)
          final_composite_bytes_np = final_composite_bytes.numpy()
  
          neural_voxel_sample = data_pb2.NeuralVoxelPlusSample()
          neural_voxel_sample.name = filename_np
          neural_voxel_sample.voxel_data.extend(
              [interpolated_voxels_np.tobytes()]
          )
          neural_voxel_sample.rerendering_data.extend(
              [final_composite_bytes_np]
          )
          neural_voxel_sample.image_data.extend([image_data_np])
          neural_voxel_sample.light_position.extend(
              list(light_position)
          )
          tfrecord_id = sample_id // np.ceil(
              DATASET_SIZE[mode] / n_tfrecords
          )
          writer[tfrecord_id].write(
              neural_voxel_sample.SerializeToString()
          )
  
      for tfrecord_id in range(n_tfrecords):
          writer[tfrecord_id].close()
  
file_path: graphics/tensorflow_graphics/projects/neural_voxel_renderer/prepare_tfrecords/generate_tfrecords_nvr_plus.py
imports: import os
  import re
  from absl import app
  from absl import flags
  import numpy
  import tensorflow
  from tensorflow_graphics.projects.neural_voxel_renderer import helpers
  from tensorflow_graphics.projects.neural_voxel_renderer.prepare_tfrecords import data_pb2
label_desc: Use image node_id 5 for calling a global function with example usage: main(_) without return types