diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index a895fb7b8..b1ad3fcf0 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/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 diff --git a/nerfstudio/data/datamanagers/semantic_datamanager.py b/nerfstudio/data/datamanagers/semantic_datamanager.py index d19769308..b3cc223ae 100644 --- a/nerfstudio/data/datamanagers/semantic_datamanager.py +++ b/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, + ) diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py index 15d4296d0..1d1a88a3f 100644 --- a/nerfstudio/data/datasets/base_dataset.py +++ b/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 diff --git a/nerfstudio/data/datasets/semantic_dataset.py b/nerfstudio/data/datasets/semantic_dataset.py index f2c2a0910..4c93a5da5 100644 --- a/nerfstudio/data/datasets/semantic_dataset.py +++ b/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"] diff --git a/nerfstudio/data/utils/data_utils.py b/nerfstudio/data/utils/data_utils.py index 953e7951b..cba914532 100644 --- a/nerfstudio/data/utils/data_utils.py +++ b/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 diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index 5eec5e530..a4717dd9c 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/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) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index d5a8b1ec0..d51e2dfbd 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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) diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index e0ea6e36e..040bd61e7 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/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