From 687906a6d9db3d381119e48f3a593a72fea25db4 Mon Sep 17 00:00:00 2001 From: lefreud Date: Mon, 29 Jan 2024 19:16:47 +0000 Subject: [PATCH] Quick hacky fix for pixel sampler --- nerfstudio/data/pixel_samplers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 27480dcc6..c2fe2407d 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -97,14 +97,14 @@ class PixelSampler: mask: mask of possible pixels in an image to sample from. """ # TODO Defining a flag for sampling even when there is a mask - if isinstance(mask, torch.Tensor): + if isinstance(mask, torch.Tensor) and mask.shape[0] >= 4: # mask size: batch * h * w * 1 # Init state, cache masks indices. if self.states_of_mask == 0: with profiler.time_function("mask squeeze: whole"): mask = mask.squeeze(-1) num_masks = mask.shape[0] - size_of_disk = 100 + size_of_disk = 2 num_disks = ceil(num_masks / size_of_disk) # slice mask to 100 disk random_idx_disk = int(torch.randint(0, num_disks, (1,)))