From de2cdb6ccadad483d93e9270da21d28a3f66fab4 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Sun, 11 Sep 2022 12:52:27 -0700 Subject: [PATCH] Friends dataset loading (#389) command line dataparsing and friends dataset loading improvement --- .vscode/launch.json | 46 ++++++++++--- nerfactory/configs/base.py | 21 ++++-- .../dataparsers/friends_parser.py | 64 +++++++++---------- 3 files changed, 85 insertions(+), 46 deletions(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index e8a656fbd..5cbcd1a53 100644 --- a/.vscode/launch.json +++ b/.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" + ] } ] -} +} \ No newline at end of file diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index a34dec8d8..cc0dd84b6 100644 --- a/nerfactory/configs/base.py +++ b/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""" diff --git a/nerfactory/datamanagers/dataparsers/friends_parser.py b/nerfactory/datamanagers/dataparsers/friends_parser.py index c5ff3be1d..bbb41187f 100644 --- a/nerfactory/datamanagers/dataparsers/friends_parser.py +++ b/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,