From 56556d22fa0955b79270568e2aae92f18e071af4 Mon Sep 17 00:00:00 2001 From: Nicholas Pfaff <53228351+nepfaff@users.noreply.github.com> Date: Tue, 1 Oct 2024 12:54:28 -0400 Subject: [PATCH] Rejection sample mask for patch pixel sampler (#3457) Use mask rejection sampling in PatchPixelSampler --- nerfstudio/data/pixel_samplers.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index a0f1dd67c..63b4a35b0 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -420,6 +420,10 @@ class PatchPixelSamplerConfig(PixelSamplerConfig): patch_size: int = 32 """Side length of patch. This must be consistent in the method config in order for samples to be reshaped into patches correctly.""" + rejection_sample_mask: bool = True + """Whether or not to use rejection sampling when sampling images with masks""" + max_num_iterations: int = 100 + """If rejection sampling masks, the maximum number of times to sample""" class PatchPixelSampler(PixelSampler): @@ -458,9 +462,20 @@ class PatchPixelSampler(PixelSampler): sub_bs = batch_size // (self.config.patch_size**2) half_patch_size = int(self.config.patch_size / 2) m = erode_mask(mask.permute(0, 3, 1, 2).float(), pixel_radius=half_patch_size) - nonzero_indices = torch.nonzero(m[:, 0], as_tuple=False).to(device) - chosen_indices = random.sample(range(len(nonzero_indices)), k=sub_bs) - indices = nonzero_indices[chosen_indices] + + if self.config.rejection_sample_mask: + indices = self.rejection_sample_mask( + mask=m, + num_samples=sub_bs, + num_images=num_images, + image_height=image_height, + image_width=image_width, + device=device, + ) + else: + nonzero_indices = torch.nonzero(m[:, 0], as_tuple=False).to(device) + chosen_indices = random.sample(range(len(nonzero_indices)), k=sub_bs) + indices = nonzero_indices[chosen_indices] indices = ( indices.view(sub_bs, 1, 1, 3)