From 2357332e66db67a65966a977f940188a4d32bd77 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Wed, 15 Feb 2023 15:35:08 -0800 Subject: [PATCH] Proposal network initial sampler (#1435) Allow first proposal sampling method --- nerfstudio/model_components/ray_samplers.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index 5df52d2fa..ef0079ffd 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -509,7 +509,16 @@ class VolumetricSampler(Sampler): class ProposalNetworkSampler(Sampler): - """Sampler that uses a proposal network to generate samples.""" + """Sampler that uses a proposal network to generate samples. + + Args: + num_proposal_samples_per_ray: Number of samples to generate per ray for each proposal step. + num_nerf_samples_per_ray: Number of samples to generate per ray for the NERF model. + num_proposal_network_iterations: Number of proposal network iterations to run. + 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. + """ def __init__( self, @@ -518,6 +527,7 @@ class ProposalNetworkSampler(Sampler): num_proposal_network_iterations: int = 2, single_jitter: bool = False, update_sched: Callable = lambda x: 1, + initial_sampler: Optional[Sampler] = None, ) -> None: super().__init__() self.num_proposal_samples_per_ray = num_proposal_samples_per_ray @@ -528,7 +538,10 @@ class ProposalNetworkSampler(Sampler): raise ValueError("num_proposal_network_iterations must be >= 1") # samplers - self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter) + if initial_sampler is None: + 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) self._anneal = 1.0