Datamanager fields typing (#609)

nerfstudio/data/datamanagers.py

@@ -25,6 +25,7 @@ import dcargs
 import torch
 from torch import nn
 from torch.nn import Parameter
+from torch.utils.data import Dataset
 from torch.utils.data.distributed import DistributedSampler
 
 from nerfstudio.cameras.rays import RayBundle
@@ -101,16 +102,16 @@ class DataManager(nn.Module):
     Attributes:
         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_input_dataset (InputDataset): the input dataset for the train dataset
-        eval_input_dataset (InputDataset): the input dataset for the eval dataset
+        train_dataset (Dataset): the dataset for the train dataset
+        eval_dataset (Dataset): the dataset for the eval dataset
 
         Additional attributes specific to each subclass are defined in the setup_train and setup_eval
         functions.
 
     """
 
-    train_input_dataset: Optional[InputDataset] = None
-    eval_input_dataset: Optional[InputDataset] = None
+    train_dataset: Optional[Dataset] = None
+    eval_dataset: Optional[Dataset] = None
     train_sampler: Optional[DistributedSampler] = None
     eval_sampler: Optional[DistributedSampler] = None
 
@@ -126,9 +127,9 @@ class DataManager(nn.Module):
         super().__init__()
         self.train_count = 0
         self.eval_count = 0
-        if self.train_input_dataset and self.train_input_dataset.dataset_inputs:
+        if self.train_dataset:
             self.setup_train()
-        if self.eval_input_dataset and self.eval_input_dataset.dataset_inputs:
+        if self.eval_dataset:
             self.setup_eval()
 
     def forward(self):
@@ -262,6 +263,8 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
     """
 
     config: VanillaDataManagerConfig
+    train_dataset: InputDataset
+    eval_dataset: InputDataset
 
     def __init__(
         self,
@@ -278,23 +281,23 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         self.local_rank = local_rank
         self.sampler = None
 
-        self.train_input_dataset = InputDataset(config.dataparser.setup().get_dataset_inputs(split="train"))
-        self.eval_input_dataset = InputDataset(
+        self.train_dataset = InputDataset(config.dataparser.setup().get_dataset_inputs(split="train"))
+        self.eval_dataset = InputDataset(
             config.dataparser.setup().get_dataset_inputs(split="val" if not test_mode else "test")
         )
         super().__init__()
 
     def setup_train(self):
         """Sets up the data loaders for training"""
-        assert self.train_input_dataset is not None
+        assert self.train_dataset is not None
         if self.world_size > 1:
             sampler = DistributedSampler(
-                self.train_input_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42
+                self.train_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42
             )
         else:
             sampler = None
         self.train_image_dataloader = CacheDataloader(
-            self.train_input_dataset,
+            self.train_dataset,
             num_images_to_sample_from=self.config.train_num_images_to_sample_from,
             device=self.device,
             num_workers=self.world_size * 4,
@@ -303,19 +306,19 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         )
         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_input_dataset.dataset_inputs.cameras.to(self.device))
+        self.train_ray_generator = RayGenerator(self.train_dataset.dataset_inputs.cameras.to(self.device))
 
     def setup_eval(self):
         """Sets up the data loader for evaluation"""
-        assert self.eval_input_dataset is not None
+        assert self.eval_dataset is not None
         if self.world_size > 1:
             sampler = DistributedSampler(
-                self.eval_input_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42
+                self.eval_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42
             )
         else:
             sampler = None
         self.eval_image_dataloader = CacheDataloader(
-            self.eval_input_dataset,
+            self.eval_dataset,
             num_images_to_sample_from=self.config.eval_num_images_to_sample_from,
             device=self.device,
             num_workers=self.world_size * 4,
@@ -324,15 +327,15 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         )
         self.iter_eval_image_dataloader = iter(self.eval_image_dataloader)
         self.eval_pixel_sampler = PixelSampler(self.config.eval_num_rays_per_batch)
-        self.eval_ray_generator = RayGenerator(self.eval_input_dataset.dataset_inputs.cameras.to(self.device))
+        self.eval_ray_generator = RayGenerator(self.eval_dataset.dataset_inputs.cameras.to(self.device))
         # for loading full images
         self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader(
-            input_dataset=self.eval_input_dataset,
+            input_dataset=self.eval_dataset,
             device=self.device,
             num_workers=self.world_size * 4,
         )
         self.eval_dataloader = RandIndicesEvalDataloader(
-            input_dataset=self.eval_input_dataset,
+            input_dataset=self.eval_dataset,
             image_indices=self.config.eval_image_indices,
             device=self.device,
             num_workers=self.world_size * 4,

nerfstudio/data/utils/__init__.py

@@ -0,0 +1,13 @@
+# 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.

nerfstudio/engine/trainer.py

@@ -144,7 +144,7 @@ class Trainer:
 
     def train(self) -> None:
         """Train the model."""
-        assert self.pipeline.datamanager.train_input_dataset is not None, "Missing DatsetInputs"
+        assert self.pipeline.datamanager.train_dataset is not None, "Missing DatsetInputs"
 
         self._init_viewer_scene()
         with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME):
@@ -209,9 +209,9 @@ class Trainer:
     @check_viewer_enabled
     def _init_viewer_scene(self) -> None:
         """Initializes viewer scene with given train dataset"""
-        assert self.viewer_state and self.pipeline.datamanager.train_input_dataset
+        assert self.viewer_state and self.pipeline.datamanager.train_dataset
         self.viewer_state.init_scene(
-            dataset=self.pipeline.datamanager.train_input_dataset,
+            dataset=self.pipeline.datamanager.train_dataset,
             start_train=self.config.viewer.start_train,
         )
 

nerfstudio/pipelines/base.py

@@ -211,10 +211,10 @@ class VanillaPipeline(Pipeline):
         )
         self.datamanager.to(device)
         # TODO(ethan): get rid of scene_box from the model
-        assert self.datamanager.train_input_dataset is not None, "Missing input dataset"
+        assert self.datamanager.train_dataset is not None, "Missing input dataset"
         self.model = config.model.setup(
-            scene_box=self.datamanager.train_input_dataset.dataset_inputs.scene_box,
-            num_train_data=len(self.datamanager.train_input_dataset),
+            scene_box=self.datamanager.train_dataset.dataset_inputs.scene_box,
+            num_train_data=len(self.datamanager.train_dataset),
         )
         self.model.to(device)