loading llff and blender datasets

Summary: Copy code from NeRF for loading LLFF data and blender synthetic data, and create dataset objects for them

Reviewed By: shapovalov

Differential Revision: D35581039

fbshipit-source-id: af7a6f3e9a42499700693381b5b147c991f57e5d

Author: bottler

LICENSE-3RD-PARTY

@@ -46,3 +46,26 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
+
+
+NeRF https://github.com/bmild/nerf/
+
+Copyright (c) 2020 bmild
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

projects/implicitron_trainer/README.md

@@ -5,7 +5,7 @@ Implicitron is a PyTorch3D-based framework for new-view synthesis via modeling t
 # License
 
 Implicitron is distributed as part of PyTorch3D under the [BSD license](https://github.com/facebookresearch/pytorch3d/blob/main/LICENSE).
-It includes code from [SRN](http://github.com/vsitzmann/scene-representation-networks) and [IDR](http://github.com/lioryariv/idr) repos.
+It includes code from the [NeRF](https://github.com/bmild/nerf), [SRN](http://github.com/vsitzmann/scene-representation-networks) and [IDR](http://github.com/lioryariv/idr) repos.
 See [LICENSE-3RD-PARTY](https://github.com/facebookresearch/pytorch3d/blob/main/LICENSE-3RD-PARTY) for their licenses.
 
 

projects/implicitron_trainer/experiment.py

@@ -315,7 +315,7 @@ def trainvalidate(
     epoch,
     loader,
     optimizer,
-    validation,
+    validation: bool,
     bp_var: str = "objective",
     metric_print_interval: int = 5,
     visualize_interval: int = 100,

projects/implicitron_trainer/tests/experiment.yaml

@@ -95,18 +95,18 @@ generic_model_args:
     append_coarse_samples_to_fine: true
     density_noise_std_train: 0.0
     return_weights: false
-    raymarcher_EmissionAbsorptionRaymarcher_args:
+    raymarcher_CumsumRaymarcher_args:
       surface_thickness: 1
       bg_color:
       - 0.0
-      background_opacity: 10000000000.0
+      background_opacity: 0.0
       density_relu: true
       blend_output: false
-    raymarcher_CumsumRaymarcher_args:
+    raymarcher_EmissionAbsorptionRaymarcher_args:
       surface_thickness: 1
       bg_color:
       - 0.0
-      background_opacity: 0.0
+      background_opacity: 10000000000.0
       density_relu: true
       blend_output: false
   renderer_SignedDistanceFunctionRenderer_args:
@@ -157,17 +157,12 @@ generic_model_args:
     view_sampler_args:
       masked_sampling: false
       sampling_mode: bilinear
-    feature_aggregator_IdentityFeatureAggregator_args:
-      exclude_target_view: true
-      exclude_target_view_mask_features: true
-      concatenate_output: true
-    feature_aggregator_ReductionFeatureAggregator_args:
+    feature_aggregator_AngleWeightedIdentityFeatureAggregator_args:
       exclude_target_view: true
       exclude_target_view_mask_features: true
       concatenate_output: true
-      reduction_functions:
-      - AVG
-      - STD
+      weight_by_ray_angle_gamma: 1.0
+      min_ray_angle_weight: 0.1
     feature_aggregator_AngleWeightedReductionFeatureAggregator_args:
       exclude_target_view: true
       exclude_target_view_mask_features: true
@@ -177,12 +172,17 @@ generic_model_args:
       - STD
       weight_by_ray_angle_gamma: 1.0
       min_ray_angle_weight: 0.1
-    feature_aggregator_AngleWeightedIdentityFeatureAggregator_args:
+    feature_aggregator_IdentityFeatureAggregator_args:
       exclude_target_view: true
       exclude_target_view_mask_features: true
       concatenate_output: true
-      weight_by_ray_angle_gamma: 1.0
-      min_ray_angle_weight: 0.1
+    feature_aggregator_ReductionFeatureAggregator_args:
+      exclude_target_view: true
+      exclude_target_view_mask_features: true
+      concatenate_output: true
+      reduction_functions:
+      - AVG
+      - STD
   implicit_function_IdrFeatureField_args:
     feature_vector_size: 3
     d_in: 3
@@ -203,42 +203,44 @@ generic_model_args:
     n_harmonic_functions_xyz: 0
     pooled_feature_dim: 0
     encoding_dim: 0
-  implicit_function_NeuralRadianceFieldImplicitFunction_args:
+  implicit_function_NeRFormerImplicitFunction_args:
     n_harmonic_functions_xyz: 10
     n_harmonic_functions_dir: 4
     n_hidden_neurons_dir: 128
     latent_dim: 0
     input_xyz: true
     xyz_ray_dir_in_camera_coords: false
     color_dim: 3
-    transformer_dim_down_factor: 1.0
-    n_hidden_neurons_xyz: 256
-    n_layers_xyz: 8
+    transformer_dim_down_factor: 2.0
+    n_hidden_neurons_xyz: 80
+    n_layers_xyz: 2
     append_xyz:
-    - 5
-  implicit_function_NeRFormerImplicitFunction_args:
+    - 1
+  implicit_function_NeuralRadianceFieldImplicitFunction_args:
     n_harmonic_functions_xyz: 10
     n_harmonic_functions_dir: 4
     n_hidden_neurons_dir: 128
     latent_dim: 0
     input_xyz: true
     xyz_ray_dir_in_camera_coords: false
     color_dim: 3
-    transformer_dim_down_factor: 2.0
-    n_hidden_neurons_xyz: 80
-    n_layers_xyz: 2
+    transformer_dim_down_factor: 1.0
+    n_hidden_neurons_xyz: 256
+    n_layers_xyz: 8
     append_xyz:
-    - 1
-  implicit_function_SRNImplicitFunction_args:
-    raymarch_function_args:
+    - 5
+  implicit_function_SRNHyperNetImplicitFunction_args:
+    hypernet_args:
       n_harmonic_functions: 3
       n_hidden_units: 256
       n_layers: 2
+      n_hidden_units_hypernet: 256
+      n_layers_hypernet: 1
       in_features: 3
       out_features: 256
+      latent_dim_hypernet: 0
       latent_dim: 0
       xyz_in_camera_coords: false
-      raymarch_function: null
     pixel_generator_args:
       n_harmonic_functions: 4
       n_hidden_units: 256
@@ -247,18 +249,16 @@ generic_model_args:
       in_features: 256
       out_features: 3
       ray_dir_in_camera_coords: false
-  implicit_function_SRNHyperNetImplicitFunction_args:
-    hypernet_args:
+  implicit_function_SRNImplicitFunction_args:
+    raymarch_function_args:
       n_harmonic_functions: 3
       n_hidden_units: 256
       n_layers: 2
-      n_hidden_units_hypernet: 256
-      n_layers_hypernet: 1
       in_features: 3
       out_features: 256
-      latent_dim_hypernet: 0
       latent_dim: 0
       xyz_in_camera_coords: false
+      raymarch_function: null
     pixel_generator_args:
       n_harmonic_functions: 4
       n_hidden_units: 256
@@ -282,6 +282,13 @@ solver_args:
 data_source_args:
   dataset_map_provider_class_type: ???
   data_loader_map_provider_class_type: SequenceDataLoaderMapProvider
+  dataset_map_provider_BlenderDatasetMapProvider_args:
+    base_dir: ???
+    object_name: ???
+    path_manager_factory_class_type: PathManagerFactory
+    n_known_frames_for_test: null
+    path_manager_factory_PathManagerFactory_args:
+      silence_logs: true
   dataset_map_provider_JsonIndexDatasetMapProvider_args:
     category: ???
     task_str: singlesequence
@@ -317,6 +324,13 @@ data_source_args:
       sort_frames: false
     path_manager_factory_PathManagerFactory_args:
       silence_logs: true
+  dataset_map_provider_LlffDatasetMapProvider_args:
+    base_dir: ???
+    object_name: ???
+    path_manager_factory_class_type: PathManagerFactory
+    n_known_frames_for_test: null
+    path_manager_factory_PathManagerFactory_args:
+      silence_logs: true
   data_loader_map_provider_SequenceDataLoaderMapProvider_args:
     batch_size: 1
     num_workers: 0

pytorch3d/implicitron/dataset/blender_dataset_map_provider.py

@@ -0,0 +1,52 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch
+from pytorch3d.implicitron.tools.config import registry
+
+from .load_blender import load_blender_data
+from .single_sequence_dataset import (
+    _interpret_blender_cameras,
+    SingleSceneDatasetMapProviderBase,
+)
+
+
+@registry.register
+class BlenderDatasetMapProvider(SingleSceneDatasetMapProviderBase):
+    """
+    Provides data for one scene from Blender synthetic dataset.
+    Uses the code in load_blender.py
+
+    Members:
+        base_dir: directory holding the data for the scene.
+        object_name: The name of the scene (e.g. "lego"). This is just used as a label.
+            It will typically be equal to the name of the directory self.base_dir.
+        path_manager_factory: Creates path manager which may be used for
+            interpreting paths.
+        n_known_frames_for_test: If set, training frames are included in the val
+            and test datasets, and this many random training frames are added to
+            each test batch. If not set, test batches each contain just a single
+            testing frame.
+    """
+
+    def _load_data(self) -> None:
+        path_manager = self.path_manager_factory.get()
+        images, poses, _, hwf, i_split = load_blender_data(
+            self.base_dir,
+            testskip=1,
+            path_manager=path_manager,
+        )
+        H, W, focal = hwf
+        H, W = int(H), int(W)
+        images = torch.from_numpy(images)
+
+        # pyre-ignore[16]
+        self.poses = _interpret_blender_cameras(poses, H, W, focal)
+        # pyre-ignore[16]
+        self.images = images
+        # pyre-ignore[16]
+        self.i_split = i_split

pytorch3d/implicitron/dataset/data_source.py

@@ -8,9 +8,11 @@
 
 from pytorch3d.implicitron.tools.config import ReplaceableBase, run_auto_creation
 
-from . import json_index_dataset_map_provider  # noqa
+from .blender_dataset_map_provider import BlenderDatasetMapProvider  # noqa
 from .data_loader_map_provider import DataLoaderMap, DataLoaderMapProviderBase
 from .dataset_map_provider import DatasetMap, DatasetMapProviderBase, Task
+from .json_index_dataset_map_provider import JsonIndexDatasetMapProvider  # noqa
+from .llff_dataset_map_provider import LlffDatasetMapProvider  # noqa
 
 
 class DataSourceBase(ReplaceableBase):

pytorch3d/implicitron/dataset/dataset_base.py

@@ -36,10 +36,11 @@ class FrameData(Mapping[str, Any]):
     Args:
         frame_number: The number of the frame within its sequence.
             0-based continuous integers.
-        frame_timestamp: The time elapsed since the start of a sequence in sec.
         sequence_name: The unique name of the frame's sequence.
         sequence_category: The object category of the sequence.
-        image_size_hw: The size of the image in pixels; (height, width) tuple.
+        frame_timestamp: The time elapsed since the start of a sequence in sec.
+        image_size_hw: The size of the image in pixels; (height, width) tensor
+                        of shape (2,).
         image_path: The qualified path to the loaded image (with dataset_root).
         image_rgb: A Tensor of shape `(3, H, W)` holding the RGB image
             of the frame; elements are floats in [0, 1].
@@ -81,9 +82,9 @@ class FrameData(Mapping[str, Any]):
     """
 
     frame_number: Optional[torch.LongTensor]
-    frame_timestamp: Optional[torch.Tensor]
     sequence_name: Union[str, List[str]]
     sequence_category: Union[str, List[str]]
+    frame_timestamp: Optional[torch.Tensor] = None
     image_size_hw: Optional[torch.Tensor] = None
     image_path: Union[str, List[str], None] = None
     image_rgb: Optional[torch.Tensor] = None
@@ -101,7 +102,7 @@ class FrameData(Mapping[str, Any]):
     sequence_point_cloud_path: Union[str, List[str], None] = None
     sequence_point_cloud: Optional[Pointclouds] = None
     sequence_point_cloud_idx: Optional[torch.Tensor] = None
-    frame_type: Union[str, List[str], None] = None  # seen | unseen
+    frame_type: Union[str, List[str], None] = None  # known | unseen
     meta: dict = field(default_factory=lambda: {})
 
     def to(self, *args, **kwargs):

pytorch3d/implicitron/dataset/llff_dataset_map_provider.py

@@ -0,0 +1,61 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import numpy as np
+import torch
+from pytorch3d.implicitron.tools.config import registry
+
+from .load_llff import load_llff_data
+
+from .single_sequence_dataset import (
+    _interpret_blender_cameras,
+    SingleSceneDatasetMapProviderBase,
+)
+
+
+@registry.register
+class LlffDatasetMapProvider(SingleSceneDatasetMapProviderBase):
+    """
+    Provides data for one scene from the LLFF dataset.
+
+    Members:
+        base_dir: directory holding the data for the scene.
+        object_name: The name of the scene (e.g. "fern"). This is just used as a label.
+            It will typically be equal to the name of the directory self.base_dir.
+        path_manager_factory: Creates path manager which may be used for
+            interpreting paths.
+        n_known_frames_for_test: If set, training frames are included in the val
+            and test datasets, and this many random training frames are added to
+            each test batch. If not set, test batches each contain just a single
+            testing frame.
+    """
+
+    def _load_data(self) -> None:
+        path_manager = self.path_manager_factory.get()
+        images, poses, _ = load_llff_data(
+            self.base_dir, factor=8, path_manager=path_manager
+        )
+        hwf = poses[0, :3, -1]
+        poses = poses[:, :3, :4]
+
+        i_test = np.arange(images.shape[0])[::8]
+        i_test_index = set(i_test.tolist())
+        i_train = np.array(
+            [i for i in np.arange(images.shape[0]) if i not in i_test_index]
+        )
+        i_split = (i_train, i_test, i_test)
+        H, W, focal = hwf
+        H, W = int(H), int(W)
+        images = torch.from_numpy(images)
+        poses = torch.from_numpy(poses)
+
+        # pyre-ignore[16]
+        self.poses = _interpret_blender_cameras(poses, H, W, focal)
+        # pyre-ignore[16]
+        self.images = images
+        # pyre-ignore[16]
+        self.i_split = i_split