diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 04d02b679..c4e71e23f 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -70,8 +70,8 @@ from nerfstudio.data.utils.dataloaders import ( from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfstudio.model_components.ray_generators import RayGenerator -from nerfstudio.utils.misc import IterableWrapper from nerfstudio.utils import profiler +from nerfstudio.utils.misc import IterableWrapper CONSOLE = Console(width=120) @@ -469,9 +469,12 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ) self.iter_train_image_dataloader = iter(self.train_image_dataloader) self.train_pixel_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch) - self.train_camera_optimizer = self.config.camera_optimizer.setup( - num_cameras=self.train_dataset.cameras.size, device=self.device - ) + if self.config.camera_optimizer.mode != "off": + self.train_camera_optimizer = self.config.camera_optimizer.setup( + num_cameras=self.train_dataset.cameras.size, device=self.device + ) + else: + self.train_camera_optimizer = None self.train_ray_generator = RayGenerator( self.train_dataset.cameras.to(self.device), self.train_camera_optimizer, @@ -492,9 +495,12 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ) self.iter_eval_image_dataloader = iter(self.eval_image_dataloader) self.eval_pixel_sampler = self._get_pixel_sampler(self.eval_dataset, self.config.eval_num_rays_per_batch) - self.eval_camera_optimizer = self.config.camera_optimizer.setup( - num_cameras=self.eval_dataset.cameras.size, device=self.device - ) + if self.config.camera_optimizer.mode != "off": + self.eval_camera_optimizer = self.config.camera_optimizer.setup( + num_cameras=self.eval_dataset.cameras.size, device=self.device + ) + else: + self.eval_camera_optimizer = None self.eval_ray_generator = RayGenerator( self.eval_dataset.cameras.to(self.device), self.eval_camera_optimizer, @@ -556,11 +562,9 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method """ param_groups = {} - camera_opt_params = list(self.train_camera_optimizer.parameters()) if self.config.camera_optimizer.mode != "off": + camera_opt_params = list(self.train_camera_optimizer.parameters()) assert len(camera_opt_params) > 0 param_groups[self.config.camera_optimizer.param_group] = camera_opt_params - else: - assert len(camera_opt_params) == 0 return param_groups diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 1d5a9697e..261a3795e 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -74,6 +74,7 @@ class PixelSampler: # pylint: disable=too-few-public-methods return indices + @profile def collate_image_dataset_batch(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): """ Operates on a batch of images and samples pixels to use for generating rays. @@ -112,6 +113,7 @@ class PixelSampler: # pylint: disable=too-few-public-methods return collated_batch + @profile def collate_image_dataset_batch_list(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): """ Does the same as collate_image_dataset_batch, except it will operate over a list of images / masks inside @@ -184,6 +186,7 @@ class PixelSampler: # pylint: disable=too-few-public-methods return collated_batch @profiler.time_function + @profile def sample(self, image_batch: Dict): """Sample an image batch and return a pixel batch. diff --git a/nerfstudio/model_components/ray_generators.py b/nerfstudio/model_components/ray_generators.py index 2e568ee3c..8d7170f45 100644 --- a/nerfstudio/model_components/ray_generators.py +++ b/nerfstudio/model_components/ray_generators.py @@ -15,6 +15,9 @@ """ Ray generator. """ +from typing import Optional + +import torch from torch import nn from torchtyping import TensorType @@ -33,7 +36,11 @@ class RayGenerator(nn.Module): pose_optimizer: pose optimization module, for optimizing noisy camera intrinsics/extrinsics. """ - def __init__(self, cameras: Cameras, pose_optimizer: CameraOptimizer) -> None: + image_coords: TensorType["height", "width", 2] + + def __init__( + self, cameras: Cameras, pose_optimizer: Optional[CameraOptimizer] = None + ) -> None: super().__init__() self.cameras = cameras self.pose_optimizer = pose_optimizer @@ -52,7 +59,10 @@ class RayGenerator(nn.Module): x = ray_indices[:, 2] # col indices coords = self.image_coords[y, x] - camera_opt_to_camera = self.pose_optimizer(c) + if self.pose_optimizer is not None: + camera_opt_to_camera = self.pose_optimizer(c) + else: + camera_opt_to_camera = None ray_bundle = self.cameras.generate_rays( camera_indices=c.unsqueeze(-1),