use gsplat strategy interface to simplify splatfacto

jb-ye · Aug 22, 2024 · d45dec4 · d45dec4
1 parent a281d99
commit d45dec4
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 272 deletions.
diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py
@@ -343,10 +343,11 @@ def next_train(self, step: int) -> Tuple[Cameras, Dict]:
         # Make sure to re-populate the unseen cameras list if we have exhausted it
         if len(self.train_unseen_cameras) == 0:
             self.train_unseen_cameras = self.sample_train_cameras()
-
-        data = self.cached_train[image_idx]
-        data["image"] = data["image"].to(self.device)
-
+        if self.device != self.config.cache_images:
+            data = deepcopy(self.cached_train[image_idx])
+            data["image"] = data["image"].to(self.device)
+        else:
+            data = self.cached_train[image_idx]
         assert len(self.train_cameras.shape) == 1, "Assumes single batch dimension"
         camera = self.train_cameras[image_idx : image_idx + 1].to(self.device)
         if camera.metadata is None: