minor patch for exporting pcs and meshes with parallel datamanager (#2503)

* minor patch for exporting pcs and meshes with parallel datamanager

* format

nerfstudio/data/datamanagers/parallel_datamanager.py

@@ -222,7 +222,7 @@ class ParallelDataManager(DataManager, Generic[TDataset]):
     def setup_train(self):
         """Sets up parallel python data processes for training."""
         assert self.train_dataset is not None
-        self.train_pix_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch)  # type: ignore
+        self.train_pixel_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch)  # type: ignore
         self.data_queue = mp.Manager().Queue(maxsize=self.config.queue_size)
         self.data_procs = [
             DataProcessor(
@@ -230,7 +230,7 @@ class ParallelDataManager(DataManager, Generic[TDataset]):
                 config=self.config,
                 dataparser_outputs=self.train_dataparser_outputs,
                 dataset=self.train_dataset,
-                pixel_sampler=self.train_pix_sampler,
+                pixel_sampler=self.train_pixel_sampler,
             )
             for i in range(self.config.num_processes)
         ]
@@ -319,3 +319,10 @@ class ParallelDataManager(DataManager, Generic[TDataset]):
             A list of dictionaries containing the data manager's param groups.
         """
         return {}
+
+    def __del__(self):
+        """Clean up the parallel data processes."""
+        if hasattr(self, "data_procs"):
+            for proc in self.data_procs:
+                proc.terminate()
+                proc.join()

nerfstudio/scripts/exporter.py

@@ -34,6 +34,7 @@ from typing_extensions import Annotated, Literal
 
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
+from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManager
 from nerfstudio.data.scene_box import OrientedBox
 from nerfstudio.exporter import texture_utils, tsdf_utils
 from nerfstudio.exporter.exporter_utils import (
@@ -137,7 +138,7 @@ class ExportPointCloud(Exporter):
         validate_pipeline(self.normal_method, self.normal_output_name, pipeline)
 
         # Increase the batchsize to speed up the evaluation.
-        assert isinstance(pipeline.datamanager, VanillaDataManager)
+        assert isinstance(pipeline.datamanager, (VanillaDataManager, ParallelDataManager))
         assert pipeline.datamanager.train_pixel_sampler is not None
         pipeline.datamanager.train_pixel_sampler.num_rays_per_batch = self.num_rays_per_batch
 
@@ -307,7 +308,7 @@ class ExportPoissonMesh(Exporter):
         validate_pipeline(self.normal_method, self.normal_output_name, pipeline)
 
         # Increase the batchsize to speed up the evaluation.
-        assert isinstance(pipeline.datamanager, VanillaDataManager)
+        assert isinstance(pipeline.datamanager, (VanillaDataManager, ParallelDataManager))
         assert pipeline.datamanager.train_pixel_sampler is not None
         pipeline.datamanager.train_pixel_sampler.num_rays_per_batch = self.num_rays_per_batch