From 53037ba8a790a7aa089c8a4983c2aa210e01111c Mon Sep 17 00:00:00 2001 From: Maxim Bonnaerens Date: Sun, 24 Sep 2023 12:06:16 +0200 Subject: [PATCH] Allow custom sampler in ProposalNetworkSampler (#2464) * Allow custom sampler in ProposalNetworkSampler This allows custom (pdf)samplers in distinct repositories while still using the base ProposalNetworkSampler. * sampler should inherit from pdf sampler * docstring --- nerfstudio/model_components/ray_samplers.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index e18c680ae..7a2052b63 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -530,6 +530,7 @@ class ProposalNetworkSampler(Sampler): single_jitter: Use a same random jitter for all samples along a ray. update_sched: A function that takes the iteration number of steps between updates. initial_sampler: Sampler to use for the first iteration. Uses UniformLinDispPiecewise if not set. + pdf_sampler: PDFSampler to use after the first iteration. Uses PDFSampler if not set. """ def __init__( @@ -540,6 +541,7 @@ class ProposalNetworkSampler(Sampler): single_jitter: bool = False, update_sched: Callable = lambda x: 1, initial_sampler: Optional[Sampler] = None, + pdf_sampler: Optional[PDFSampler] = None, ) -> None: super().__init__() self.num_proposal_samples_per_ray = num_proposal_samples_per_ray @@ -554,7 +556,10 @@ class ProposalNetworkSampler(Sampler): self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter) else: self.initial_sampler = initial_sampler - self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter) + if pdf_sampler is None: + self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter) + else: + self.pdf_sampler = pdf_sampler self._anneal = 1.0 self._steps_since_update = 0