Merge branch 'main' into bell-one-main

nerfstudio-project · Aug 23, 2024 · 163f3f7 · 163f3f7
2 parents 510e2a3 + e3a506f
commit 163f3f7
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 8 deletions.
diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py
@@ -332,8 +332,13 @@ def get_param_groups(self) -> Dict[str, List[Parameter]]:
         return {}
 
     def get_train_rays_per_batch(self):
-        # TODO: fix this to be the resolution of the last image rendered
-        return 800 * 800
+        """Returns resolution of the image returned from datamanager."""
+        if len(self.cached_train) != 0:
+            h = self.cached_train[0]["image"].shape[0]
+            w = self.cached_train[0]["image"].shape[1]
+            return h * w
+        else:
+            return 800 * 800
 
     def next_train(self, step: int) -> Tuple[Cameras, Dict]:
         """Returns the next training batch

diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py
@@ -229,10 +229,10 @@ def setup_optimizers(self) -> Optimizers:
     def train(self) -> None:
         """Train the model."""
         assert self.pipeline.datamanager.train_dataset is not None, "Missing DatsetInputs"
-
-        self.pipeline.datamanager.train_dataparser_outputs.save_dataparser_transform(
-            self.base_dir / "dataparser_transforms.json"
-        )
+        if hasattr(self.pipeline.datamanager, "train_dataparser_outputs"):
+            self.pipeline.datamanager.train_dataparser_outputs.save_dataparser_transform(
+                self.base_dir / "dataparser_transforms.json"
+            )
 
         self._init_viewer_state()
         with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME):

diff --git a/nerfstudio/viewer/viewer.py b/nerfstudio/viewer/viewer.py
@@ -458,9 +458,9 @@ def init_scene(
             R = R @ vtf.SO3.from_x_radians(np.pi)
             camera_handle = self.viser_server.scene.add_camera_frustum(
                 name=f"/cameras/camera_{idx:05d}",
-                fov=float(2 * np.arctan(camera.cx / camera.fx[0])),
+                fov=float(2 * np.arctan((camera.cx / camera.fx[0]).cpu())),
                 scale=self.config.camera_frustum_scale,
-                aspect=float(camera.cx[0] / camera.cy[0]),
+                aspect=float((camera.cx[0] / camera.cy[0]).cpu()),
                 image=image_uint8,
                 wxyz=R.wxyz,
                 position=c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO,