From dba20fcb26d9256c4267b578f39b3614a6ffe669 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Tue, 10 Oct 2023 17:18:16 -0700 Subject: [PATCH] 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 | 11 +++++++++-- nerfstudio/scripts/exporter.py | 5 +++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/nerfstudio/data/datamanagers/parallel_datamanager.py b/nerfstudio/data/datamanagers/parallel_datamanager.py index 8d541d79e..7a2e5ac55 100644 --- a/nerfstudio/data/datamanagers/parallel_datamanager.py +++ b/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() diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index 2e4fc68ca..b104597ab 100644 --- a/nerfstudio/scripts/exporter.py +++ b/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