Friends dataset loading (#389)

command line dataparsing and friends dataset loading improvement

.vscode/launch.json

@@ -25,7 +25,9 @@
       "request": "launch",
       "program": "scripts/run_train.py",
       "console": "integratedTerminal",
-      "args": ["vanilla_nerf"]
+      "args": [
+        "vanilla_nerf"
+      ]
     },
     {
       "name": "Python: run_train (cpu)",
@@ -33,7 +35,9 @@
       "request": "launch",
       "program": "scripts/run_train.py",
       "console": "integratedTerminal",
-      "args": ["machine.num_gpus=0"]
+      "args": [
+        "machine.num_gpus=0"
+      ]
     },
     {
       "name": "Python: run_train ngp (gpu)",
@@ -41,7 +45,28 @@
       "request": "launch",
       "program": "scripts/run_train.py",
       "console": "integratedTerminal",
-      "args": ["instant_ngp"]
+      "args": [
+        "instant_ngp"
+      ]
+    },
+    {
+      "name": "Python: run_train ngp-friends (gpu)",
+      "type": "python",
+      "request": "launch",
+      "program": "scripts/run_train.py",
+      "console": "integratedTerminal",
+      "args": [
+        "instant_ngp",
+        "--viewer.enable",
+        "--viewer.zmq-url=tcp://127.0.0.1:6001",
+        "--viewer.websocket-port=8014",
+        "--trainer.steps_per_save=1000000",
+        "--trainer.steps_per_test=1000000",
+        "--pipeline.model.density-field-params.base-scale=4",
+        "--pipeline.model.density-field-params.num-cascades=1",
+        "pipeline.datamanager.train-dataparser:friends-data-parser-config",
+        "--pipeline.datamanager.train-dataparser.scene-scale=4"
+      ]
     },
     {
       "name": "Python: run_train mipnerf 360 (gpu)",
@@ -49,7 +74,9 @@
       "request": "launch",
       "program": "scripts/run_train.py",
       "console": "integratedTerminal",
-      "args": ["mipnerf_360"]
+      "args": [
+        "mipnerf_360"
+      ]
     },
     {
       "name": "Python: run_train tensorf (gpu)",
@@ -57,7 +84,9 @@
       "request": "launch",
       "program": "scripts/run_train.py",
       "console": "integratedTerminal",
-      "args": ["tensorf"]
+      "args": [
+        "tensorf"
+      ]
     },
     {
       "name": "Python: run_train nerfw (gpu)",
@@ -67,7 +96,6 @@
       "console": "integratedTerminal",
       "args": [
         "nerfw"
-        // "++data.use_preprocessing_cache=true"
       ]
     },
     {
@@ -76,7 +104,9 @@
       "request": "launch",
       "program": "scripts/run_train.py",
       "console": "integratedTerminal",
-      "args": ["semantic_nerf", "++data.use_preprocessing_cache=true"]
+      "args": [
+        "semantic_nerf"
+      ]
     }
   ]
-}
+}

nerfactory/configs/base.py

@@ -19,7 +19,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Dict, List, Literal, Optional, Tuple, Type
+from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union
 
 import dcargs
 import torch
@@ -250,6 +250,12 @@ class FriendsDataParserConfig(DataParserConfig):
     """directory specifying location of data"""
     include_semantics: bool = True
     """whether or not to include loading of semantics data"""
+    downscale_factor: int = 8
+    scene_scale: float = 4.0
+    """
+    Sets the bounding cube to have edge length of this size.
+    The longest dimension of the Friends axis-aligned bbox will be scaled to this value.
+    """
 
 
 @dataclass
@@ -266,7 +272,7 @@ class MipNerf360DataParserConfig(DataParserConfig):
     """1/val_skip images to use for validation. Defaults to 8."""
     auto_scale: bool = True
     """Scale based on pose bounds. Defaults to True."""
-    aabb_scale = 4
+    aabb_scale: float = 4
     """Scene scale, Defaults to 1.0."""
 
 
@@ -294,7 +300,7 @@ class Record3DDataParserConfig(DataParserConfig):
     """Location of data"""
     val_skip: int = 8
     """1/val_skip images to use for validation. Defaults to 8."""
-    aabb_scale = 4.0
+    aabb_scale: float = 4.0
     """Scene scale, Defaults to 4.0."""
     max_dataset_size: int = 150
     """Max number of images to train on. If the dataset has
@@ -314,7 +320,14 @@ class VanillaDataManagerConfig(InstantiateConfig):
 
     _target: Type = VanillaDataManager
     """target class to instantiate"""
-    train_dataparser: DataParserConfig = BlenderDataParserConfig()
+    train_dataparser: Union[
+        NerfactoryDataParserConfig,
+        BlenderDataParserConfig,
+        FriendsDataParserConfig,
+        MipNerf360DataParserConfig,
+        InstantNGPDataParserConfig,
+        Record3DDataParserConfig,
+    ] = BlenderDataParserConfig()
     """specifies the dataparser used to unpack the data"""
     train_num_rays_per_batch: int = 1024
     """number of rays per batch to use per training iteration"""

nerfactory/datamanagers/dataparsers/friends_parser.py

@@ -58,32 +58,6 @@ class Friends(DataParser):
 
     config: cfg.FriendsDataParserConfig
 
-    @classmethod
-    def _get_aabb_and_transform(cls, basedir):
-        """Returns the aabb and pointcloud transform from the threejs.json file.
-
-        Args:
-            basedir: base directory to load from
-        """
-        filename = basedir / "threejs.json"
-        assert filename.exists()
-        data = load_from_json(filename)
-
-        # point cloud transformation
-        transposed_point_cloud_transform = np.array(data["object"]["children"][0]["matrix"]).reshape(4, 4).T
-        assert transposed_point_cloud_transform[3, 3] == 1.0
-
-        # bbox transformation
-        bbox_transform = np.array(data["object"]["children"][1]["matrix"]).reshape(4, 4).T
-        w, h, d = data["geometries"][1]["width"], data["geometries"][1]["height"], data["geometries"][1]["depth"]
-        temp = np.array([w, h, d]) / 2.0
-        bbox = np.array([-temp, temp])
-        bbox = np.concatenate([bbox, np.ones_like(bbox[:, 0:1])], axis=1)
-        bbox = (bbox_transform @ bbox.T).T[:, 0:3]
-
-        aabb = bbox  # rename to aabb because it's an axis-aligned bounding box
-        return torch.from_numpy(aabb).float(), torch.from_numpy(transposed_point_cloud_transform).float()
-
     def _generate_dataset_inputs(self, split="train"):  # pylint: disable=unused-argument
 
         abs_dir = get_absolute_path(self.config.data_directory)
@@ -92,6 +66,10 @@ class Friends(DataParser):
         frames = cameras_json["frames"]
         bbox = torch.tensor(cameras_json["bbox"])
 
+        downscale_suffix = f"_{self.config.downscale_factor}" if self.config.downscale_factor != 1 else ""
+        images_folder = f"images{downscale_suffix}"
+        segmentations_folder = f"segmentations{downscale_suffix}"
+
         image_filenames = []
         fx = []
         fy = []
@@ -100,7 +78,7 @@ class Friends(DataParser):
         camera_to_worlds = []
         for frame in frames:
             # unpack data
-            image_filename = abs_dir / "images" / frame["image_name"]
+            image_filename = abs_dir / images_folder / frame["image_name"]
             intrinsics = torch.tensor(frame["intrinsics"])
             camtoworld = torch.tensor(frame["camtoworld"])[:3]
             # append data
@@ -121,22 +99,39 @@ class Friends(DataParser):
         camera_to_worlds[:, :3] = rotation @ camera_to_worlds[:, :3]
         bbox = (rotation @ bbox.T).T
 
-        # -- set the bounding box ---
+        scene_scale = self.config.scene_scale
+
+        # -- set the scene bounds ---
         scene_bounds = SceneBounds(aabb=bbox)
-        # # for shifting and rescale accoding to scene bounds
-        # box_center = scene_bounds_original.get_center()
-        # box_scale_factor = 5.0 / scene_bounds_original.get_diagonal_length()  # the target diagonal length
-        # scene_bounds = scene_bounds_original.get_centered_and_scaled_scene_bounds(box_scale_factor)
+        # center the box and adjust the cameras too
+        center = scene_bounds.get_center()
+        scene_bounds.aabb -= center
+        camera_to_worlds[..., 3] -= center
+        # scale the longest dimension to match the cube size
+        lengths = scene_bounds.aabb[1] - scene_bounds.aabb[0]
+        longest_dim = torch.argmax(lengths)
+        longest_length = lengths[longest_dim]
+        scale = scene_scale / longest_length
+        scene_bounds.aabb = scene_bounds.aabb * scale  # box
+        camera_to_worlds[..., 3] *= scale  # cameras
 
         # --- semantics ---
         semantics = None
         if self.config.include_semantics:
             thing_filenames = [
-                Path(str(image_filename).replace("/images/", "/segmentations/thing/").replace(".jpg", ".png"))
+                Path(
+                    str(image_filename)
+                    .replace(f"/{images_folder}/", f"/{segmentations_folder}/thing/")
+                    .replace(".jpg", ".png")
+                )
                 for image_filename in image_filenames
             ]
             stuff_filenames = [
-                Path(str(image_filename).replace("/images/", "/segmentations/stuff/").replace(".jpg", ".png"))
+                Path(
+                    str(image_filename)
+                    .replace(f"/{images_folder}/", f"/{segmentations_folder}/stuff/")
+                    .replace(".jpg", ".png")
+                )
                 for image_filename in image_filenames
             ]
             panoptic_classes = load_from_json(abs_dir / "panoptic_classes.json")
@@ -164,6 +159,7 @@ class Friends(DataParser):
             camera_to_worlds=camera_to_worlds,
             camera_type=CameraType.PERSPECTIVE,
         )
+        cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor)
 
         dataset_inputs = DatasetInputs(
             image_filenames=image_filenames,