diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 5e1897659..21029e1b0 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -53,7 +53,7 @@ class CameraOptimizerConfig(cfg.InstantiateConfig): optimizer: AdamOptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15) """ADAM parameters for camera optimization.""" - scheduler: SchedulerConfig = SchedulerConfig(max_steps=10000) + scheduler: SchedulerConfig = SchedulerConfig(max_steps=10000, lr_final=1e-5) """Learning rate scheduler for camera optimizer..""" param_group: tyro.conf.Suppress[str] = "camera_opt" @@ -77,6 +77,7 @@ class CameraOptimizer(nn.Module): self.config = config self.num_cameras = num_cameras self.device = device + self.i = 0 # Initialize learnable parameters. if self.config.mode == "off": @@ -134,10 +135,15 @@ class CameraOptimizer(nn.Module): """ Samples the blur parameters and returns transformation matrices for the origins and directions """ + self.i += 1 + # blur_adjs = self.blur_adjustment[indices, :] * torch.rand((indices.shape[0], 1), device=self.device) blur_adjs = self.blur_adjustment[indices, :] * torch.normal( torch.zeros((indices.shape[0], 1), device=self.device), # mean torch.ones((indices.shape[0], 1), device=self.device), # std ) + # if self.i % 500 == 0: + # # print the directions + # print(self.blur_adjustment) # print( # "trans,rot norm", # torch.mean(torch.linalg.norm(blur_adjs[:, :3], dim=0)), diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 377c782aa..cf9f4bfec 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -63,10 +63,13 @@ method_configs["nerfacto"] = Config( pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), - train_num_rays_per_batch=4096, + train_num_rays_per_batch=1500, + num_blur_samples=5, eval_num_rays_per_batch=4096, camera_optimizer=CameraOptimizerConfig( - mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) + mode="SO3xR3", + optimizer=AdamOptimizerConfig(lr=1e-2, eps=1e-8, weight_decay=0), + scheduler=SchedulerConfig(max_steps=10000, lr_final=1e-3), ), ), model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 0e971d0c1..a76354735 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -264,6 +264,7 @@ 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.""" + num_blur_samples: int = 1 class VanillaDataManager(DataManager): # pylint: disable=abstract-method @@ -365,12 +366,19 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self.train_count += 1 image_batch = next(self.iter_train_image_dataloader) batch = self.train_pixel_sampler.sample(image_batch) - ray_indices = batch["indices"] - ray_bundle = self.train_ray_generator(ray_indices) + ray_bundle = self.train_ray_generator(batch["indices"]) + if self.config.num_blur_samples > 1: + # duplicate ray_bundle num_blur_samples times and stack + os = ray_bundle.origins.repeat(self.config.num_blur_samples, 1) + ds = ray_bundle.directions.repeat(self.config.num_blur_samples, 1) + ps = ray_bundle.pixel_area.repeat(self.config.num_blur_samples, 1) + cs = ray_bundle.camera_indices.repeat(self.config.num_blur_samples, 1) + ray_bundle = RayBundle(origins=os, directions=ds, pixel_area=ps, camera_indices=cs) + camera_ids = ray_bundle.camera_indices[:, 0] # Multiply in the camera optimization deltas to the rays - cam_deltas = to4x4(self.train_camera_optimizer(ray_indices[:, 0])) # B x 3 x 4 - blur_deltas = self.train_camera_optimizer.sample_blur_correction(ray_indices[:, 0]) + cam_deltas = to4x4(self.train_camera_optimizer(camera_ids)) # B x 3 x 4 + blur_deltas = self.train_camera_optimizer.sample_blur_correction(camera_ids) # TODO(Justin) is this allowed mathematically cam_deltas = torch.bmm(blur_deltas, cam_deltas) origins_homogeneous = torch.cat( diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 6929c2fda..bbfce5104 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -216,6 +216,11 @@ class NerfactoModel(Model): ray_samples_list.append(ray_samples) rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) + s = 1500 + if rgb.shape[0] == s * 5: + rgb = (1 / 5.0) * ( + rgb[:s, :] + rgb[s : 2 * s, :] + rgb[2 * s : 3 * s, :] + rgb[3 * s : 4 * s, :] + rgb[4 * s : 5 * s, :] + ) depth = self.renderer_depth(weights=weights, ray_samples=ray_samples) accumulation = self.renderer_accumulation(weights=weights) diff --git a/nerfstudio/utils/tensor_dataclass.py b/nerfstudio/utils/tensor_dataclass.py index cd85eac1c..c20fac1ef 100644 --- a/nerfstudio/utils/tensor_dataclass.py +++ b/nerfstudio/utils/tensor_dataclass.py @@ -197,7 +197,7 @@ class TensorDataclass: """ return self._apply_fn_to_fields(lambda x: x.broadcast_to((*shape, x.shape[-1]))) - def to(self: TensorDataclassT, device) -> TensorDataclassT: + def to(self: TensorDataclassT, *args, **kwargs) -> TensorDataclassT: """Returns a new TensorDataclass with the same data but on the specified device. Args: @@ -206,7 +206,18 @@ class TensorDataclass: Returns: A new TensorDataclass with the same data but on the specified device. """ - return self._apply_fn_to_fields(lambda x: x.to(device)) + return self._apply_fn_to_fields(lambda x: x.to(*args, **kwargs)) + + def pin_memory(self: TensorDataclassT) -> TensorDataclassT: + """Returns a new TensorDataclass with the same data but on the specified device. + + Args: + device: The device to place the tensor dataclass. + + Returns: + A new TensorDataclass with the same data but on the specified device. + """ + return self._apply_fn_to_fields(lambda x: x.pin_memory()) def _apply_fn_to_fields( self: TensorDataclassT, fn: Callable, dataclass_fn: Optional[Callable] = None