diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 0ed2473f9..77b7e2b47 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -20,7 +20,7 @@ from __future__ import annotations import functools from dataclasses import dataclass, field -from typing import Literal, Type, Union +from typing import Literal, Optional, Type, Union import torch import tyro @@ -74,12 +74,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": @@ -121,11 +123,13 @@ 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: + outputs[0][self.non_trainable_camera_indices] = torch.eye(4, device=self.device)[:3, :4] # Apply initial pose noise. if self.pose_noise is not None: outputs.append(self.pose_noise[indices, :, :]) - # Return: identity if no transforms are needed, otherwise multiply transforms together. if len(outputs) == 0: # Note that using repeat() instead of tile() here would result in unnecessary copies.