realitycapture nerfacto variable resolution (#1651)

* allow periods within filename, not only for file extension

* don't assume all images are the same resolution

* black reformatting

* Move variable resolution collate logic into vanilla datamanager

* good catch pylint

* now that vanilla supports variable resolution, remove VariableResDataManager

* remove comment

* CONSOLE.print

* avoid `IndexError: tensors used as indices must be long, byte or bool tensors`

* use large dataset parameters for phototourism

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),
     ),

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:

nerfstudio/data/datamanagers/variable_res_datamanager.py

@@ -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)

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

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