diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 27cad2389..37520dbba 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -58,12 +58,14 @@ class CameraOptimizer(nn.Module): config: CameraOptimizerConfig, num_cameras: int, device: Union[torch.device, str], + non_trainable_camera_indices: Optional[Int[Tensor, "num_non_trainable_cameras"]] = None, **kwargs, ) -> None: super().__init__() self.config = config self.num_cameras = num_cameras self.device = device + self.non_trainable_camera_indices = non_trainable_camera_indices # Initialize learnable parameters. if self.config.mode == "off": @@ -95,6 +97,11 @@ class CameraOptimizer(nn.Module): outputs.append(exp_map_SE3(self.pose_adjustment[indices, :])) else: assert_never(self.config.mode) + # Detach non-trainable indices by setting to identity transform + if self.non_trainable_camera_indices is not None: + if self.non_trainable_camera_indices.device != self.pose_adjustment.device: + self.non_trainable_camera_indices = self.non_trainable_camera_indices.to(self.pose_adjustment.device) + outputs[0][self.non_trainable_camera_indices] = torch.eye(4, device=self.pose_adjustment.device)[:3, :4] # Return: identity if no transforms are needed, otherwise multiply transforms together. if len(outputs) == 0: @@ -114,4 +121,7 @@ class CameraOptimizer(nn.Module): """ Add a regularizer """ - loss_dict['camera_opt_regularizer'] = self.pose_adjustment[:,:3].norm(dim=-1).mean() * self.config.trans_l2_penalty + self.pose_adjustment[:,3:].norm(dim=-1).mean() * self.config.rot_l2_penalty \ No newline at end of file + loss_dict["camera_opt_regularizer"] = ( + self.pose_adjustment[:, :3].norm(dim=-1).mean() * self.config.trans_l2_penalty + + self.pose_adjustment[:, 3:].norm(dim=-1).mean() * self.config.rot_l2_penalty + )