From ce06935f3df92adf027c7421cbc1883c50e7252a Mon Sep 17 00:00:00 2001 From: Ji Shi Date: Mon, 16 Oct 2023 01:05:22 -0500 Subject: [PATCH] Fix datamanagers' get_rays_per_batch functions in case of using DynamicBatchPipeline (#2522) --- nerfstudio/data/datamanagers/base_datamanager.py | 4 ++++ nerfstudio/data/datamanagers/parallel_datamanager.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 0b937f64c..b4cbeaa12 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -556,9 +556,13 @@ class VanillaDataManager(DataManager, Generic[TDataset]): raise ValueError("No more eval images") def get_train_rays_per_batch(self) -> int: + if self.train_pixel_sampler is not None: + return self.train_pixel_sampler.num_rays_per_batch return self.config.train_num_rays_per_batch def get_eval_rays_per_batch(self) -> int: + if self.eval_pixel_sampler is not None: + return self.eval_pixel_sampler.num_rays_per_batch return self.config.eval_num_rays_per_batch def get_datapath(self) -> Path: diff --git a/nerfstudio/data/datamanagers/parallel_datamanager.py b/nerfstudio/data/datamanagers/parallel_datamanager.py index 7a2e5ac55..5c4d5c4a2 100644 --- a/nerfstudio/data/datamanagers/parallel_datamanager.py +++ b/nerfstudio/data/datamanagers/parallel_datamanager.py @@ -303,10 +303,14 @@ class ParallelDataManager(DataManager, Generic[TDataset]): def get_train_rays_per_batch(self) -> int: """Returns the number of rays per batch for training.""" + if self.train_pixel_sampler is not None: + return self.train_pixel_sampler.num_rays_per_batch return self.config.train_num_rays_per_batch def get_eval_rays_per_batch(self) -> int: """Returns the number of rays per batch for evaluation.""" + if self.eval_pixel_sampler is not None: + return self.eval_pixel_sampler.num_rays_per_batch return self.config.eval_num_rays_per_batch def get_datapath(self) -> Path: