diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 22635a911..1e5a24ec0 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -29,9 +29,6 @@ from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConf from nerfstudio.data.datamanagers.depth_datamanager import DepthDataManagerConfig from nerfstudio.data.datamanagers.sdf_datamanager import SDFDataManagerConfig from nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig -from nerfstudio.data.datamanagers.variable_res_datamanager import ( - VariableResDataManagerConfig, -) from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig from nerfstudio.data.dataparsers.dycheck_dataparser import DycheckDataParserConfig @@ -358,13 +355,18 @@ method_configs["phototourism"] = TrainerConfig( max_num_iterations=30000, mixed_precision=True, pipeline=VanillaPipelineConfig( - datamanager=VariableResDataManagerConfig( # NOTE: one of the only differences with nerfacto + datamanager=VanillaDataManagerConfig( dataparser=PhototourismDataParserConfig(), # NOTE: one of the only differences with nerfacto train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, camera_optimizer=CameraOptimizerConfig( mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) ), + # Large dataset, so using prior values from VariableResDataManager. + train_num_images_to_sample_from=40, + train_num_times_to_repeat_images=100, + eval_num_images_to_sample_from=40, + eval_num_times_to_repeat_images=100, ), model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), ), diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index dc96e22dc..8bcd21bbe 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -74,6 +74,31 @@ from nerfstudio.utils.misc import IterableWrapper CONSOLE = Console(width=120) + +def variable_res_collate(batch: List[Dict]) -> Dict: + """Default collate function for the cached dataloader. + Args: + batch: Batch of samples from the dataset. + Returns: + Collated batch. + """ + images = [] + masks = [] + for data in batch: + image = data.pop("image") + mask = data.pop("mask", None) + images.append(image) + if mask: + masks.append(mask) + + new_batch: dict = nerfstudio_collate(batch) + new_batch["image"] = images + if masks: + new_batch["mask"] = masks + + return new_batch + + AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of flags in subcommands. tyro.extras.subcommand_type_from_defaults( { @@ -385,6 +410,16 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self.train_dataset = self.create_train_dataset() self.eval_dataset = self.create_eval_dataset() + + if self.train_dataparser_outputs is not None: + cameras = self.train_dataparser_outputs.cameras + if len(cameras) > 1: + for i in range(1, len(cameras)): + if cameras[0].width != cameras[i].width or cameras[0].height != cameras[i].height: + CONSOLE.print("Variable resolution, using variable_res_collate") + self.config.collate_fn = variable_res_collate + break + super().__init__() def create_train_dataset(self) -> InputDataset: diff --git a/nerfstudio/data/datamanagers/variable_res_datamanager.py b/nerfstudio/data/datamanagers/variable_res_datamanager.py deleted file mode 100644 index b821bc034..000000000 --- a/nerfstudio/data/datamanagers/variable_res_datamanager.py +++ /dev/null @@ -1,66 +0,0 @@ -# Copyright 2022 The Nerfstudio 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 loader for variable resolution datasets, where batching raw image tensors isn't possible. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Dict, List - -from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig -from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate - - -def variable_res_collate(batch: List[Dict]) -> Dict: - """Default collate function for the cached dataloader. - Args: - batch: Batch of samples from the dataset. - Returns: - Collated batch. - """ - images = [] - masks = [] - for data in batch: - image = data.pop("image") - mask = data.pop("mask", None) - images.append(image) - if mask: - masks.append(mask) - - new_batch: dict = nerfstudio_collate(batch) - new_batch["image"] = images - if masks: - new_batch["mask"] = masks - - return new_batch - - -@dataclass -class VariableResDataManagerConfig(VanillaDataManagerConfig): - """A datamanager for variable resolution datasets, with presets to optimize - for the fact that we are now dealing with lists of images and masks. - """ - - train_num_images_to_sample_from: int = 40 - """Number of images to sample during training iteration.""" - train_num_times_to_repeat_images: int = 100 - """When not training on all images, number of iterations before picking new - images. If -1, never pick new images.""" - eval_num_images_to_sample_from: int = 40 - """Number of images to sample during eval iteration.""" - eval_num_times_to_repeat_images: int = 100 - collate_fn = staticmethod(variable_res_collate) diff --git a/nerfstudio/data/dataparsers/phototourism_dataparser.py b/nerfstudio/data/dataparsers/phototourism_dataparser.py index e5a88e10e..d470215ed 100644 --- a/nerfstudio/data/dataparsers/phototourism_dataparser.py +++ b/nerfstudio/data/dataparsers/phototourism_dataparser.py @@ -136,8 +136,8 @@ class Phototourism(DataParser): ) # equally spaced training images starting and ending at 0 and num_images-1 i_eval = np.setdiff1d(i_all, i_train) # eval images are the remaining images i_all = torch.tensor(i_all) - i_train = torch.tensor(i_train) - i_eval = torch.tensor(i_eval) + i_train = torch.tensor(i_train, dtype=torch.long) + i_eval = torch.tensor(i_eval, dtype=torch.long) assert len(i_eval) == num_eval_images if split == "train": indices = i_train diff --git a/nerfstudio/process_data/realitycapture_utils.py b/nerfstudio/process_data/realitycapture_utils.py index 87ce9240c..4bc02417e 100644 --- a/nerfstudio/process_data/realitycapture_utils.py +++ b/nerfstudio/process_data/realitycapture_utils.py @@ -59,17 +59,6 @@ def realitycapture_to_json( for column, value in row.items(): cameras.setdefault(column, []).append(value) - for name in cameras["#name"]: - camera_label = name.rpartition(".")[0] - if camera_label in image_filename_map: - img = np.array(Image.open(output_dir / image_filename_map[camera_label])) - break - - height, width, _ = img.shape - - data["h"] = int(height) - data["w"] = int(width) - missing_image_data = 0 for i, name in enumerate(cameras["#name"]): @@ -79,7 +68,12 @@ def realitycapture_to_json( CONSOLE.print(f"Missing image for camera data {basename}, Skipping") missing_image_data += 1 continue + frame = {} + img = np.array(Image.open(output_dir / image_filename_map[basename])) + height, width, _ = img.shape + frame["h"] = int(height) + frame["w"] = int(width) frame["file_path"] = image_filename_map[basename].as_posix() frame["fl_x"] = float(cameras["f"][i]) * max(width, height) / 36 frame["fl_y"] = float(cameras["f"][i]) * max(width, height) / 36