Support downsample in dataset (#973)

* init

* tests passing

* pass through arguments

* verbose

* remove debug prints

* switch to bilinear

* update viewer

* make parser outputs private

* change name

nerfstudio/data/datamanagers/base_datamanager.py

@@ -266,6 +266,10 @@ class VanillaDataManagerConfig(InstantiateConfig):
     camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
     """Specifies the camera pose optimizer used during training. Helpful if poses are noisy, such as for data from
     Record3D."""
+    camera_res_scale_factor: float = 1.0
+    """The scale factor for scaling spatial data such as images, mask, semantics
+    along with relevant information about camera intrinsics
+    """
 
 
 class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
@@ -301,6 +305,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         self.sampler = None
         self.test_mode = test_mode
         self.test_split = "test" if test_mode in ["test", "inference"] else "val"
+        self.dataparser = self.config.dataparser.setup()
 
         self.train_dataset = self.create_train_dataset()
         self.eval_dataset = self.create_eval_dataset()
@@ -308,11 +313,17 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
 
     def create_train_dataset(self) -> InputDataset:
         """Sets up the data loaders for training"""
-        return InputDataset(self.config.dataparser.setup().get_dataparser_outputs(split="train"))
+        return InputDataset(
+            dataparser_outputs=self.dataparser.get_dataparser_outputs(split="train"),
+            scale_factor=self.config.camera_res_scale_factor,
+        )
 
     def create_eval_dataset(self) -> InputDataset:
         """Sets up the data loaders for evaluation"""
-        return InputDataset(self.config.dataparser.setup().get_dataparser_outputs(split=self.test_split))
+        return InputDataset(
+            dataparser_outputs=self.dataparser.get_dataparser_outputs(split=self.test_split),
+            scale_factor=self.config.camera_res_scale_factor,
+        )
 
     def setup_train(self):
         """Sets up the data loaders for training"""
@@ -329,10 +340,10 @@ 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_camera_optimizer = self.config.camera_optimizer.setup(
-            num_cameras=self.train_dataset.dataparser_outputs.cameras.size, device=self.device
+            num_cameras=self.train_dataset.cameras.size, device=self.device
         )
         self.train_ray_generator = RayGenerator(
-            self.train_dataset.dataparser_outputs.cameras.to(self.device),
+            self.train_dataset.cameras.to(self.device),
             self.train_camera_optimizer,
         )
 
@@ -351,7 +362,7 @@ 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_dataset.dataparser_outputs.cameras.to(self.device),
+            self.eval_dataset.cameras.to(self.device),
             self.train_camera_optimizer,  # should be shared between train and eval.
         )
         # for loading full images

nerfstudio/data/datamanagers/semantic_datamanager.py

@@ -41,7 +41,13 @@ class SemanticDataManager(VanillaDataManager):  # pylint: disable=abstract-metho
     """
 
     def create_train_dataset(self) -> SemanticDataset:
-        return SemanticDataset(self.config.dataparser.setup().get_dataparser_outputs(split="train"))
+        return SemanticDataset(
+            dataparser_outputs=self.dataparser.get_dataparser_outputs(split="train"),
+            scale_factor=self.config.camera_res_scale_factor,
+        )
 
     def create_eval_dataset(self) -> SemanticDataset:
-        return SemanticDataset(self.config.dataparser.setup().get_dataparser_outputs(split=self.test_split))
+        return SemanticDataset(
+            dataparser_outputs=self.dataparser.get_dataparser_outputs(split=self.test_split),
+            scale_factor=self.config.camera_res_scale_factor,
+        )

nerfstudio/data/datasets/base_dataset.py

@@ -17,6 +17,7 @@ Dataset.
 """
 from __future__ import annotations
 
+from copy import deepcopy
 from typing import Dict
 
 import numpy as np
@@ -35,15 +36,21 @@ class InputDataset(Dataset):
 
     Args:
         dataparser_outputs: description of where and how to read input images.
+        scale_factor: The scaling factor for the dataparser outputs
     """
 
-    def __init__(self, dataparser_outputs: DataparserOutputs):
+    def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0):
         super().__init__()
-        self.dataparser_outputs = dataparser_outputs
-        self.has_masks = self.dataparser_outputs.mask_filenames is not None
+        self._dataparser_outputs = dataparser_outputs
+        self.has_masks = dataparser_outputs.mask_filenames is not None
+        self.scale_factor = scale_factor
+        self.scene_box = deepcopy(dataparser_outputs.scene_box)
+        self.metadata = deepcopy(dataparser_outputs.metadata)
+        self.cameras = deepcopy(dataparser_outputs.cameras)
+        self.cameras.rescale_output_resolution(scaling_factor=scale_factor)
 
     def __len__(self):
-        return len(self.dataparser_outputs.image_filenames)
+        return len(self._dataparser_outputs.image_filenames)
 
     def get_numpy_image(self, image_idx: int) -> npt.NDArray[np.uint8]:
         """Returns the image of shape (H, W, 3 or 4).
@@ -51,8 +58,12 @@ class InputDataset(Dataset):
         Args:
             image_idx: The image index in the dataset.
         """
-        image_filename = self.dataparser_outputs.image_filenames[image_idx]
+        image_filename = self._dataparser_outputs.image_filenames[image_idx]
         pil_image = Image.open(image_filename)
+        if self.scale_factor != 1.0:
+            width, height = pil_image.size
+            newsize = (int(width * self.scale_factor), int(height * self.scale_factor))
+            pil_image = pil_image.resize(newsize, resample=Image.BILINEAR)
         image = np.array(pil_image, dtype="uint8")  # shape is (h, w, 3 or 4)
         assert len(image.shape) == 3
         assert image.dtype == np.uint8
@@ -66,9 +77,9 @@ class InputDataset(Dataset):
             image_idx: The image index in the dataset.
         """
         image = torch.from_numpy(self.get_numpy_image(image_idx).astype("float32") / 255.0)
-        if self.dataparser_outputs.alpha_color is not None and image.shape[-1] == 4:
+        if self._dataparser_outputs.alpha_color is not None and image.shape[-1] == 4:
             assert image.shape[-1] == 4
-            image = image[:, :, :3] * image[:, :, -1:] + self.dataparser_outputs.alpha_color * (1.0 - image[:, :, -1:])
+            image = image[:, :, :3] * image[:, :, -1:] + self._dataparser_outputs.alpha_color * (1.0 - image[:, :, -1:])
         else:
             image = image[:, :, :3]
         return image
@@ -83,8 +94,8 @@ class InputDataset(Dataset):
         data = {"image_idx": image_idx}
         data["image"] = image
         if self.has_masks:
-            mask_filepath = self.dataparser_outputs.mask_filenames[image_idx]
-            data["mask"] = get_image_mask_tensor_from_path(filepath=mask_filepath)
+            mask_filepath = self._dataparser_outputs.mask_filenames[image_idx]
+            data["mask"] = get_image_mask_tensor_from_path(filepath=mask_filepath, scale_factor=self.scale_factor)
         metadata = self.get_metadata(data)
         data.update(metadata)
         return data

nerfstudio/data/datasets/semantic_dataset.py

@@ -32,12 +32,10 @@ class SemanticDataset(InputDataset):
         dataparser_outputs: description of where and how to read input images.
     """
 
-    def __init__(self, dataparser_outputs: DataparserOutputs):
-        super().__init__(dataparser_outputs)
-        assert "semantics" in dataparser_outputs.metadata.keys() and isinstance(
-            dataparser_outputs.metadata["semantics"], Semantics
-        )
-        self.semantics = dataparser_outputs.metadata["semantics"]
+    def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0):
+        super().__init__(dataparser_outputs, scale_factor)
+        assert "semantics" in dataparser_outputs.metadata.keys() and isinstance(self.metadata["semantics"], Semantics)
+        self.semantics = self.metadata["semantics"]
         self.mask_indices = torch.tensor(
             [self.semantics.classes.index(mask_class) for mask_class in self.semantics.mask_classes]
         ).view(1, 1, -1)
@@ -46,7 +44,7 @@ class SemanticDataset(InputDataset):
         # handle mask
         filepath = self.semantics.filenames[data["image_idx"]]
         semantic_label, mask = get_semantics_and_mask_tensors_from_path(
-            filepath=filepath, mask_indices=self.mask_indices
+            filepath=filepath, mask_indices=self.mask_indices, scale_factor=self.scale_factor
         )
         if "mask" in data.keys():
             mask = mask & data["mask"]

nerfstudio/data/utils/data_utils.py

@@ -21,17 +21,21 @@ import torch
 from PIL import Image
 
 
-def get_image_mask_tensor_from_path(filepath: Path) -> torch.Tensor:
+def get_image_mask_tensor_from_path(filepath: Path, scale_factor: float = 1.0) -> torch.Tensor:
     """
     Utility function to read a mask image from the given path and return a boolean tensor
     """
     pil_mask = Image.open(filepath)
+    if scale_factor != 1.0:
+        width, height = pil_mask.size
+        newsize = (int(width * scale_factor), int(height * scale_factor))
+        pil_mask = pil_mask.resize(newsize, resample=Image.NEAREST)
     mask_tensor = torch.from_numpy(np.array(pil_mask)).unsqueeze(-1).bool()
     return mask_tensor
 
 
 def get_semantics_and_mask_tensors_from_path(
-    filepath: Path, mask_indices: Union[List, torch.Tensor]
+    filepath: Path, mask_indices: Union[List, torch.Tensor], scale_factor: float = 1.0
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Utility function to read segmentation from the given filepath
@@ -40,6 +44,10 @@ def get_semantics_and_mask_tensors_from_path(
     if isinstance(mask_indices, List):
         mask_indices = torch.tensor(mask_indices, dtype="int64").view(1, 1, -1)
     pil_image = Image.open(filepath)
+    if scale_factor != 1.0:
+        width, height = pil_image.size
+        newsize = (int(width * scale_factor), int(height * scale_factor))
+        pil_image = pil_image.resize(newsize, resample=Image.NEAREST)
     semantics = torch.from_numpy(np.array(pil_image, dtype="int64"))[..., None]
     mask = torch.sum(semantics == mask_indices, dim=-1, keepdim=True) == 0
     return semantics, mask

nerfstudio/data/utils/dataloaders.py

@@ -132,7 +132,7 @@ class EvalDataloader(DataLoader):
         **kwargs,
     ):
         self.input_dataset = input_dataset
-        self.cameras = input_dataset.dataparser_outputs.cameras.to(device)
+        self.cameras = input_dataset.cameras.to(device)
         self.device = device
         self.kwargs = kwargs
         super().__init__(dataset=input_dataset)

nerfstudio/pipelines/base_pipeline.py

@@ -226,9 +226,9 @@ class VanillaPipeline(Pipeline):
         assert self.datamanager.train_dataset is not None, "Missing input dataset"
 
         self._model = config.model.setup(
-            scene_box=self.datamanager.train_dataset.dataparser_outputs.scene_box,
+            scene_box=self.datamanager.train_dataset.scene_box,
             num_train_data=len(self.datamanager.train_dataset),
-            metadata=self.datamanager.train_dataset.dataparser_outputs.metadata,
+            metadata=self.datamanager.train_dataset.metadata,
         )
         self.model.to(device)
 

nerfstudio/viewer/server/viewer_utils.py

@@ -310,11 +310,11 @@ class ViewerState:
         for idx in image_indices:
             image = dataset[idx]["image"]
             bgr = image[..., [2, 1, 0]]
-            camera_json = dataset.dataparser_outputs.cameras.to_json(camera_idx=idx, image=bgr, max_size=100)
+            camera_json = dataset.cameras.to_json(camera_idx=idx, image=bgr, max_size=100)
             self.vis[f"sceneState/cameras/{idx:06d}"].write(camera_json)
 
         # draw the scene box (i.e., the bounding box)
-        json_ = dataset.dataparser_outputs.scene_box.to_json()
+        json_ = dataset.scene_box.to_json()
         self.vis["sceneState/sceneBox"].write(json_)
 
         # set the initial state whether to train or not