Merge branch 'dataloading-revamp' of https://github.com/AntonioMacaronio/nerfstudio into dataloading-revamp

Author: AntonioMacaronio

nerfstudio/data/datamanagers/base_datamanager.py

@@ -357,7 +357,7 @@ def __post_init__(self):
             try:
                 torch.multiprocessing.set_start_method("spawn")
             except RuntimeError:
-                pass
+                assert torch.multiprocessing.get_start_method() == "spawn", 'start method must be "spawn"'
             self.dataloader_num_workers = 4 if self.dataloader_num_workers == 0 else self.dataloader_num_workers
 
 

nerfstudio/data/datamanagers/full_images_datamanager.py

@@ -97,7 +97,7 @@ def __post_init__(self):
             try:
                 torch.multiprocessing.set_start_method("spawn")
             except RuntimeError:
-                pass
+                assert torch.multiprocessing.get_start_method() == "spawn", 'start method must be "spawn"'
             if self.prefetch_factor == 0:
                 CONSOLE.log('cache_images set to "disk" with no prefetch factor, defaulting to 4')
                 self.prefetch_factor = 4
@@ -352,7 +352,7 @@ def setup_eval(self):
             self.eval_image_dataloader = DataLoader(
                 self.eval_imagebatch_stream,
                 batch_size=1,
-                num_workers=0,
+                num_workers=0,  # This must be 0 otherwise there is a crash when trying to pickle custom_view_processor
                 collate_fn=identity_collate,
             )
             self.iter_eval_image_dataloader = iter(self.eval_image_dataloader)
@@ -434,6 +434,9 @@ def next_eval_image(self, step: int) -> Tuple[Cameras, Dict]:
         """Returns the next evaluation batch
         Returns a Camera instead of raybundle
         TODO: Make sure this logic is consistent with the vanilladatamanager"""
+        if self.config.cache_images == "disk":
+            camera, data = next(self.iter_eval_image_dataloader)[0]
+            return camera, data
         image_idx = self.eval_unseen_cameras.pop(random.randint(0, len(self.eval_unseen_cameras) - 1))
         # Make sure to re-populate the unseen cameras list if we have exhausted it
         if len(self.eval_unseen_cameras) == 0:

nerfstudio/data/datamanagers/parallel_datamanager.py

@@ -174,7 +174,7 @@ def setup_eval(self):
         self.eval_ray_dataloader = DataLoader(
             self.eval_raybatchstream,
             batch_size=1,
-            num_workers=0,
+            num_workers=0,  # This must be 0 otherwise there is a crash when trying to pickle custom_view_processor
             shuffle=False,
             collate_fn=identity_collate,  # Our dataset handles batching / collation of rays
         )

nerfstudio/viewer/render_panel.py

@@ -30,6 +30,7 @@
 import viser.transforms as tf
 from scipy import interpolate
 
+from nerfstudio.utils.rich_utils import CONSOLE
 from nerfstudio.viewer.control_panel import ControlPanel
 
 
@@ -1130,8 +1131,15 @@ def _(event: viser.GuiEvent) -> None:
                 }
 
         # now write the json file
-        json_outfile = datapath / "camera_paths" / f"{render_name_text.value}.json"
-        json_outfile.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            json_outfile = datapath / "camera_paths" / f"{render_name_text.value}.json"
+            json_outfile.parent.mkdir(parents=True, exist_ok=True)
+        except Exception:
+            CONSOLE.print(
+                "[bold yellow]Warning: Failed to write the camera path to the data directory. Saving to the output directory instead."
+            )
+            json_outfile = config_path.parent / "camera_paths" / f"{render_name_text.value}.json"
+            json_outfile.parent.mkdir(parents=True, exist_ok=True)
         with open(json_outfile.absolute(), "w") as outfile:
             json.dump(json_data, outfile)
         # now show the command