From 56113e6c750d680c0de2d09508832475f7a50c99 Mon Sep 17 00:00:00 2001 From: evonneng Date: Tue, 6 Sep 2022 13:54:35 -0700 Subject: [PATCH] [WIP] dataloading cleanup (#332) * testing; need to merge * fixing wrong base file in models/base * pushed input parser into datasets class to abstract away fetching data in dataloader * slight name chage and all * nits * added dataparser * added licnese --- docs/conf.py | 2 +- docs/reference/api/data.rst | 16 +- .../pipelines/creating_pipelines.ipynb | 8 +- nerfactory/configs/base.py | 32 +- nerfactory/configs/base_configs.py | 24 +- nerfactory/dataloaders/data_parsers.py | 740 ------------------ .../{dataloaders => datamanagers}/__init__.py | 0 .../{dataloaders => datamanagers}/base.py | 94 +-- .../dataloaders.py} | 35 +- nerfactory/datamanagers/dataparsers/base.py | 106 +++ .../dataparsers/blender_parser.py | 95 +++ .../dataparsers/friends_parser.py | 183 +++++ .../dataparsers/instant_ngp_parser.py | 106 +++ .../dataparsers/mipnerf_parser.py | 151 ++++ .../dataparsers/nerfactory_parser.py | 112 +++ .../dataparsers/record3d_parser.py | 148 ++++ .../datasets.py} | 67 +- .../pixel_sampler.py | 0 .../{dataloaders => datamanagers}/structs.py | 0 nerfactory/engine/trainer.py | 10 +- nerfactory/fields/instant_ngp_field.py | 2 +- nerfactory/models/base.py | 2 +- nerfactory/models/modules/scene_colliders.py | 2 +- nerfactory/models/semantic_nerf.py | 2 +- nerfactory/pipelines/base.py | 30 +- nerfactory/viewer/server/viewer_utils.py | 16 +- scripts/benchmarking/launch_train_blender.sh | 2 +- scripts/run_eval.py | 10 +- tests/test_run_train.py | 13 +- 29 files changed, 1048 insertions(+), 960 deletions(-) delete mode 100644 nerfactory/dataloaders/data_parsers.py rename nerfactory/{dataloaders => datamanagers}/__init__.py (100%) rename nerfactory/{dataloaders => datamanagers}/base.py (74%) rename nerfactory/{dataloaders/image_dataloader.py => datamanagers/dataloaders.py} (89%) create mode 100644 nerfactory/datamanagers/dataparsers/base.py create mode 100644 nerfactory/datamanagers/dataparsers/blender_parser.py create mode 100644 nerfactory/datamanagers/dataparsers/friends_parser.py create mode 100644 nerfactory/datamanagers/dataparsers/instant_ngp_parser.py create mode 100644 nerfactory/datamanagers/dataparsers/mipnerf_parser.py create mode 100644 nerfactory/datamanagers/dataparsers/nerfactory_parser.py create mode 100644 nerfactory/datamanagers/dataparsers/record3d_parser.py rename nerfactory/{dataloaders/image_dataset.py => datamanagers/datasets.py} (64%) rename nerfactory/{dataloaders => datamanagers}/pixel_sampler.py (100%) rename nerfactory/{dataloaders => datamanagers}/structs.py (100%) diff --git a/docs/conf.py b/docs/conf.py index 177b8b410..d987c8e2c 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -16,7 +16,7 @@ import sys sys.path.insert(0, os.path.abspath("..")) import nerfactory.configs.base -import nerfactory.dataloaders.base +import nerfactory.datamanagers.base # -- Project information ----------------------------------------------------- diff --git a/docs/reference/api/data.rst b/docs/reference/api/data.rst index b3759b89c..fff8e624d 100644 --- a/docs/reference/api/data.rst +++ b/docs/reference/api/data.rst @@ -6,34 +6,34 @@ Data Data Parsers ---------------- -.. automodule:: nerfactory.dataloaders.data_parsers +.. automodule:: nerfactory.datamanagers.dataparsers :members: :show-inheritance: -Image Dataset +Datasets ---------------- -.. automodule:: nerfactory.dataloaders.image_dataset +.. automodule:: nerfactory.datamanagers.datasets :members: :show-inheritance: -Image Dataloader +Dataloaders ---------------- -.. automodule:: nerfactory.dataloaders.image_dataloader +.. automodule:: nerfactory.datamanagers.dataloaders :members: :show-inheritance: -Pixel Sampler +Pixel Samplers ---------------- -.. automodule:: nerfactory.dataloaders.pixel_sampler +.. automodule:: nerfactory.datamanagers.pixel_sampler :members: :show-inheritance: Structures ---------------- -.. automodule:: nerfactory.dataloaders.structs +.. automodule:: nerfactory.datamanagers.structs :members: :show-inheritance: \ No newline at end of file diff --git a/docs/tutorials/pipelines/creating_pipelines.ipynb b/docs/tutorials/pipelines/creating_pipelines.ipynb index b9fad4c24..294bcb953 100644 --- a/docs/tutorials/pipelines/creating_pipelines.ipynb +++ b/docs/tutorials/pipelines/creating_pipelines.ipynb @@ -72,14 +72,12 @@ "DataManager Components\n", "
\n", "Datasets Parsers:
\n", - " Takes a data location and returns a DatasetInputs instance. A DatasetInputs will contain all data needed throughout the pipeline.\n", + " Takes a data location and returns a DatsetInputs instance. A DatsetInputs will contain all data needed throughout the pipeline.\n", "
  • Vanilla Image DataParser
    • \n", - "
  • Panoptic Image DataParser
    • \n", "
    \n", "Image Datasets:
    \n", " Takes image filename information and can return images.\n", "
  • Vanilla Image Dataset
    • \n", - "
  • Panoptic Image Dataset
    • \n", "
    \n", "Image Dataloader:
    \n", " Takes an Image Dataset and samples images from it.\n", @@ -119,8 +117,8 @@ "import torch\n", "\n", "from nerfactory.utils.config import DataManagerConfig, ModelConfig, PipelineConfig, ViewerConfig\n", - "from nerfactory.dataloaders.base import DataManager\n", - "from nerfactory.dataloaders.structs import DatasetInputs\n", + "from nerfactory.datamanagers.base import DataManager\n", + "from nerfactory.datamanagers.structs import DatasetInputs\n", "from nerfactory.models.base import Model\n", "from nerfactory.pipelines.base import Pipeline\n", "from nerfactory.models.instant_ngp import NGPModel\n", diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 11a4ae65a..93677be81 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -24,16 +24,18 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Type import torch from nerfactory.configs.utils import to_immutable_dict -from nerfactory.dataloaders.base import VanillaDataManager -from nerfactory.dataloaders.data_parsers import ( - Blender, - DataParser, - Friends, - InstantNGP, - Mipnerf360, - Nerfactory, - Record3D, -) + +# data instances +from nerfactory.datamanagers.base import VanillaDataManager +from nerfactory.datamanagers.dataparsers.base import DataParser +from nerfactory.datamanagers.dataparsers.blender_parser import Blender +from nerfactory.datamanagers.dataparsers.friends_parser import Friends +from nerfactory.datamanagers.dataparsers.instant_ngp_parser import InstantNGP +from nerfactory.datamanagers.dataparsers.mipnerf_parser import Mipnerf360 +from nerfactory.datamanagers.dataparsers.nerfactory_parser import Nerfactory +from nerfactory.datamanagers.dataparsers.record3d_parser import Record3D + +# model instances from nerfactory.models.base import Model from nerfactory.models.instant_ngp import NGPModel from nerfactory.models.nerfw import NerfWModel @@ -238,11 +240,10 @@ class VanillaDataManagerConfig(InstantiateConfig): """Configuration for data manager instantiation""" _target: Type = VanillaDataManager - train_dataset: DataParserConfig = BlenderDataParserConfig() - image_dataset_type: str = "rgb" + train_dataparser: DataParserConfig = BlenderDataParserConfig() train_num_rays_per_batch: int = 1024 train_num_images_to_sample_from: int = -1 - eval_dataset: Optional[InstantiateConfig] = None + eval_dataparser: Optional[InstantiateConfig] = None eval_image_indices: Optional[Tuple[int, ...]] = (0,) eval_num_rays_per_chunk: int = 4096 @@ -252,8 +253,7 @@ class FriendsDataManagerConfig(VanillaDataManagerConfig): """Friends data manager config""" _target: Type = VanillaDataManager - train_dataset: DataParserConfig = FriendsDataParserConfig() - image_dataset_type: str = "panoptic" + train_dataparser: DataParserConfig = FriendsDataParserConfig() # Model related configs @@ -322,7 +322,7 @@ class PipelineConfig(InstantiateConfig): """Configuration for pipeline instantiation""" _target: Type = Pipeline - data_manager: VanillaDataManagerConfig = VanillaDataManagerConfig() + datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig() model: ModelConfig = ModelConfig() diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index ea6eff6a3..a708c1d5c 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -43,8 +43,8 @@ base_configs["instant_ngp"] = Config( method_name="instant_ngp", trainer=TrainerConfig(mixed_precision=True), pipeline=PipelineConfig( - data_manager=VanillaDataManagerConfig( - train_dataset=BlenderDataParserConfig(), + datamanager=VanillaDataManagerConfig( + train_dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192, eval_num_rays_per_chunk=8192, ), @@ -63,8 +63,8 @@ base_configs["mipnerf_360"] = Config( method_name="mipnerf_360", trainer=TrainerConfig(steps_per_test=200), pipeline=PipelineConfig( - data_manager=VanillaDataManagerConfig( - train_dataset=MipNerf360DataParserConfig(), + datamanager=VanillaDataManagerConfig( + train_dataparser=MipNerf360DataParserConfig(), train_num_rays_per_batch=8192, eval_num_rays_per_chunk=8192, ), @@ -81,8 +81,8 @@ base_configs["mipnerf_360"] = Config( base_configs["mipnerf"] = Config( method_name="mipnerf", pipeline=PipelineConfig( - data_manager=VanillaDataManagerConfig( - train_dataset=BlenderDataParserConfig(), + datamanager=VanillaDataManagerConfig( + train_dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192, eval_num_rays_per_chunk=8192, ), @@ -98,14 +98,14 @@ base_configs["mipnerf"] = Config( base_configs["nerfw"] = Config( experiment_name="friends_TBBT-big_living_room", method_name="nerfw", - pipeline=PipelineConfig(data_manager=FriendsDataManagerConfig(), model=NerfWModelConfig()), + pipeline=PipelineConfig(datamanager=FriendsDataManagerConfig(), model=NerfWModelConfig()), ) base_configs["semantic_nerf"] = Config( experiment_name="friends_TBBT-big_living_room", method_name="semantic_nerf", pipeline=PipelineConfig( - data_manager=FriendsDataManagerConfig(), + datamanager=FriendsDataManagerConfig(), model=ModelConfig( _target=SemanticNerfModel, loss_coefficients={"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0, "semantic_loss_fine": 0.05}, @@ -118,8 +118,8 @@ base_configs["semantic_nerf"] = Config( base_configs["vanilla_nerf"] = Config( method_name="vanilla_nerf", pipeline=PipelineConfig( - data_manager=VanillaDataManagerConfig( - train_dataset=BlenderDataParserConfig(), + datamanager=VanillaDataManagerConfig( + train_dataparser=BlenderDataParserConfig(), ), model=ModelConfig(_target=NeRFModel), ), @@ -129,8 +129,8 @@ base_configs["tensorf"] = Config( method_name="tensorf", trainer=TrainerConfig(mixed_precision=True), pipeline=PipelineConfig( - data_manager=VanillaDataManagerConfig( - train_dataset=BlenderDataParserConfig(), + datamanager=VanillaDataManagerConfig( + train_dataparser=BlenderDataParserConfig(), ), model=TensoRFModelConfig(), ), diff --git a/nerfactory/dataloaders/data_parsers.py b/nerfactory/dataloaders/data_parsers.py deleted file mode 100644 index f44fdccfb..000000000 --- a/nerfactory/dataloaders/data_parsers.py +++ /dev/null @@ -1,740 +0,0 @@ -# Copyright 2022 The Plenoptix Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""A set of standard datasets.""" - -from __future__ import annotations - -import dataclasses -import logging -from abc import abstractmethod -from dataclasses import dataclass -from pathlib import Path - -import imageio -import numpy as np -import torch -from scipy.spatial.transform import Rotation - -from nerfactory.cameras import utils as camera_utils -from nerfactory.cameras.cameras import Cameras, CameraType -from nerfactory.configs import base as cfg -from nerfactory.dataloaders.structs import ( - DatasetInputs, - PointCloud, - SceneBounds, - Semantics, -) -from nerfactory.utils.colmap_utils import ( - read_cameras_binary, - read_images_binary, - read_points3d_binary, -) -from nerfactory.utils.colors import get_color -from nerfactory.utils.io import ( - get_absolute_path, - load_from_json, - load_from_pkl, - write_to_pkl, -) -from nerfactory.utils.misc import get_hash_str_from_dict - - -@dataclass -class DataParser: - """A dataset.""" - - def __init__(self, config: cfg.DataParserConfig): - super().__init__() - self.config = config - - @abstractmethod - def _generate_dataset_inputs(self, split: str = "train") -> DatasetInputs: - """Returns the dataset inputs for the given split. - - Args: - split: Which dataset split to generate. - - Returns: - DatasetInputs - """ - - def get_dataset_inputs(self, split: str = "train", use_preprocessing_cache: bool = False) -> DatasetInputs: - """Returns the dataset inputs for the given split. - - Args: - split: Which dataset split to generate. - use_preprocessing_cache: Whether to use the cached dataset inputs. Defaults to False. - - Returns: - DatasetInputs - """ - if use_preprocessing_cache: - dataset_inputs = self._get_dataset_inputs_from_cache(split) - else: - dataset_inputs = self._generate_dataset_inputs(split) - return dataset_inputs - - def _get_cache_filename(self, split: str) -> Path: - """Creates a cache filename from the dataset inputs arguments. - - Args: - split: Which dataset split to generate filename for. - - Returns: - filename for cache. - """ - dataset_config_hash = get_hash_str_from_dict(dataclasses.asdict(self)) - dataset_config_hash_filename = get_absolute_path( - f"cache/dataset_inputs/{dataset_config_hash}-{split}.pkl" - ).parent.mkdir() - assert dataset_config_hash_filename is not None, "dataset hash is None" - return dataset_config_hash_filename - - def save_dataset_inputs_to_cache(self, split: str): - """Saves the dataset inputs to cache. - - Args: - split: Which dataset split to save. - """ - dataset_inputs = self.get_dataset_inputs(split=split) - dataset_config_hash_filename = self._get_cache_filename(split) - write_to_pkl(dataset_config_hash_filename, dataset_inputs) - - def _get_dataset_inputs_from_cache(self, split: str) -> DatasetInputs: - """Loads the dataset inputs from cache. If the cache does not exist, it will be created. - - Args: - split: Which dataset split to load. - """ - dataset_config_hash_filename = self._get_cache_filename(split) - if dataset_config_hash_filename.exists(): - logging.info("Loading dataset from cache.") - dataset_inputs = load_from_pkl(dataset_config_hash_filename) - else: - logging.info("Cache file not found. Generating and saving dataset to cache.") - dataset_inputs = self._generate_dataset_inputs(split=split) - self.save_dataset_inputs_to_cache(split) - return dataset_inputs - - -@dataclass -class Nerfactory(DataParser): - """Nerfactory Dataset - - Args: - data_directory: Location of data - scale_factor: How much to scale the camera origins by. - downscale_factor: How much to downscale images. - scene_scale: How much to scale the scene. - orientation_method: The method to use for orientation. Either "pca" or "up". - """ - - config: cfg.NerfactoryDataParserConfig - - def _generate_dataset_inputs(self, split="train"): - - abs_dir = get_absolute_path(self.config.data_directory) - - meta = load_from_json(abs_dir / "transforms.json") - image_filenames = [] - poses = [] - num_skipped_image_filenames = 0 - for frame in meta["frames"]: - fname = abs_dir / Path(frame["file_path"]) - if self.config.downscale_factor > 1: - fname = abs_dir / f"images_{self.config.downscale_factor}" / Path(frame["file_path"]).name - else: - fname = abs_dir / Path(frame["file_path"]) - if not fname: - num_skipped_image_filenames += 1 - else: - image_filenames.append(fname) - poses.append(np.array(frame["transform_matrix"])) - if num_skipped_image_filenames >= 0: - logging.info("Skipping %s files in dataset split %s.", num_skipped_image_filenames, split) - assert ( - len(image_filenames) != 0 - ), """ - No image files found. - You should check the file_paths in the transforms.json file to make sure they are correct. - """ - poses = torch.from_numpy(np.array(poses).astype(np.float32)) - poses = camera_utils.auto_orient_poses(poses, method=self.config.orientation_method) - - # Scale poses - scale_factor = 1.0 / torch.max(torch.abs(poses[:, :3, 3])) - poses[:, :3, 3] *= scale_factor - - distortion_params = camera_utils.get_distortion_params( - k1=float(meta["k1"]), k2=float(meta["k2"]), p1=float(meta["p1"]), p2=float(meta["p2"]) - ) - - # in x,y,z order - # assumes that the scene is centered at the origin - aabb_scale = self.config.scene_scale - scene_bounds = SceneBounds( - aabb=torch.tensor( - [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 - ) - ) - - cameras = Cameras( - fx=float(meta["fl_x"]), - fy=float(meta["fl_y"]), - cx=float(meta["cx"]), - cy=float(meta["cy"]), - distortion_params=distortion_params, - height=int(meta["h"]), - width=int(meta["w"]), - camera_to_worlds=poses[:, :3, :4], - camera_type=CameraType.PERSPECTIVE, - ) - - cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor) - - dataset_inputs = DatasetInputs( - image_filenames=image_filenames, - cameras=cameras, - scene_bounds=scene_bounds, - ) - return dataset_inputs - - -@dataclass -class Blender(DataParser): - """Blender Dataset - Some of this code comes from https://github.com/yenchenlin/nerf-pytorch/blob/master/load_blender.py#L37. - - Args: - data_directory: Location of data - alpha_color: Sets transparent regions to specified color, otherwise black. - scale_factor: How much to scale the camera origins by. - """ - - def __init__(self, config: cfg.BlenderDataParserConfig): - super().__init__(config=config) - self.data_directory: Path = config.data_directory - self.scale_factor: float = config.scale_factor - self.alpha_color = config.alpha_color - - def _generate_dataset_inputs(self, split="train"): - if self.alpha_color is not None: - alpha_color_tensor = get_color(self.alpha_color) - else: - alpha_color_tensor = None - - abs_dir = get_absolute_path(self.data_directory) - meta = load_from_json(abs_dir / f"transforms_{split}.json") - image_filenames = [] - poses = [] - for frame in meta["frames"]: - fname = abs_dir / Path(frame["file_path"].replace("./", "") + ".png") - image_filenames.append(fname) - poses.append(np.array(frame["transform_matrix"])) - poses = np.array(poses).astype(np.float32) - - img_0 = imageio.imread(image_filenames[0]) - image_height, image_width = img_0.shape[:2] - camera_angle_x = float(meta["camera_angle_x"]) - focal_length = 0.5 * image_width / np.tan(0.5 * camera_angle_x) - - cx = image_width / 2.0 - cy = image_height / 2.0 - camera_to_world = torch.from_numpy(poses[:, :3]) # camera to world transform - - # in x,y,z order - camera_to_world[..., 3] *= self.scale_factor - scene_bounds = SceneBounds(aabb=torch.tensor([[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]], dtype=torch.float32)) - - cameras = Cameras( - camera_to_worlds=camera_to_world, - fx=focal_length, - fy=focal_length, - cx=cx, - cy=cy, - camera_type=CameraType.PERSPECTIVE, - ) - - dataset_inputs = DatasetInputs( - image_filenames=image_filenames, - cameras=cameras, - alpha_color=alpha_color_tensor, - scene_bounds=scene_bounds, - ) - - return dataset_inputs - - -@dataclass -class InstantNGP(DataParser): - """Instant NGP Dataset - - Args: - data_directory: Location of data - scale_factor: How much to scale the camera origins by. - scene_scale: How much to scale the scene. Defaults to 0.33 - """ - - config: cfg.InstantNGPDataParserConfig - - def _generate_dataset_inputs(self, split="train"): - - abs_dir = get_absolute_path(self.config.data_directory) - - meta = load_from_json(abs_dir / "transforms.json") - image_filenames = [] - poses = [] - num_skipped_image_filenames = 0 - for frame in meta["frames"]: - fname = abs_dir / Path(frame["file_path"]) - if not fname: - num_skipped_image_filenames += 1 - else: - image_filenames.append(fname) - poses.append(np.array(frame["transform_matrix"])) - if num_skipped_image_filenames >= 0: - logging.info("Skipping %s files in dataset split %s.", num_skipped_image_filenames, split) - assert ( - len(image_filenames) != 0 - ), """ - No image files found. - You should check the file_paths in the transforms.json file to make sure they are correct. - """ - poses = np.array(poses).astype(np.float32) - poses[:3, 3] *= self.config.scene_scale - - camera_to_world = torch.from_numpy(poses[:, :3]) # camera to world transform - - distortion_params = camera_utils.get_distortion_params( - k1=float(meta["k1"]), k2=float(meta["k2"]), p1=float(meta["p1"]), p2=float(meta["p2"]) - ) - - # in x,y,z order - # assumes that the scene is centered at the origin - aabb_scale = meta["aabb_scale"] - scene_bounds = SceneBounds( - aabb=torch.tensor( - [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 - ) - ) - - cameras = Cameras( - fx=float(meta["fl_x"]), - fy=float(meta["fl_y"]), - cx=float(meta["cx"]), - cy=float(meta["cy"]), - distortion_params=distortion_params, - height=int(meta["h"]), - width=int(meta["w"]), - camera_to_worlds=camera_to_world, - camera_type=CameraType.PERSPECTIVE, - ) - - # TODO(ethan): add alpha background color - dataset_inputs = DatasetInputs( - image_filenames=image_filenames, - cameras=cameras, - scene_bounds=scene_bounds, - ) - - return dataset_inputs - - -@dataclass -class Mipnerf360(DataParser): - """MipNeRF 360 Dataset - - Args: - data_directory: Location of data - downscale_factor: How much to downscale images. Defaults to 1. - val_skip: 1/val_skip images to use for validation. Defaults to 8. - auto_scale: Scale based on pose bounds. Defaults to True. - aabb_scale: Scene scale, Defaults to 1.0. - """ - - config: cfg.MipNerf360DataParserConfig - - @classmethod - def normalize_orientation(cls, poses: np.ndarray): - """Set the _up_ direction to be in the positive Y direction. - - Args: - poses: Numpy array of poses. - """ - poses_orig = poses.copy() - bottom = np.reshape([0, 0, 0, 1.0], [1, 4]) - center = poses[:, :3, 3].mean(0) - vec2 = poses[:, :3, 2].sum(0) / np.linalg.norm(poses[:, :3, 2].sum(0)) - up = poses[:, :3, 1].sum(0) - vec0 = np.cross(up, vec2) / np.linalg.norm(np.cross(up, vec2)) - vec1 = np.cross(vec2, vec0) / np.linalg.norm(np.cross(vec2, vec0)) - c2w = np.stack([vec0, vec1, vec2, center], -1) # [3, 4] - c2w = np.concatenate([c2w[:3, :4], bottom], -2) # [4, 4] - bottom = np.tile(np.reshape(bottom, [1, 1, 4]), [poses.shape[0], 1, 1]) # [BS, 1, 4] - poses = np.concatenate([poses[:, :3, :4], bottom], -2) # [BS, 4, 4] - poses = np.linalg.inv(c2w) @ poses - poses_orig[:, :3, :4] = poses[:, :3, :4] - return poses_orig - - def _generate_dataset_inputs(self, split="train"): - abs_dir = get_absolute_path(self.config.data_directory) - image_dir = "images" - if self.config.downscale_factor > 1: - image_dir += f"_{self.config.downscale_factor}" - image_dir = abs_dir / image_dir - if not image_dir.exists(): - raise ValueError(f"Image directory {image_dir} doesn't exist") - - valid_formats = [".jpg", ".png"] - image_filenames = [] - for f in image_dir.iterdir(): - ext = f.suffix - if ext.lower() not in valid_formats: - continue - image_filenames.append(image_dir / f) - image_filenames = sorted(image_filenames) - num_images = len(image_filenames) - - poses_data = np.load(abs_dir / "poses_bounds.npy") - poses = poses_data[:, :-2].reshape([-1, 3, 5]).astype(np.float32) - bounds = poses_data[:, -2:].transpose([1, 0]) - - if num_images != poses.shape[0]: - raise RuntimeError(f"Different number of images ({num_images}), and poses ({poses.shape[0]})") - - idx_test = np.arange(num_images)[:: self.config.val_skip] - idx_train = np.array([i for i in np.arange(num_images) if i not in idx_test]) - idx = idx_train if split == "train" else idx_test - - image_filenames = np.array(image_filenames)[idx] - poses = poses[idx] - - img_0 = imageio.imread(image_filenames[0]) - image_height, image_width = img_0.shape[:2] - - poses[:, :2, 4] = np.array([image_height, image_width]) - poses[:, 2, 4] = poses[:, 2, 4] - - # Reorder pose to match our convention - poses = np.concatenate([poses[:, :, 1:2], -poses[:, :, 0:1], poses[:, :, 2:]], axis=-1) - - # Center poses and rotate. (Compute up from average of all poses) - poses = self.normalize_orientation(poses) - - # Scale factor used in mipnerf - if self.config.auto_scale: - scale_factor = 1 / (np.min(bounds) * 0.75) - poses[:, :3, 3] *= scale_factor - bounds *= scale_factor - - # Center poses - poses[:, :3, 3] = poses[:, :3, 3] - np.mean(poses[:, :3, :], axis=0)[:, 3] - - focal_length = poses[0, -1, -1] - - cx = image_width / 2.0 - cy = image_height / 2.0 - camera_to_world = torch.from_numpy(poses[:, :3, :4]) # camera to world transform - num_cameras = len(image_filenames) - num_intrinsics_params = 3 - intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32) - intrinsics *= torch.tensor([cx, cy, focal_length]) - - aabb = torch.tensor([[-4, -4, -4], [4, 4, 4]], dtype=torch.float32) * self.config.aabb_scale - scene_bounds = SceneBounds(aabb=aabb) - - cameras = Cameras( - fx=focal_length, - fy=focal_length, - cx=cx, - cy=cy, - camera_to_worlds=camera_to_world, - camera_type=CameraType.PERSPECTIVE, - ) - - cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor) - - dataset_inputs = DatasetInputs( - image_filenames=image_filenames, - cameras=cameras, - scene_bounds=scene_bounds, - ) - - return dataset_inputs - - -@dataclass -class Record3D(DataParser): - """Record3D Dataset - - Args: - data_directory: Location of data - val_skip: 1/val_skip images to use for validation. Defaults to 8. - aabb_scale: Scene scale, Defaults to 4.0. - max_dataset_size: Max number of images to train on. If the dataset has - more, images will be sampled approximately evenly. Defaults to 150. - """ - - config: cfg.Record3DDataParserConfig - - def _generate_dataset_inputs(self, split: str = "train") -> DatasetInputs: - abs_dir = get_absolute_path(self.config.data_directory) - - image_dir = abs_dir / "rgb" - - if not image_dir.exists(): - raise ValueError(f"Image directory {image_dir} doesn't exist") - - image_filenames = [] - for f in image_dir.iterdir(): - if f.stem.isdigit(): # removes possible duplicate images (for example, 123(3).jpg) - image_filenames.append(image_dir / f) - - image_filenames = sorted(image_filenames, key=lambda fn: int(fn.stem)) - image_filenames = np.array(image_filenames) - num_images = len(image_filenames) - - metadata_path = abs_dir / "metadata.json" - metadata_dict = load_from_json(metadata_path) - - poses_data = np.array(metadata_dict["poses"]) - # (N, 3, 4) - poses = np.concatenate( - [Rotation.from_quat(poses_data[:, :4]).as_matrix(), poses_data[:, 4:, None]], - axis=-1, - ).astype(np.float32) - - if num_images > self.config.max_dataset_size: - # Evenly select max_dataset_size images from dataset, including first - # and last indices. - idx = np.round(np.linspace(0, num_images - 1, self.config.max_dataset_size)).astype(int) - poses = poses[idx] - image_filenames = image_filenames[idx] - num_images = len(image_filenames) - - # Normalization similar to Mipnerf360 - poses = Mipnerf360.normalize_orientation(poses) - - bottom = np.reshape([0, 0, 0, 1.0], [1, 4]) - bottom = np.tile(np.reshape(bottom, [1, 1, 4]), [poses.shape[0], 1, 1]) - poses = np.concatenate([poses[:, :3, :4], bottom], -2).astype(np.float32) - - rotation_matrix = np.array( - [ - [1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, -1.0, 0.0], - [0.0, 1.0, 0.0, 0.0], - [0.0, 0.0, 0.0, 1.0], - ], - dtype=np.float32, - ) - poses = rotation_matrix @ poses - - idx_test = np.arange(num_images)[:: self.config.val_skip] - idx_train = np.array([i for i in np.arange(num_images) if i not in idx_test]) - idx = idx_train if split == "train" else idx_test - if num_images != poses.shape[0]: - raise RuntimeError(f"Different number of images ({num_images}), and poses ({poses.shape[0]})") - - image_filenames = image_filenames[idx] - poses = poses[idx] - - # Centering poses - poses[:, :3, 3] = poses[:, :3, 3] - np.mean(poses[:, :3, 3], axis=0) - - camera_to_world = torch.from_numpy(poses[:, :3, :4]) # camera to world transform - - # Camera intrinsics - K = np.array(metadata_dict["K"]).reshape((3, 3)).T - focal_length = K[0, 0] - - H = metadata_dict["h"] - W = metadata_dict["w"] - - # TODO(akristoffersen): The metadata dict comes with principle points, - # but caused errors in image coord indexing. Should update once that is fixed. - cx, cy = W / 2, H / 2 - - num_cameras = len(image_filenames) - num_intrinsics_params = 3 - intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32) - intrinsics *= torch.tensor([cx, cy, focal_length]) - - # scene_bounds = SceneBounds.from_camera_poses(camera_to_world, self.config.aabb_scale) - aabb = torch.tensor([[-1, -1, -1], [1, 1, 1]], dtype=torch.float32) * self.config.aabb_scale - scene_bounds = SceneBounds(aabb=aabb) - - cameras = Cameras( - fx=focal_length, - fy=focal_length, - cx=cx, - cy=cy, - camera_to_worlds=camera_to_world, - camera_type=CameraType.PERSPECTIVE, - ) - - dataset_inputs = DatasetInputs( - image_filenames=image_filenames, - cameras=cameras, - scene_bounds=scene_bounds, - ) - - return dataset_inputs - - -@dataclass -class Friends(DataParser): - """Friends Dataset - - Args: - data_directory: Location of data - include_semantics: whether or not to include the semantics. Defaults to False. - include_point_cloud: whether or not to include the point cloud. Defaults to False. - """ - - data_directory: Path - include_semantics: bool = True - include_point_cloud: bool = False - - @classmethod - def _get_aabb_and_transform(cls, basedir): - """Returns the aabb and pointcloud transform from the threejs.json file.""" - 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=too-many-statements - - abs_dir = get_absolute_path(self.data_directory) - - images_data = read_images_binary(abs_dir / "colmap" / "images.bin") - # `image_path` is only the end of the filename, including the extension e.g., `.jpg` - image_paths = sorted((abs_dir / "images").iterdir()) - - image_path_to_image_id = {} - image_id_to_image_path = {} - for v in images_data.values(): - image_path_to_image_id[v.name] = v.id - image_id_to_image_path[v.id] = v.name - # TODO: handle the splits differently - image_filenames = [abs_dir / "images" / image_path for image_path in image_paths] - - # -- set the bounding box --- - aabb, transposed_point_cloud_transform = self._get_aabb_and_transform(abs_dir) - scene_bounds_original = SceneBounds(aabb=aabb) - # 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) - - # --- intrinsics --- - cameras_data = read_cameras_binary(abs_dir / "colmap" / "cameras.bin") - focal_lengths = [] - for image_path in image_paths: - cam = cameras_data[image_path_to_image_id[image_path]] - assert len(cam.params) == 3 - focal_lengths.append(cam.params[0]) # f (fx and fy) - focal_lengths = torch.stack(focal_lengths, dim=0) - - cam = cameras_data[image_path_to_image_id[image_paths[0]]] - cx = cam.params[1] # cx - cy = cam.params[2] # cy - - # --- camera_to_world (extrinsics) --- - camera_to_world = [] - bottom_row = np.array([0, 0, 0, 1.0]).reshape(1, 4) - for image_path in image_paths: - image_data = images_data[image_path_to_image_id[image_path]] - rot = image_data.qvec2rotmat() - trans = image_data.tvec.reshape(3, 1) - c2w = np.concatenate([np.concatenate([rot, trans], 1), bottom_row], 0) - camera_to_world.append(c2w) - camera_to_world = torch.tensor(np.array(camera_to_world)).float() - camera_to_world = torch.inverse(camera_to_world) - camera_to_world[..., 1:3] *= -1 - camera_to_world = transposed_point_cloud_transform @ camera_to_world - camera_to_world = camera_to_world[:, :3] - camera_to_world[..., 3] = (camera_to_world[..., 3] - box_center) * box_scale_factor # center and rescale - - # --- semantics --- - semantics = None - if self.include_semantics: - thing_filenames = [ - Path(str(image_filename).replace("/images/", "/segmentations/thing/").replace(".jpg", ".png")) - for image_filename in image_filenames - ] - stuff_filenames = [ - Path(str(image_filename).replace("/images/", "/segmentations/stuff/").replace(".jpg", ".png")) - for image_filename in image_filenames - ] - panoptic_classes = load_from_json(abs_dir / "panoptic_classes.json") - stuff_classes = panoptic_classes["stuff"] - stuff_colors = torch.tensor(panoptic_classes["stuff_colors"], dtype=torch.float32) / 255.0 - thing_classes = panoptic_classes["thing"] - thing_colors = torch.tensor(panoptic_classes["thing_colors"], dtype=torch.float32) / 255.0 - semantics = Semantics( - stuff_classes=stuff_classes, - stuff_colors=stuff_colors, - stuff_filenames=stuff_filenames, - thing_classes=thing_classes, - thing_colors=thing_colors, - thing_filenames=thing_filenames, - ) - - # Possibly include the sparse point cloud from COLMAP in the dataset inputs. - # NOTE(ethan): this will be common across the different splits. - point_cloud = PointCloud() - if self.include_point_cloud: - points_3d = read_points3d_binary(abs_dir / "colmap" / "points3D.bin") - xyz = torch.tensor(np.array([p_value.xyz for p_id, p_value in points_3d.items()])).float() - rgb = torch.tensor(np.array([p_value.rgb for p_id, p_value in points_3d.items()])).float() - xyz_h = torch.cat([xyz, torch.ones_like(xyz[..., :1])], -1) - xyz = (xyz_h @ transposed_point_cloud_transform.T)[..., :3] - xyz = (xyz - box_center) * box_scale_factor # center and rescale - point_cloud.xyz = xyz - point_cloud.rgb = rgb - - cameras = Cameras( - fx=focal_lengths, - fy=focal_lengths, - cx=cx, - cy=cy, - camera_to_worlds=camera_to_world, - camera_type=CameraType.PERSPECTIVE, - ) - - dataset_inputs = DatasetInputs( - image_filenames=image_filenames, - cameras=cameras, - semantics=semantics, - point_cloud=point_cloud, - scene_bounds=scene_bounds, - ) - return dataset_inputs diff --git a/nerfactory/dataloaders/__init__.py b/nerfactory/datamanagers/__init__.py similarity index 100% rename from nerfactory/dataloaders/__init__.py rename to nerfactory/datamanagers/__init__.py diff --git a/nerfactory/dataloaders/base.py b/nerfactory/datamanagers/base.py similarity index 74% rename from nerfactory/dataloaders/base.py rename to nerfactory/datamanagers/base.py index d4df44aa8..2d8bb0c4f 100644 --- a/nerfactory/dataloaders/base.py +++ b/nerfactory/datamanagers/base.py @@ -28,14 +28,12 @@ from torch.utils.data.distributed import DistributedSampler from nerfactory.cameras.rays import RayBundle from nerfactory.configs import base as cfg -from nerfactory.dataloaders.image_dataloader import ( +from nerfactory.datamanagers.dataloaders import ( CacheImageDataloader, - EvalDataloader, FixedIndicesEvalDataloader, ) -from nerfactory.dataloaders.image_dataset import ImageDataset, PanopticImageDataset -from nerfactory.dataloaders.pixel_sampler import PixelSampler -from nerfactory.dataloaders.structs import DatasetInputs +from nerfactory.datamanagers.datasets import InputDataset +from nerfactory.datamanagers.pixel_sampler import PixelSampler from nerfactory.models.modules.ray_generator import RayGenerator from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfactory.utils.misc import IterableWrapper @@ -75,47 +73,36 @@ class DataManager(nn.Module): next_eval: will be called on __next__() for the eval iterator get_eval_iterable: utility that gets a clean pythonic iterator for your eval data - Args: - use_train: whether this data manager is being used for training - use_eval: whether this data manager is being used for evaluation Attributes: - use_train (bool): whether or not we are using train - use_eval (bool): whether or not we are using eval train_count (int): the step number of our train iteration, needs to be incremented manually eval_count (int): the step number of our eval iteration, needs to be incremented manually - train_datasetinputs (DatasetInputs): the inputs for the train dataset - eval_datasetinputs (DatasetInputs): the inputs for the eval dataset - train_image_dataset (ImageDataset): the image dataset for the train dataset + train_input_dataset (InputDataset): the input dataset for the train dataset + eval_input_dataset (InputDataset): the input dataset for the eval dataset Additional attributes specific to each subclass are defined in the setup_train and setup_eval functions. """ - train_datasetinputs: Optional[DatasetInputs] = None - eval_datasetinputs: Optional[DatasetInputs] = None - train_image_dataset: Optional[ImageDataset] = None - eval_dataloader: Optional[EvalDataloader] = None + train_input_dataset: Optional[InputDataset] = None + eval_input_dataset: Optional[InputDataset] = None - def __init__(self, use_train: bool, use_eval: bool): + def __init__(self): """Constructor for the DataManager class. Subclassed DataManagers will likely need to override this constructor. If you aren't manually calling the setup_train and setup_eval functions from an overriden - constructor, that you call super().__init__(use_train, use_eval) BEFORE you initialize any + constructor, that you call super().__init__() BEFORE you initialize any nn.Modules or nn.Parameters, but AFTER you've already set all the attributes you need for the setup functions.""" super().__init__() - self.use_train = use_train - self.use_eval = use_eval self.train_count = 0 self.eval_count = 0 - assert use_train or use_eval - if use_train: + if self.train_input_dataset and self.train_input_dataset.inputs: self.setup_train() - if use_eval: + if self.eval_input_dataset and self.eval_input_dataset.inputs: self.setup_eval() def forward(self): @@ -144,7 +131,7 @@ class DataManager(nn.Module): as __iter__ and __next__ methods respectivley. This basically is just a little utility if you want to do something like: - | for ray_bundle, batch in data_manager.get_train_iterable(): + | for ray_bundle, batch in datamanager.get_train_iterable(): | since the returned IterableWrapper is just an iterator with the __iter__ and __next__ methods (methods bound to our DataManager instance in this case) specified in the constructor. @@ -156,7 +143,7 @@ class DataManager(nn.Module): as __iter__ and __next__ methods respectivley. This basically is just a little utility if you want to do something like: - | for ray_bundle, batch in data_manager.get_eval_iterable(): + | for ray_bundle, batch in datamanager.get_eval_iterable(): | since the returned IterableWrapper is just an iterator with the __iter__ and __next__ methods (methods bound to our DataManager instance in this case) specified in the constructor. @@ -233,70 +220,43 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self.local_rank = local_rank self.sampler = None - dataset_train = config.train_dataset.setup() - self.train_datasetinputs = dataset_train.get_dataset_inputs(split="train") - self.train_cameras = self.train_datasetinputs.cameras.to(device) - if config.eval_dataset is not None: - dataset_eval = config.eval_dataset.setup() - else: + if config.eval_dataparser is None: logging.info("No eval dataset specified so using train dataset for eval.") - dataset_eval = dataset_train - self.eval_datasetinputs = dataset_eval.get_dataset_inputs(split="val" if not test_mode else "test") - self.eval_cameras = self.eval_datasetinputs.cameras.to(device) - use_train = self.train_datasetinputs is not None - use_eval = self.eval_datasetinputs is not None - super().__init__(use_train, use_eval) + config.eval_dataparser = config.train_dataparser + self.train_input_dataset = InputDataset(config.train_dataparser, split="train") + self.eval_input_dataset = InputDataset(config.eval_dataparser, split="val" if not test_mode else "test") + super().__init__() def setup_train(self): - """Sets up the data manager for training""" - assert self.train_datasetinputs is not None - if self.config.image_dataset_type == "rgb": - self.train_image_dataset = ImageDataset( - image_filenames=self.train_datasetinputs.image_filenames, - alpha_color=self.train_datasetinputs.alpha_color, - ) - elif self.config.image_dataset_type == "panoptic": - self.train_image_dataset = PanopticImageDataset(**self.train_datasetinputs.as_dict()) - else: - raise ValueError(f"Unknown image dataset type {self.config.image_dataset_type}") - + """Sets up the data loaders for training""" + assert self.train_input_dataset is not None if self.world_size > 1: self.sampler = DistributedSampler( - self.train_image_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42 + self.train_input_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42 ) self.train_image_dataloader = CacheImageDataloader( - self.train_image_dataset, + self.train_input_dataset, num_images_to_sample_from=self.config.train_num_images_to_sample_from, device=self.device, num_workers=self.world_size * 4, pin_memory=True, sampler=self.sampler, ) # TODO(ethan): pass this in - else: self.train_image_dataloader = CacheImageDataloader( - self.train_image_dataset, + self.train_input_dataset, num_images_to_sample_from=self.config.train_num_images_to_sample_from, device=self.device, ) self.iter_train_image_dataloader = iter(self.train_image_dataloader) self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch) - self.train_ray_generator = RayGenerator(self.train_cameras) + self.train_ray_generator = RayGenerator(self.train_input_dataset.inputs.cameras.to(self.device)) def setup_eval(self): - """Sets up the data manager for evaluation""" - assert self.eval_datasetinputs is not None - if self.config.image_dataset_type == "rgb": - self.eval_image_dataset = ImageDataset( - image_filenames=self.eval_datasetinputs.image_filenames, - alpha_color=self.eval_datasetinputs.alpha_color, - ) - elif self.config.image_dataset_type == "panoptic": - self.eval_image_dataset = PanopticImageDataset(**self.eval_datasetinputs.as_dict()) - + """Sets up the data loader for evaluation""" + assert self.eval_input_dataset is not None self.eval_dataloader = FixedIndicesEvalDataloader( - image_dataset=self.eval_image_dataset, - cameras=self.eval_cameras, + input_dataset=self.eval_input_dataset, num_rays_per_chunk=self.config.eval_num_rays_per_chunk, image_indices=self.config.eval_image_indices, device=self.device, diff --git a/nerfactory/dataloaders/image_dataloader.py b/nerfactory/datamanagers/dataloaders.py similarity index 89% rename from nerfactory/dataloaders/image_dataloader.py rename to nerfactory/datamanagers/dataloaders.py index a90d26ad0..0614fb800 100644 --- a/nerfactory/dataloaders/image_dataloader.py +++ b/nerfactory/datamanagers/dataloaders.py @@ -26,13 +26,13 @@ from torch.utils.data.dataloader import DataLoader from nerfactory.cameras.cameras import Cameras from nerfactory.cameras.rays import RayBundle -from nerfactory.dataloaders.image_dataset import ImageDataset +from nerfactory.datamanagers.datasets import InputDataset from nerfactory.utils.misc import get_dict_to_torch class CacheImageDataloader(DataLoader): """Collated image dataset that implements caching of images. - Creates batches of the ImageDataset return type. + Creates batches of the InputDataset return type. Args: dataset (torch.utils.data.Dataset): Dataset to sample from. @@ -43,7 +43,7 @@ class CacheImageDataloader(DataLoader): def __init__( self, - dataset: ImageDataset, + dataset: InputDataset, num_images_to_sample_from: int = -1, num_times_to_repeat_images: int = 0, device: Union[torch.device, str] = "cpu", @@ -103,7 +103,7 @@ class EvalDataloader(DataLoader): """Evaluation dataloader base class Args: - image_dataset: ImageDataset to load data from + input_dataset: InputDataset to load data from cameras: Cameras to use for generating rays num_rays_per_chunk: Number of camera rays to generate per chunk device: Device to load data to @@ -111,18 +111,17 @@ class EvalDataloader(DataLoader): def __init__( self, - image_dataset: ImageDataset, - cameras: Cameras, + input_dataset: InputDataset, num_rays_per_chunk: int, device: Union[torch.device, str] = "cpu", **kwargs, ): - self.image_dataset = image_dataset - self.cameras = cameras + self.input_dataset = input_dataset + self.cameras = input_dataset.inputs.cameras.to(device) self.num_rays_per_chunk = num_rays_per_chunk self.device = device self.kwargs = kwargs - super().__init__(dataset=image_dataset) + super().__init__(dataset=input_dataset) @abstractmethod def __iter__(self): @@ -159,7 +158,7 @@ class EvalDataloader(DataLoader): ray_bundle = self.cameras.generate_rays(camera_indices=image_idx) ray_bundle.num_rays_per_chunk = self.num_rays_per_chunk ray_bundle.camera_indices = torch.Tensor([image_idx])[..., None].int() - batch = self.image_dataset[image_idx] + batch = self.input_dataset[image_idx] batch = get_dict_to_torch(batch, device=self.device) return ray_bundle, batch @@ -169,8 +168,7 @@ class FixedIndicesEvalDataloader(EvalDataloader): def __init__( self, - image_dataset: ImageDataset, - cameras: Cameras, + input_dataset: InputDataset, num_rays_per_chunk: int, image_indices: Optional[Tuple[int]] = None, device: Union[torch.device, str] = "cpu", @@ -178,15 +176,15 @@ class FixedIndicesEvalDataloader(EvalDataloader): ): """ Args: - image_dataset: ImageDataset to load data from + input_dataset: InputDataset to load data from cameras: Cameras to use for generating rays num_rays_per_chunk: Number of camera rays to generate per chunk image_indices: List of image indices to load data from. If None, then use all images. device: Device to load data to """ - super().__init__(image_dataset, cameras, num_rays_per_chunk, device, **kwargs) + super().__init__(input_dataset, num_rays_per_chunk, device, **kwargs) if image_indices is None: - self.image_indices = list(range(len(image_dataset))) + self.image_indices = list(range(len(input_dataset))) else: self.image_indices = image_indices self.num_rays_per_chunk = num_rays_per_chunk @@ -210,7 +208,7 @@ class RandIndicesEvalDataloader(EvalDataloader): """Dataloader that returns random images. Args: - image_dataset: ImageDataset to load data from + input_dataset: InputDataset to load data from cameras: Cameras to use for generating rays num_rays_per_chunk: Number of camera rays to generate per chunk device: Device to load data to @@ -218,13 +216,12 @@ class RandIndicesEvalDataloader(EvalDataloader): def __init__( self, - image_dataset: ImageDataset, - cameras: Cameras, + input_dataset: InputDataset, num_rays_per_chunk: int, device: Union[torch.device, str] = "cpu", **kwargs, ): - super().__init__(image_dataset, cameras, num_rays_per_chunk, device, **kwargs) + super().__init__(input_dataset, num_rays_per_chunk, device, **kwargs) self.count = 0 def __iter__(self): diff --git a/nerfactory/datamanagers/dataparsers/base.py b/nerfactory/datamanagers/dataparsers/base.py new file mode 100644 index 000000000..4f18b7679 --- /dev/null +++ b/nerfactory/datamanagers/dataparsers/base.py @@ -0,0 +1,106 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A set of standard datasets.""" + +from __future__ import annotations + +import dataclasses +import logging +from abc import abstractmethod +from dataclasses import dataclass +from pathlib import Path + +from nerfactory.configs import base as cfg +from nerfactory.datamanagers.structs import DatasetInputs +from nerfactory.utils.io import get_absolute_path, load_from_pkl, write_to_pkl +from nerfactory.utils.misc import get_hash_str_from_dict + + +@dataclass +class DataParser: + """A dataset.""" + + def __init__(self, config: cfg.DataParserConfig): + super().__init__() + self.config = config + + @abstractmethod + def _generate_dataset_inputs(self, split: str = "train") -> DatasetInputs: + """Returns the dataset inputs for the given split. + + Args: + split: Which dataset split to generate. + + Returns: + DatasetInputs + """ + + def get_dataset_inputs(self, split: str = "train", use_preprocessing_cache: bool = False) -> DatasetInputs: + """Returns the dataset inputs for the given split. + + Args: + split: Which dataset split to generate. + use_preprocessing_cache: Whether to use the cached dataset inputs. Defaults to False. + + Returns: + DatasetInputs + """ + if use_preprocessing_cache: + dataset_inputs = self._get_dataset_inputs_from_cache(split) + else: + dataset_inputs = self._generate_dataset_inputs(split) + return dataset_inputs + + def _get_cache_filename(self, split: str) -> Path: + """Creates a cache filename from the dataset inputs arguments. + + Args: + split: Which dataset split to generate filename for. + + Returns: + filename for cache. + """ + dataset_config_hash = get_hash_str_from_dict(dataclasses.asdict(self)) + dataset_config_hash_filename = get_absolute_path( + f"cache/dataset_inputs/{dataset_config_hash}-{split}.pkl" + ).parent.mkdir() + assert dataset_config_hash_filename is not None, "dataset hash is None" + return dataset_config_hash_filename + + def save_dataset_inputs_to_cache(self, split: str): + """Saves the dataset inputs to cache. + + Args: + split: Which dataset split to save. + """ + dataset_inputs = self.get_dataset_inputs(split=split) + dataset_config_hash_filename = self._get_cache_filename(split) + write_to_pkl(dataset_config_hash_filename, dataset_inputs) + + def _get_dataset_inputs_from_cache(self, split: str) -> DatasetInputs: + """Loads the dataset inputs from cache. If the cache does not exist, it will be created. + + Args: + split: Which dataset split to load. + """ + dataset_config_hash_filename = self._get_cache_filename(split) + if dataset_config_hash_filename.exists(): + logging.info("Loading dataset from cache.") + dataset_inputs = load_from_pkl(dataset_config_hash_filename) + else: + logging.info("Cache file not found. Generating and saving dataset to cache.") + dataset_inputs = self._generate_dataset_inputs(split=split) + self.save_dataset_inputs_to_cache(split) + return dataset_inputs diff --git a/nerfactory/datamanagers/dataparsers/blender_parser.py b/nerfactory/datamanagers/dataparsers/blender_parser.py new file mode 100644 index 000000000..69a6ee9f9 --- /dev/null +++ b/nerfactory/datamanagers/dataparsers/blender_parser.py @@ -0,0 +1,95 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Data parser for blender dataset""" +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import imageio +import numpy as np +import torch + +from nerfactory.cameras.cameras import Cameras, CameraType +from nerfactory.configs import base as cfg +from nerfactory.datamanagers.dataparsers.base import DataParser +from nerfactory.datamanagers.structs import DatasetInputs, SceneBounds +from nerfactory.utils.colors import get_color +from nerfactory.utils.io import get_absolute_path, load_from_json + + +@dataclass +class Blender(DataParser): + """Blender Dataset + Some of this code comes from https://github.com/yenchenlin/nerf-pytorch/blob/master/load_blender.py#L37. + + Args: + data_directory: Location of data + alpha_color: Sets transparent regions to specified color, otherwise black. + scale_factor: How much to scale the camera origins by. + """ + + def __init__(self, config: cfg.BlenderDataParserConfig): + super().__init__(config=config) + self.data_directory: Path = config.data_directory + self.scale_factor: float = config.scale_factor + self.alpha_color = config.alpha_color + + def _generate_dataset_inputs(self, split="train"): + if self.alpha_color is not None: + alpha_color_tensor = get_color(self.alpha_color) + else: + alpha_color_tensor = None + + abs_dir = get_absolute_path(self.data_directory) + meta = load_from_json(abs_dir / f"transforms_{split}.json") + image_filenames = [] + poses = [] + for frame in meta["frames"]: + fname = abs_dir / Path(frame["file_path"].replace("./", "") + ".png") + image_filenames.append(fname) + poses.append(np.array(frame["transform_matrix"])) + poses = np.array(poses).astype(np.float32) + + img_0 = imageio.imread(image_filenames[0]) + image_height, image_width = img_0.shape[:2] + camera_angle_x = float(meta["camera_angle_x"]) + focal_length = 0.5 * image_width / np.tan(0.5 * camera_angle_x) + + cx = image_width / 2.0 + cy = image_height / 2.0 + camera_to_world = torch.from_numpy(poses[:, :3]) # camera to world transform + + # in x,y,z order + camera_to_world[..., 3] *= self.scale_factor + scene_bounds = SceneBounds(aabb=torch.tensor([[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]], dtype=torch.float32)) + + cameras = Cameras( + camera_to_worlds=camera_to_world, + fx=focal_length, + fy=focal_length, + cx=cx, + cy=cy, + camera_type=CameraType.PERSPECTIVE, + ) + + dataset_inputs = DatasetInputs( + image_filenames=image_filenames, + cameras=cameras, + alpha_color=alpha_color_tensor, + scene_bounds=scene_bounds, + ) + + return dataset_inputs diff --git a/nerfactory/datamanagers/dataparsers/friends_parser.py b/nerfactory/datamanagers/dataparsers/friends_parser.py new file mode 100644 index 000000000..728c2c743 --- /dev/null +++ b/nerfactory/datamanagers/dataparsers/friends_parser.py @@ -0,0 +1,183 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Data parser for friends dataset""" +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch + +from nerfactory.cameras.cameras import Cameras, CameraType +from nerfactory.datamanagers.dataparsers.base import DataParser +from nerfactory.datamanagers.structs import ( + DatasetInputs, + PointCloud, + SceneBounds, + Semantics, +) +from nerfactory.utils.colmap_utils import ( + read_cameras_binary, + read_images_binary, + read_points3d_binary, +) +from nerfactory.utils.io import get_absolute_path, load_from_json + + +@dataclass +class Friends(DataParser): + """Friends Dataset + + Args: + data_directory: Location of data + include_semantics: whether or not to include the semantics. Defaults to False. + include_point_cloud: whether or not to include the point cloud. Defaults to False. + """ + + data_directory: Path + include_semantics: bool = True + include_point_cloud: bool = False + + @classmethod + def _get_aabb_and_transform(cls, basedir): + """Returns the aabb and pointcloud transform from the threejs.json file.""" + 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=too-many-statements + + abs_dir = get_absolute_path(self.data_directory) + + images_data = read_images_binary(abs_dir / "colmap" / "images.bin") + # `image_path` is only the end of the filename, including the extension e.g., `.jpg` + image_paths = sorted((abs_dir / "images").iterdir()) + + image_path_to_image_id = {} + image_id_to_image_path = {} + for v in images_data.values(): + image_path_to_image_id[v.name] = v.id + image_id_to_image_path[v.id] = v.name + # TODO: handle the splits differently + image_filenames = [abs_dir / "images" / image_path for image_path in image_paths] + + # -- set the bounding box --- + aabb, transposed_point_cloud_transform = self._get_aabb_and_transform(abs_dir) + scene_bounds_original = SceneBounds(aabb=aabb) + # 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) + + # --- intrinsics --- + cameras_data = read_cameras_binary(abs_dir / "colmap" / "cameras.bin") + focal_lengths = [] + for image_path in image_paths: + cam = cameras_data[image_path_to_image_id[image_path]] + assert len(cam.params) == 3 + focal_lengths.append(cam.params[0]) # f (fx and fy) + focal_lengths = torch.stack(focal_lengths, dim=0) + + cam = cameras_data[image_path_to_image_id[image_paths[0]]] + cx = cam.params[1] # cx + cy = cam.params[2] # cy + + # --- camera_to_world (extrinsics) --- + camera_to_world = [] + bottom_row = np.array([0, 0, 0, 1.0]).reshape(1, 4) + for image_path in image_paths: + image_data = images_data[image_path_to_image_id[image_path]] + rot = image_data.qvec2rotmat() + trans = image_data.tvec.reshape(3, 1) + c2w = np.concatenate([np.concatenate([rot, trans], 1), bottom_row], 0) + camera_to_world.append(c2w) + camera_to_world = torch.tensor(np.array(camera_to_world)).float() + camera_to_world = torch.inverse(camera_to_world) + camera_to_world[..., 1:3] *= -1 + camera_to_world = transposed_point_cloud_transform @ camera_to_world + camera_to_world = camera_to_world[:, :3] + camera_to_world[..., 3] = (camera_to_world[..., 3] - box_center) * box_scale_factor # center and rescale + + # --- semantics --- + semantics = None + if self.include_semantics: + thing_filenames = [ + Path(str(image_filename).replace("/images/", "/segmentations/thing/").replace(".jpg", ".png")) + for image_filename in image_filenames + ] + stuff_filenames = [ + Path(str(image_filename).replace("/images/", "/segmentations/stuff/").replace(".jpg", ".png")) + for image_filename in image_filenames + ] + panoptic_classes = load_from_json(abs_dir / "panoptic_classes.json") + stuff_classes = panoptic_classes["stuff"] + stuff_colors = torch.tensor(panoptic_classes["stuff_colors"], dtype=torch.float32) / 255.0 + thing_classes = panoptic_classes["thing"] + thing_colors = torch.tensor(panoptic_classes["thing_colors"], dtype=torch.float32) / 255.0 + semantics = Semantics( + stuff_classes=stuff_classes, + stuff_colors=stuff_colors, + stuff_filenames=stuff_filenames, + thing_classes=thing_classes, + thing_colors=thing_colors, + thing_filenames=thing_filenames, + ) + + # Possibly include the sparse point cloud from COLMAP in the dataset inputs. + # NOTE(ethan): this will be common across the different splits. + point_cloud = PointCloud() + if self.include_point_cloud: + points_3d = read_points3d_binary(abs_dir / "colmap" / "points3D.bin") + xyz = torch.tensor(np.array([p_value.xyz for p_id, p_value in points_3d.items()])).float() + rgb = torch.tensor(np.array([p_value.rgb for p_id, p_value in points_3d.items()])).float() + xyz_h = torch.cat([xyz, torch.ones_like(xyz[..., :1])], -1) + xyz = (xyz_h @ transposed_point_cloud_transform.T)[..., :3] + xyz = (xyz - box_center) * box_scale_factor # center and rescale + point_cloud.xyz = xyz + point_cloud.rgb = rgb + + cameras = Cameras( + fx=focal_lengths, + fy=focal_lengths, + cx=cx, + cy=cy, + camera_to_worlds=camera_to_world, + camera_type=CameraType.PERSPECTIVE, + ) + + dataset_inputs = DatasetInputs( + image_filenames=image_filenames, + cameras=cameras, + semantics=semantics, + point_cloud=point_cloud, + scene_bounds=scene_bounds, + ) + return dataset_inputs diff --git a/nerfactory/datamanagers/dataparsers/instant_ngp_parser.py b/nerfactory/datamanagers/dataparsers/instant_ngp_parser.py new file mode 100644 index 000000000..9e22d007e --- /dev/null +++ b/nerfactory/datamanagers/dataparsers/instant_ngp_parser.py @@ -0,0 +1,106 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Data parser for instant ngp data""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch + +from nerfactory.cameras import utils as camera_utils +from nerfactory.cameras.cameras import Cameras, CameraType +from nerfactory.configs import base as cfg +from nerfactory.datamanagers.dataparsers.base import DataParser +from nerfactory.datamanagers.structs import DatasetInputs, SceneBounds +from nerfactory.utils.io import get_absolute_path, load_from_json + + +@dataclass +class InstantNGP(DataParser): + """Instant NGP Dataset + + Args: + data_directory: Location of data + scale_factor: How much to scale the camera origins by. + scene_scale: How much to scale the scene. Defaults to 0.33 + """ + + config: cfg.InstantNGPDataParserConfig + + def _generate_dataset_inputs(self, split="train"): + + abs_dir = get_absolute_path(self.config.data_directory) + + meta = load_from_json(abs_dir / "transforms.json") + image_filenames = [] + poses = [] + num_skipped_image_filenames = 0 + for frame in meta["frames"]: + fname = abs_dir / Path(frame["file_path"]) + if not fname: + num_skipped_image_filenames += 1 + else: + image_filenames.append(fname) + poses.append(np.array(frame["transform_matrix"])) + if num_skipped_image_filenames >= 0: + logging.info("Skipping %s files in dataset split %s.", num_skipped_image_filenames, split) + assert ( + len(image_filenames) != 0 + ), """ + No image files found. + You should check the file_paths in the transforms.json file to make sure they are correct. + """ + poses = np.array(poses).astype(np.float32) + poses[:3, 3] *= self.config.scene_scale + + camera_to_world = torch.from_numpy(poses[:, :3]) # camera to world transform + + distortion_params = camera_utils.get_distortion_params( + k1=float(meta["k1"]), k2=float(meta["k2"]), p1=float(meta["p1"]), p2=float(meta["p2"]) + ) + + # in x,y,z order + # assumes that the scene is centered at the origin + aabb_scale = meta["aabb_scale"] + scene_bounds = SceneBounds( + aabb=torch.tensor( + [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 + ) + ) + + cameras = Cameras( + fx=float(meta["fl_x"]), + fy=float(meta["fl_y"]), + cx=float(meta["cx"]), + cy=float(meta["cy"]), + distortion_params=distortion_params, + height=int(meta["h"]), + width=int(meta["w"]), + camera_to_worlds=camera_to_world, + camera_type=CameraType.PERSPECTIVE, + ) + + # TODO(ethan): add alpha background color + dataset_inputs = DatasetInputs( + image_filenames=image_filenames, + cameras=cameras, + scene_bounds=scene_bounds, + ) + + return dataset_inputs diff --git a/nerfactory/datamanagers/dataparsers/mipnerf_parser.py b/nerfactory/datamanagers/dataparsers/mipnerf_parser.py new file mode 100644 index 000000000..ecfe9172f --- /dev/null +++ b/nerfactory/datamanagers/dataparsers/mipnerf_parser.py @@ -0,0 +1,151 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Data parser for mipnerf dataset""" +from __future__ import annotations + +from dataclasses import dataclass + +import imageio +import numpy as np +import torch + +from nerfactory.cameras.cameras import Cameras, CameraType +from nerfactory.configs import base as cfg +from nerfactory.datamanagers.dataparsers.base import DataParser +from nerfactory.datamanagers.structs import DatasetInputs, SceneBounds +from nerfactory.utils.io import get_absolute_path + + +@dataclass +class Mipnerf360(DataParser): + """MipNeRF 360 Dataset + + Args: + data_directory: Location of data + downscale_factor: How much to downscale images. Defaults to 1. + val_skip: 1/val_skip images to use for validation. Defaults to 8. + auto_scale: Scale based on pose bounds. Defaults to True. + aabb_scale: Scene scale, Defaults to 1.0. + """ + + config: cfg.MipNerf360DataParserConfig + + @classmethod + def normalize_orientation(cls, poses: np.ndarray): + """Set the _up_ direction to be in the positive Y direction. + + Args: + poses: Numpy array of poses. + """ + poses_orig = poses.copy() + bottom = np.reshape([0, 0, 0, 1.0], [1, 4]) + center = poses[:, :3, 3].mean(0) + vec2 = poses[:, :3, 2].sum(0) / np.linalg.norm(poses[:, :3, 2].sum(0)) + up = poses[:, :3, 1].sum(0) + vec0 = np.cross(up, vec2) / np.linalg.norm(np.cross(up, vec2)) + vec1 = np.cross(vec2, vec0) / np.linalg.norm(np.cross(vec2, vec0)) + c2w = np.stack([vec0, vec1, vec2, center], -1) # [3, 4] + c2w = np.concatenate([c2w[:3, :4], bottom], -2) # [4, 4] + bottom = np.tile(np.reshape(bottom, [1, 1, 4]), [poses.shape[0], 1, 1]) # [BS, 1, 4] + poses = np.concatenate([poses[:, :3, :4], bottom], -2) # [BS, 4, 4] + poses = np.linalg.inv(c2w) @ poses + poses_orig[:, :3, :4] = poses[:, :3, :4] + return poses_orig + + def _generate_dataset_inputs(self, split="train"): + abs_dir = get_absolute_path(self.config.data_directory) + image_dir = "images" + if self.config.downscale_factor > 1: + image_dir += f"_{self.config.downscale_factor}" + image_dir = abs_dir / image_dir + if not image_dir.exists(): + raise ValueError(f"Image directory {image_dir} doesn't exist") + + valid_formats = [".jpg", ".png"] + image_filenames = [] + for f in image_dir.iterdir(): + ext = f.suffix + if ext.lower() not in valid_formats: + continue + image_filenames.append(image_dir / f) + image_filenames = sorted(image_filenames) + num_images = len(image_filenames) + + poses_data = np.load(abs_dir / "poses_bounds.npy") + poses = poses_data[:, :-2].reshape([-1, 3, 5]).astype(np.float32) + bounds = poses_data[:, -2:].transpose([1, 0]) + + if num_images != poses.shape[0]: + raise RuntimeError(f"Different number of images ({num_images}), and poses ({poses.shape[0]})") + + idx_test = np.arange(num_images)[:: self.config.val_skip] + idx_train = np.array([i for i in np.arange(num_images) if i not in idx_test]) + idx = idx_train if split == "train" else idx_test + + image_filenames = np.array(image_filenames)[idx] + poses = poses[idx] + + img_0 = imageio.imread(image_filenames[0]) + image_height, image_width = img_0.shape[:2] + + poses[:, :2, 4] = np.array([image_height, image_width]) + poses[:, 2, 4] = poses[:, 2, 4] + + # Reorder pose to match our convention + poses = np.concatenate([poses[:, :, 1:2], -poses[:, :, 0:1], poses[:, :, 2:]], axis=-1) + + # Center poses and rotate. (Compute up from average of all poses) + poses = self.normalize_orientation(poses) + + # Scale factor used in mipnerf + if self.config.auto_scale: + scale_factor = 1 / (np.min(bounds) * 0.75) + poses[:, :3, 3] *= scale_factor + bounds *= scale_factor + + # Center poses + poses[:, :3, 3] = poses[:, :3, 3] - np.mean(poses[:, :3, :], axis=0)[:, 3] + + focal_length = poses[0, -1, -1] + + cx = image_width / 2.0 + cy = image_height / 2.0 + camera_to_world = torch.from_numpy(poses[:, :3, :4]) # camera to world transform + num_cameras = len(image_filenames) + num_intrinsics_params = 3 + intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32) + intrinsics *= torch.tensor([cx, cy, focal_length]) + + aabb = torch.tensor([[-4, -4, -4], [4, 4, 4]], dtype=torch.float32) * self.config.aabb_scale + scene_bounds = SceneBounds(aabb=aabb) + + cameras = Cameras( + fx=focal_length, + fy=focal_length, + cx=cx, + cy=cy, + camera_to_worlds=camera_to_world, + camera_type=CameraType.PERSPECTIVE, + ) + + cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor) + + dataset_inputs = DatasetInputs( + image_filenames=image_filenames, + cameras=cameras, + scene_bounds=scene_bounds, + ) + + return dataset_inputs diff --git a/nerfactory/datamanagers/dataparsers/nerfactory_parser.py b/nerfactory/datamanagers/dataparsers/nerfactory_parser.py new file mode 100644 index 000000000..68538fc05 --- /dev/null +++ b/nerfactory/datamanagers/dataparsers/nerfactory_parser.py @@ -0,0 +1,112 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch + +from nerfactory.cameras import utils as camera_utils +from nerfactory.cameras.cameras import Cameras, CameraType +from nerfactory.configs import base as cfg +from nerfactory.datamanagers.dataparsers.base import DataParser +from nerfactory.datamanagers.structs import DatasetInputs, SceneBounds +from nerfactory.utils.io import get_absolute_path, load_from_json + + +@dataclass +class Nerfactory(DataParser): + """Nerfactory Dataset + + Args: + data_directory: Location of data + scale_factor: How much to scale the camera origins by. + downscale_factor: How much to downscale images. + scene_scale: How much to scale the scene. + orientation_method: The method to use for orientation. Either "pca" or "up". + """ + + config: cfg.NerfactoryDataParserConfig + + def _generate_dataset_inputs(self, split="train"): + + abs_dir = get_absolute_path(self.config.data_directory) + + meta = load_from_json(abs_dir / "transforms.json") + image_filenames = [] + poses = [] + num_skipped_image_filenames = 0 + for frame in meta["frames"]: + fname = abs_dir / Path(frame["file_path"]) + if self.config.downscale_factor > 1: + fname = abs_dir / f"images_{self.config.downscale_factor}" / Path(frame["file_path"]).name + else: + fname = abs_dir / Path(frame["file_path"]) + if not fname: + num_skipped_image_filenames += 1 + else: + image_filenames.append(fname) + poses.append(np.array(frame["transform_matrix"])) + if num_skipped_image_filenames >= 0: + logging.info("Skipping %s files in dataset split %s.", num_skipped_image_filenames, split) + assert ( + len(image_filenames) != 0 + ), """ + No image files found. + You should check the file_paths in the transforms.json file to make sure they are correct. + """ + poses = torch.from_numpy(np.array(poses).astype(np.float32)) + poses = camera_utils.auto_orient_poses(poses, method=self.config.orientation_method) + + # Scale poses + scale_factor = 1.0 / torch.max(torch.abs(poses[:, :3, 3])) + poses[:, :3, 3] *= scale_factor + + distortion_params = camera_utils.get_distortion_params( + k1=float(meta["k1"]), k2=float(meta["k2"]), p1=float(meta["p1"]), p2=float(meta["p2"]) + ) + + # in x,y,z order + # assumes that the scene is centered at the origin + aabb_scale = self.config.scene_scale + scene_bounds = SceneBounds( + aabb=torch.tensor( + [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 + ) + ) + + cameras = Cameras( + fx=float(meta["fl_x"]), + fy=float(meta["fl_y"]), + cx=float(meta["cx"]), + cy=float(meta["cy"]), + distortion_params=distortion_params, + height=int(meta["h"]), + width=int(meta["w"]), + camera_to_worlds=poses[:, :3, :4], + camera_type=CameraType.PERSPECTIVE, + ) + + cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor) + + dataset_inputs = DatasetInputs( + image_filenames=image_filenames, + cameras=cameras, + scene_bounds=scene_bounds, + ) + return dataset_inputs diff --git a/nerfactory/datamanagers/dataparsers/record3d_parser.py b/nerfactory/datamanagers/dataparsers/record3d_parser.py new file mode 100644 index 000000000..f966142e5 --- /dev/null +++ b/nerfactory/datamanagers/dataparsers/record3d_parser.py @@ -0,0 +1,148 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Data parser for record3d dataset""" +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import torch +from scipy.spatial.transform import Rotation + +from nerfactory.cameras.cameras import Cameras, CameraType +from nerfactory.configs import base as cfg +from nerfactory.datamanagers.dataparsers.base import DataParser +from nerfactory.datamanagers.dataparsers.mipnerf_parser import Mipnerf360 +from nerfactory.datamanagers.structs import DatasetInputs, SceneBounds +from nerfactory.utils.io import get_absolute_path, load_from_json + + +@dataclass +class Record3D(DataParser): + """Record3D Dataset + + Args: + data_directory: Location of data + val_skip: 1/val_skip images to use for validation. Defaults to 8. + aabb_scale: Scene scale, Defaults to 4.0. + max_dataset_size: Max number of images to train on. If the dataset has + more, images will be sampled approximately evenly. Defaults to 150. + """ + + config: cfg.Record3DDataParserConfig + + def _generate_dataset_inputs(self, split: str = "train") -> DatasetInputs: + abs_dir = get_absolute_path(self.config.data_directory) + + image_dir = abs_dir / "rgb" + + if not image_dir.exists(): + raise ValueError(f"Image directory {image_dir} doesn't exist") + + image_filenames = [] + for f in image_dir.iterdir(): + if f.stem.isdigit(): # removes possible duplicate images (for example, 123(3).jpg) + image_filenames.append(image_dir / f) + + image_filenames = sorted(image_filenames, key=lambda fn: int(fn.stem)) + image_filenames = np.array(image_filenames) + num_images = len(image_filenames) + + metadata_path = abs_dir / "metadata.json" + metadata_dict = load_from_json(metadata_path) + + poses_data = np.array(metadata_dict["poses"]) + # (N, 3, 4) + poses = np.concatenate( + [Rotation.from_quat(poses_data[:, :4]).as_matrix(), poses_data[:, 4:, None]], + axis=-1, + ).astype(np.float32) + + if num_images > self.config.max_dataset_size: + # Evenly select max_dataset_size images from dataset, including first + # and last indices. + idx = np.round(np.linspace(0, num_images - 1, self.config.max_dataset_size)).astype(int) + poses = poses[idx] + image_filenames = image_filenames[idx] + num_images = len(image_filenames) + + # Normalization similar to Mipnerf360 + poses = Mipnerf360.normalize_orientation(poses) + + bottom = np.reshape([0, 0, 0, 1.0], [1, 4]) + bottom = np.tile(np.reshape(bottom, [1, 1, 4]), [poses.shape[0], 1, 1]) + poses = np.concatenate([poses[:, :3, :4], bottom], -2).astype(np.float32) + + rotation_matrix = np.array( + [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, -1.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + poses = rotation_matrix @ poses + + idx_test = np.arange(num_images)[:: self.config.val_skip] + idx_train = np.array([i for i in np.arange(num_images) if i not in idx_test]) + idx = idx_train if split == "train" else idx_test + if num_images != poses.shape[0]: + raise RuntimeError(f"Different number of images ({num_images}), and poses ({poses.shape[0]})") + + image_filenames = image_filenames[idx] + poses = poses[idx] + + # Centering poses + poses[:, :3, 3] = poses[:, :3, 3] - np.mean(poses[:, :3, 3], axis=0) + + camera_to_world = torch.from_numpy(poses[:, :3, :4]) # camera to world transform + + # Camera intrinsics + K = np.array(metadata_dict["K"]).reshape((3, 3)).T + focal_length = K[0, 0] + + H = metadata_dict["h"] + W = metadata_dict["w"] + + # TODO(akristoffersen): The metadata dict comes with principle points, + # but caused errors in image coord indexing. Should update once that is fixed. + cx, cy = W / 2, H / 2 + + num_cameras = len(image_filenames) + num_intrinsics_params = 3 + intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32) + intrinsics *= torch.tensor([cx, cy, focal_length]) + + # scene_bounds = SceneBounds.from_camera_poses(camera_to_world, self.config.aabb_scale) + aabb = torch.tensor([[-1, -1, -1], [1, 1, 1]], dtype=torch.float32) * self.config.aabb_scale + scene_bounds = SceneBounds(aabb=aabb) + + cameras = Cameras( + fx=focal_length, + fy=focal_length, + cx=cx, + cy=cy, + camera_to_worlds=camera_to_world, + camera_type=CameraType.PERSPECTIVE, + ) + + dataset_inputs = DatasetInputs( + image_filenames=image_filenames, + cameras=cameras, + scene_bounds=scene_bounds, + ) + + return dataset_inputs diff --git a/nerfactory/dataloaders/image_dataset.py b/nerfactory/datamanagers/datasets.py similarity index 64% rename from nerfactory/dataloaders/image_dataset.py rename to nerfactory/datamanagers/datasets.py index c266ffa0f..6f752b078 100644 --- a/nerfactory/dataloaders/image_dataset.py +++ b/nerfactory/datamanagers/datasets.py @@ -15,10 +15,10 @@ """ Some dataset code. """ +from __future__ import annotations from abc import abstractmethod -from pathlib import Path -from typing import Dict, List, Optional, Union +from typing import Dict, Union import numpy as np import numpy.typing as npt @@ -27,18 +27,15 @@ from PIL import Image from torch.utils.data import Dataset from torchtyping import TensorType -from nerfactory.dataloaders.structs import Semantics +from nerfactory.configs import base as cfg +from nerfactory.datamanagers.structs import DatasetInputs from nerfactory.utils.misc import is_not_none -class ImageDataset(Dataset): +class InputDataset(Dataset): """Dataset that returns images.""" - def __init__( - self, - image_filenames: List[Path], - alpha_color: Optional[TensorType[3]] = None, - ): + def __init__(self, config: Union[cfg.DataParserConfig, cfg.InstantiateConfig], split: str): """_summary_ Args: @@ -46,11 +43,10 @@ class ImageDataset(Dataset): alpha_color (TensorType[3], optional): Sets transparent regions to specified color, otherwise black. """ super().__init__() - self.image_filenames = image_filenames - self.alpha_color = alpha_color + self.inputs: DatasetInputs = config.setup().get_dataset_inputs(split=split) def __len__(self): - return len(self.image_filenames) + return len(self.inputs.image_filenames) def get_numpy_image(self, image_idx: int) -> npt.NDArray[np.uint8]: """Returns the image. @@ -61,7 +57,7 @@ class ImageDataset(Dataset): Returns: np.uint8: an image of shape (H, W, 3 or 4) """ - image_filename = self.image_filenames[image_idx] + image_filename = self.inputs.image_filenames[image_idx] pil_image = Image.open(image_filename) image = np.array(pil_image, dtype="uint8") # shape is (h, w, 3 or 4) assert len(image.shape) == 3 @@ -72,9 +68,9 @@ class ImageDataset(Dataset): def get_image(self, image_idx: int) -> TensorType["image_height", "image_width", "num_channels"]: """Returns a 3 channel image.""" image = torch.from_numpy(self.get_numpy_image(image_idx).astype("float32") / 255.0) - if self.alpha_color is not None and image.shape[-1] == 4: + if self.inputs.alpha_color is not None and image.shape[-1] == 4: assert image.shape[-1] == 4 - image = image[:, :, :3] * image[:, :, -1:] + self.alpha_color * (1.0 - image[:, :, -1:]) + image = image[:, :, :3] * image[:, :, -1:] + self.inputs.alpha_color * (1.0 - image[:, :, -1:]) else: image = image[:, :, :3] return image @@ -82,11 +78,23 @@ class ImageDataset(Dataset): @abstractmethod def get_mask(self, image_idx: int) -> Union[TensorType["image_height", "image_width", 1], None]: """Returns a mask, which indicates which pixels are valid to use with nerf.""" + if self.inputs.semantics: + person_index = self.inputs.semantics.thing_classes.index("person") + thing_image_filename = self.inputs.semantics.thing_filenames[image_idx] + pil_image = Image.open(thing_image_filename) + thing_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] + mask = (thing_semantics != person_index).to(torch.float32) # 1 where valid + return mask return None @abstractmethod def get_semantics(self, image_idx: int) -> Union[TensorType["image_height", "image_width", "num_classes"], None]: """Returns an image with semantic class values.""" + if self.inputs.semantics: + stuff_image_filename = self.inputs.semantics.stuff_filenames[image_idx] + pil_image = Image.open(stuff_image_filename) + stuff_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] + return stuff_semantics return None def get_data(self, image_idx) -> Dict: @@ -108,32 +116,3 @@ class ImageDataset(Dataset): def __getitem__(self, image_idx): data = self.get_data(image_idx) return data - - -class PanopticImageDataset(ImageDataset): - """Panoptic image dataset that masks out people.""" - - def __init__( - self, - semantics: Semantics, - image_filenames: List[Path], - alpha_color: Optional[TensorType[3]] = None, - **kwargs, - ): - self.semantics = semantics - self.person_index = self.semantics.thing_classes.index("person") - super().__init__(image_filenames, alpha_color, **kwargs) - - def get_mask(self, image_idx): - """Mask out the people. Valid only where there aren't people.""" - thing_image_filename = self.semantics.thing_filenames[image_idx] - pil_image = Image.open(thing_image_filename) - thing_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] - mask = (thing_semantics != self.person_index).to(torch.float32) # 1 where valid - return mask - - def get_semantics(self, image_idx): - stuff_image_filename = self.semantics.stuff_filenames[image_idx] - pil_image = Image.open(stuff_image_filename) - stuff_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] - return stuff_semantics diff --git a/nerfactory/dataloaders/pixel_sampler.py b/nerfactory/datamanagers/pixel_sampler.py similarity index 100% rename from nerfactory/dataloaders/pixel_sampler.py rename to nerfactory/datamanagers/pixel_sampler.py diff --git a/nerfactory/dataloaders/structs.py b/nerfactory/datamanagers/structs.py similarity index 100% rename from nerfactory/dataloaders/structs.py rename to nerfactory/datamanagers/structs.py diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 458ae51b8..822948b06 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -114,13 +114,9 @@ class Trainer: def train(self) -> None: """Train the model.""" - assert self.pipeline.data_manager.train_datasetinputs is not None, "Missing DatasetInputs" - assert self.pipeline.data_manager.train_image_dataset is not None, "Missing ImageDataset" + assert self.pipeline.datamanager.train_input_dataset is not None, "Missing DatsetInputs" - self.visualizer_state.init_scene( - image_dataset=self.pipeline.data_manager.train_image_dataset, - dataset_inputs=self.pipeline.data_manager.train_datasetinputs, - ) + self.visualizer_state.init_scene(dataset=self.pipeline.datamanager.train_input_dataset) with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME): num_iterations = self.config.trainer.max_num_iterations for step in range(self.start_step, self.start_step + num_iterations): @@ -143,7 +139,7 @@ class Trainer: with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _: self.visualizer_state.update_scene(step, self.pipeline.model) - train_num_rays_per_batch = self.config.pipeline.data_manager.train_num_rays_per_batch + train_num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch writer.put_scalar(name=EventName.RAYS_PER_SEC, scalar=train_num_rays_per_batch / t.duration, step=step) if step != 0 and step % self.config.logging.steps_per_log == 0: diff --git a/nerfactory/fields/instant_ngp_field.py b/nerfactory/fields/instant_ngp_field.py index 553cb1584..808edbad0 100644 --- a/nerfactory/fields/instant_ngp_field.py +++ b/nerfactory/fields/instant_ngp_field.py @@ -24,7 +24,7 @@ from torch.nn.parameter import Parameter from torchtyping import TensorType from nerfactory.cameras.rays import RaySamples -from nerfactory.dataloaders.structs import SceneBounds +from nerfactory.datamanagers.structs import SceneBounds from nerfactory.fields.base import Field from nerfactory.fields.modules.encoding import Encoding, HashEncoding, SHEncoding from nerfactory.fields.modules.field_heads import FieldHeadNames diff --git a/nerfactory/models/base.py b/nerfactory/models/base.py index 402c896e7..e5207d88a 100644 --- a/nerfactory/models/base.py +++ b/nerfactory/models/base.py @@ -28,7 +28,7 @@ from torch.nn import Parameter from nerfactory.cameras.rays import RayBundle from nerfactory.configs import base as cfg -from nerfactory.dataloaders.structs import SceneBounds +from nerfactory.datamanagers.structs import SceneBounds from nerfactory.fields.density_fields.density_grid import DensityGrid from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfactory.utils.misc import get_masked_dict, is_not_none diff --git a/nerfactory/models/modules/scene_colliders.py b/nerfactory/models/modules/scene_colliders.py index df3f550fe..8c316f470 100644 --- a/nerfactory/models/modules/scene_colliders.py +++ b/nerfactory/models/modules/scene_colliders.py @@ -23,7 +23,7 @@ from torch import nn from torchtyping import TensorType from nerfactory.cameras.rays import RayBundle -from nerfactory.dataloaders.structs import SceneBounds +from nerfactory.datamanagers.structs import SceneBounds class SceneBoundsCollider(nn.Module): diff --git a/nerfactory/models/semantic_nerf.py b/nerfactory/models/semantic_nerf.py index 8a84363ab..09a1d0eed 100644 --- a/nerfactory/models/semantic_nerf.py +++ b/nerfactory/models/semantic_nerf.py @@ -22,7 +22,7 @@ from torch import nn from nerfactory.cameras.rays import RayBundle from nerfactory.configs import base as cfg -from nerfactory.dataloaders.structs import Semantics +from nerfactory.datamanagers.structs import Semantics from nerfactory.fields.modules.encoding import NeRFEncoding from nerfactory.fields.modules.field_heads import FieldHeadNames from nerfactory.fields.nerf_field import NeRFField diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 038866049..1334195f7 100644 --- a/nerfactory/pipelines/base.py +++ b/nerfactory/pipelines/base.py @@ -27,7 +27,7 @@ from torch.nn import Parameter from torch.nn.parallel import DistributedDataParallel as DDP from nerfactory.configs import base as cfg -from nerfactory.dataloaders.base import DataManager +from nerfactory.datamanagers.base import DataManager from nerfactory.models.base import Model from nerfactory.utils import profiler from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes @@ -66,11 +66,11 @@ class Pipeline(nn.Module): Args: model: The model to be used in the pipeline. - data_manager: The data_manager to be used in the pipeline. + datamanager: The datamanager to be used in the pipeline. loss_coefficients: A dictionary of loss coefficients that will be used Attributes: - self.data_manager (DataManager): The data manager that will be used + self.datamanager (DataManager): The data manager that will be used self.model (Model): The model that will be used """ @@ -78,13 +78,13 @@ class Pipeline(nn.Module): self, config: cfg.PipelineConfig, device: str, test_mode: bool = False, world_size: int = 1, local_rank: int = 0 ): super().__init__() - self.data_manager: DataManager = config.data_manager.setup( + self.datamanager: DataManager = config.datamanager.setup( device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank ) - self.data_manager.to(device) + self.datamanager.to(device) # TODO(ethan): get rid of scene_bounds from the model - assert self.data_manager.train_datasetinputs is not None, "Missing DatasetInputs" - self.model: Model = config.model.setup(scene_bounds=self.data_manager.train_datasetinputs.scene_bounds) + assert self.datamanager.train_input_dataset is not None, "Missing input dataset" + self.model: Model = config.model.setup(scene_bounds=self.datamanager.train_input_dataset.inputs.scene_bounds) self.model.to(device) self.world_size = world_size @@ -106,8 +106,8 @@ class Pipeline(nn.Module): getting the next batch of data from the DataManager and interfacing with the Model class, feeding the data to the model's forward function.""" if self.world_size > 1: - self.data_manager.sampler.set_epoch(step) - ray_bundle, batch = self.data_manager.next_train() + self.datamanager.sampler.set_epoch(step) + ray_bundle, batch = self.datamanager.next_train() model_outputs, loss_dict, metrics_dict = self.model(ray_bundle, batch) return model_outputs, loss_dict, metrics_dict @@ -118,8 +118,8 @@ class Pipeline(nn.Module): from the DataManager and feed it to the model's forward function""" self.eval() # NOTE(ethan): next_eval() is not being used right now - assert self.data_manager.eval_dataloader is not None - for camera_ray_bundle, batch in self.data_manager.eval_dataloader: + assert self.datamanager.eval_dataloader is not None + for camera_ray_bundle, batch in self.datamanager.eval_dataloader: assert camera_ray_bundle.camera_indices is not None image_idx = int(camera_ray_bundle.camera_indices[0, 0]) if self.world_size > 1: @@ -144,12 +144,12 @@ class Pipeline(nn.Module): self, training_callback_attributes: TrainingCallbackAttributes ) -> List[TrainingCallback]: """Returns the training callbacks from both the Dataloader and the Model.""" - data_manager_callbacks = self.data_manager.get_training_callbacks(training_callback_attributes) + datamanager_callbacks = self.datamanager.get_training_callbacks(training_callback_attributes) if self.world_size > 1: model_callbacks = self.model.module.get_training_callbacks(training_callback_attributes) else: model_callbacks = self.model.get_training_callbacks(training_callback_attributes) - callbacks = data_manager_callbacks + model_callbacks + callbacks = datamanager_callbacks + model_callbacks return callbacks def get_param_groups(self) -> Dict[str, List[Parameter]]: @@ -158,10 +158,10 @@ class Pipeline(nn.Module): Returns: A list of dictionaries containing the pipeline's param groups. """ - data_manager_params = self.data_manager.get_param_groups() + datamanager_params = self.datamanager.get_param_groups() if self.world_size > 1: model_params = self.model.module.get_param_groups() else: model_params = self.model.get_param_groups() # TODO(ethan): assert that key names don't overlap - return {**data_manager_params, **model_params} + return {**datamanager_params, **model_params} diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 47ee70445..f0d2f0ab0 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -28,8 +28,7 @@ import torch from nerfactory.cameras.cameras import Cameras from nerfactory.cameras.rays import RayBundle from nerfactory.configs import base as cfg -from nerfactory.dataloaders.image_dataset import ImageDataset -from nerfactory.dataloaders.structs import DatasetInputs +from nerfactory.datamanagers.datasets import InputDataset from nerfactory.models.base import Model from nerfactory.utils import profiler, visualization from nerfactory.utils.decorators import check_visualizer_enabled, decorate_all @@ -230,12 +229,11 @@ class VisualizerState: self.outputs_set = False - def init_scene(self, image_dataset: ImageDataset, dataset_inputs: DatasetInputs) -> None: + def init_scene(self, dataset: InputDataset) -> None: """Draw some images and the scene aabb in the viewer. Args: - image_dataset: dataset to render in the scene - dataset_inputs: inputs to the image dataset and ray generator + dataset: dataset to render in the scene """ # clear the current scene @@ -243,15 +241,15 @@ class VisualizerState: self.vis["sceneState/cameras"].delete() # draw the training cameras and images - image_indices = range(len(image_dataset)) + image_indices = range(len(dataset)) for idx in image_indices: - image = image_dataset[idx]["image"] + image = dataset[idx]["image"] bgr = image[..., [2, 1, 0]] - camera_json = dataset_inputs.cameras.to_json(camera_idx=idx, image=bgr, resize_shape=(100, 100)) + camera_json = dataset.inputs.cameras.to_json(camera_idx=idx, image=bgr, resize_shape=(100, 100)) self.vis[f"sceneState/cameras/{idx:06d}"].write(camera_json) # draw the scene bounds (i.e., the bounding box) - json_ = dataset_inputs.scene_bounds.to_json() + json_ = dataset.inputs.scene_bounds.to_json() self.vis["sceneState/sceneBounds"].write(json_) # set the properties of the camera diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh index a28373d79..d77a8c013 100755 --- a/scripts/benchmarking/launch_train_blender.sh +++ b/scripts/benchmarking/launch_train_blender.sh @@ -55,7 +55,7 @@ for dataset in ${DATASETS[@]}; do python scripts/run_train.py \ ${config_name} \ --logging.writer.2.no-enable \ - --pipeline.data_manager.train-dataset.data-directory=data/blender/${dataset} \ + --pipeline.datamanager.train-dataset.data-directory=data/blender/${dataset} \ --experiment-name=blender_${dataset}_${tag} \ --trainer.model-dir=nerfactory_models/ \ --trainer.steps-per-save=16000 \ diff --git a/scripts/run_eval.py b/scripts/run_eval.py index bb59b4fbc..498b47d20 100644 --- a/scripts/run_eval.py +++ b/scripts/run_eval.py @@ -79,7 +79,7 @@ def _render_stats_dict(pipeline: Pipeline) -> Dict[str, float]: avg_psnr = 0 avg_rays_per_sec = 0 avg_fps = 0 - for step, (camera_ray_bundle, batch) in tqdm(enumerate(pipeline.data_manager.eval_dataloader)): + for step, (camera_ray_bundle, batch) in tqdm(enumerate(pipeline.datamanager.eval_dataloader)): with TimeWriter(writer=None, name=None, write=False) as t: with torch.no_grad(): image_idx = int(camera_ray_bundle.camera_indices[0, 0]) @@ -142,7 +142,7 @@ def _eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]: # load checkpoints from wherever they were saved # TODO: expose the ability to choose an arbitrary checkpoint config.trainer.load_dir = config.trainer.model_dir - config.pipeline.data_manager.eval_image_indices = None + config.pipeline.datamanager.eval_image_indices = None # setup pipeline (which includes the DataManager) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -210,11 +210,11 @@ class RenderTrajectory: # TODO(ethan): use camera information from parsing args if self.traj == "spiral": - camera_start = pipeline.data_manager.eval_dataloader.get_camera(image_idx=0) + camera_start = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0) # TODO(ethan): pass in the up direction of the camera camera_path = get_spiral_path(camera_start, steps=30, radius=0.1) elif self.traj == "interp": - cameras = pipeline.data_manager.eval_dataloader.get_camera(image_idx=[0, 10]) + cameras = pipeline.datamanager.eval_dataloader.get_camera(image_idx=[0, 10]) camera_path = get_interpolated_camera_path(cameras, steps=30) else: assert_never(self.traj) @@ -225,7 +225,7 @@ class RenderTrajectory: output_filename=self.output_path, rendered_output_name=self.rendered_output_name, rendered_resolution_scaling_factor=self.rendered_resolution_scaling_factor, - num_rays_per_chunk=config.pipeline.data_manager.eval_num_rays_per_chunk, + num_rays_per_chunk=config.pipeline.datamanager.eval_num_rays_per_chunk, ) diff --git a/tests/test_run_train.py b/tests/test_run_train.py index 0e2a5774d..637c14313 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -21,15 +21,14 @@ def set_reduced_config(config: cfg.Config): config.machine.num_gpus = 0 config.trainer.max_num_iterations = 2 # reduce dataset factors; set dataset to test - # switch to using the vanilla ImageDataset class - config.pipeline.data_manager.image_dataset_type = "rgb" - - config.pipeline.data_manager.train_dataset = cfg.BlenderDataParserConfig( + config.pipeline.datamanager.train_dataparser = cfg.BlenderDataParserConfig( + data_directory=Path("tests/data/lego_test") + ) + config.pipeline.datamanager.train_num_images_to_sample_from = 1 + config.pipeline.datamanager.train_num_rays_per_batch = 4 + config.pipeline.datamanager.eval_dataparser = cfg.BlenderDataParserConfig( data_directory=Path("tests/data/lego_test") ) - config.pipeline.data_manager.train_num_images_to_sample_from = 1 - config.pipeline.data_manager.train_num_rays_per_batch = 4 - config.pipeline.data_manager.eval_dataset = cfg.BlenderDataParserConfig(data_directory=Path("tests/data/lego_test")) # reduce model factors config.pipeline.model.num_coarse_samples = 4