From e0e66ce9fb4fe834caf1e815b40b66ecc4f30650 Mon Sep 17 00:00:00 2001 From: Matias Turkulainen <30566358+maturk@users.noreply.github.com> Date: Sun, 18 Jun 2023 19:55:13 +0300 Subject: [PATCH] Add non-trainable camera indices to CameraOptimizer (#2080) * add non-trainable indices to camera optimizer * improve comment * create on device --------- Co-authored-by: Matthew Tancik --- nerfstudio/cameras/camera_optimizers.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) 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.