Main training script

Summary: Implements the training script of NeRF.

Reviewed By: nikhilaravi

Differential Revision: D25684439

fbshipit-source-id: 8b19b6dc282eb6bf6e46ec4476bb0f13a84c90dd

Author: davnov134

projects/nerf/__init__.py

@@ -0,0 +1,44 @@
+seed: 3
+resume: True
+stats_print_interval: 10
+validation_epoch_interval: 150
+checkpoint_epoch_interval: 150
+checkpoint_path: 'checkpoints/fern_pt3d.pth'
+data:
+  dataset_name: 'fern'
+  image_size: [378, 504] # [height, width]
+  precache_rays: True
+test:
+  mode: 'evaluation'
+  trajectory_type: 'figure_eight'
+  up: [0.0, 1.0, 0.0]
+  scene_center: [0.0, 0.0, -2.0]
+  n_frames: 100
+  fps: 20
+optimizer:
+  max_epochs: 37500
+  lr: 0.0005
+  lr_scheduler_step_size: 12500
+  lr_scheduler_gamma: 0.1
+visualization:
+  history_size: 10
+  visdom: True
+  visdom_server: 'localhost'
+  visdom_port: 8097
+  visdom_env: 'nerf_pytorch3d'
+raysampler:
+  n_pts_per_ray: 64
+  n_pts_per_ray_fine: 64
+  n_rays_per_image: 1024
+  min_depth: 1.2
+  max_depth: 6.28
+  stratified: True
+  stratified_test: False
+  chunk_size_test: 6000
+implicit_function:
+  n_harmonic_functions_xyz: 10
+  n_harmonic_functions_dir: 4
+  n_hidden_neurons_xyz: 256
+  n_hidden_neurons_dir: 128
+  density_noise_std: 0.0
+  n_layers_xyz: 8

projects/nerf/configs/fern.yaml

@@ -0,0 +1,44 @@
+seed: 3
+resume: True
+stats_print_interval: 10
+validation_epoch_interval: 30
+checkpoint_epoch_interval: 30
+checkpoint_path: 'checkpoints/lego_pt3d.pth'
+data:
+  dataset_name: 'lego'
+  image_size: [800, 800] # [height, width]
+  precache_rays: True
+test:
+  mode: 'evaluation'
+  trajectory_type: 'circular'
+  up: [0.0, 0.0, 1.0]
+  scene_center: [0.0, 0.0, 0.0]
+  n_frames: 100
+  fps: 20
+optimizer:
+  max_epochs: 20000
+  lr: 0.0005
+  lr_scheduler_step_size: 5000
+  lr_scheduler_gamma: 0.1
+visualization:
+  history_size: 10
+  visdom: True
+  visdom_server: 'localhost'
+  visdom_port: 8097
+  visdom_env: 'nerf_pytorch3d'
+raysampler:
+  n_pts_per_ray: 64
+  n_pts_per_ray_fine: 64
+  n_rays_per_image: 1024
+  min_depth: 2.0
+  max_depth: 6.0
+  stratified: True
+  stratified_test: False
+  chunk_size_test: 6000
+implicit_function:
+  n_harmonic_functions_xyz: 10
+  n_harmonic_functions_dir: 4
+  n_hidden_neurons_xyz: 256
+  n_hidden_neurons_dir: 128
+  density_noise_std: 0.0
+  n_layers_xyz: 8

projects/nerf/configs/lego.yaml

@@ -0,0 +1,44 @@
+seed: 3
+resume: True
+stats_print_interval: 10
+validation_epoch_interval: 30
+checkpoint_epoch_interval: 30
+checkpoint_path: 'checkpoints/pt3logo_pt3d.pth'
+data:
+  dataset_name: 'pt3logo'
+  image_size: [512, 1024] # [height, width]
+  precache_rays: True
+test:
+  mode: 'export_video'
+  trajectory_type: 'figure_eight'
+  up: [0.0, -1.0, 0.0]
+  scene_center: [0.0, 0.0, 0.0]
+  n_frames: 100
+  fps: 20
+optimizer:
+  max_epochs: 100000
+  lr: 0.0005
+  lr_scheduler_step_size: 10000
+  lr_scheduler_gamma: 0.1
+visualization:
+  history_size: 20
+  visdom: True
+  visdom_server: 'localhost'
+  visdom_port: 8097
+  visdom_env: 'nerf_pytorch3d'
+raysampler:
+  n_pts_per_ray: 64
+  n_pts_per_ray_fine: 64
+  n_rays_per_image: 1024
+  min_depth: 8.0
+  max_depth: 23.0
+  stratified: True
+  stratified_test: False
+  chunk_size_test: 6000
+implicit_function:
+  n_harmonic_functions_xyz: 10
+  n_harmonic_functions_dir: 4
+  n_hidden_neurons_xyz: 256
+  n_hidden_neurons_dir: 128
+  density_noise_std: 0.0
+  n_layers_xyz: 8

projects/nerf/configs/pt3logo.yaml

@@ -2,8 +2,11 @@
 from typing import Tuple, List, Optional
 
 import torch
-from pytorch3d.renderer import ImplicitRenderer
+from pytorch3d.renderer import ImplicitRenderer, ray_bundle_to_ray_points
 from pytorch3d.renderer.cameras import CamerasBase
+from pytorch3d.structures import Pointclouds
+from pytorch3d.vis.plotly_vis import plot_scene
+from visdom import Visdom
 
 from .implicit_function import NeuralRadianceField
 from .raymarcher import EmissionAbsorptionNeRFRaymarcher
@@ -357,3 +360,68 @@ def forward(
                     )
 
         return out, metrics
+
+
+def visualize_nerf_outputs(
+    nerf_out: dict, output_cache: List, viz: Visdom, visdom_env: str
+):
+    """
+    Visualizes the outputs of the `RadianceFieldRenderer`.
+
+    Args:
+        nerf_out: An output of the validation rendering pass.
+        output_cache: A list with outputs of several training render passes.
+        viz: A visdom connection object.
+        visdom_env: The name of visdom environment for visualization.
+    """
+
+    # Show the training images.
+    ims = torch.stack([o["image"] for o in output_cache])
+    ims = torch.cat(list(ims), dim=1)
+    viz.image(
+        ims.permute(2, 0, 1),
+        env=visdom_env,
+        win="images",
+        opts={"title": "train_images"},
+    )
+
+    # Show the coarse and fine renders together with the ground truth images.
+    ims_full = torch.cat(
+        [
+            nerf_out[imvar][0].permute(2, 0, 1).detach().cpu().clamp(0.0, 1.0)
+            for imvar in ("rgb_coarse", "rgb_fine", "rgb_gt")
+        ],
+        dim=2,
+    )
+    viz.image(
+        ims_full,
+        env=visdom_env,
+        win="images_full",
+        opts={"title": "coarse | fine | target"},
+    )
+
+    # Make a 3D plot of training cameras and their emitted rays.
+    camera_trace = {
+        f"camera_{ci:03d}": o["camera"].cpu() for ci, o in enumerate(output_cache)
+    }
+    ray_pts_trace = {
+        f"ray_pts_{ci:03d}": Pointclouds(
+            ray_bundle_to_ray_points(o["coarse_ray_bundle"])
+            .detach()
+            .cpu()
+            .view(1, -1, 3)
+        )
+        for ci, o in enumerate(output_cache)
+    }
+    plotly_plot = plot_scene(
+        {
+            "training_scene": {
+                **camera_trace,
+                **ray_pts_trace,
+            },
+        },
+        pointcloud_max_points=5000,
+        pointcloud_marker_size=1,
+        camera_scale=0.3,
+    )
+    viz.plotlyplot(plotly_plot, env=visdom_env, win="scenes")