diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index 9b2057357..6e041b33d 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -180,26 +180,30 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): return data CONSOLE.log("Caching / undistorting train images") - with ThreadPoolExecutor() as executor: + with ThreadPoolExecutor(max_workers=2) as executor: cached_train = list( - executor.map( - process_train_data, - track( + track( + executor.map( + process_train_data, range(len(self.train_dataset)), - description="Caching / undistorting train images", - transient=True, ), + description="Caching / undistorting train images", + transient=True, + total=len(self.train_dataset), ) ) CONSOLE.log("Caching / undistorting eval images") - with ThreadPoolExecutor() as executor: + with ThreadPoolExecutor(max_workers=2) as executor: cached_eval = list( - executor.map( - process_eval_data, - track( - range(len(self.eval_dataset)), description="Caching / undistorting eval images", transient=True + track( + executor.map( + process_eval_data, + range(len(self.eval_dataset)), ), + description="Caching / undistorting eval images", + transient=True, + total=len(self.eval_dataset), ) ) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index ef6238aaf..c24bb4280 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -178,22 +178,50 @@ class PixelSampler: if isinstance(mask, torch.Tensor) and not self.config.ignore_mask: indices = self.sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device) else: - rand_samples = torch.rand((batch_size, 3), device=device) - # convert random samples tto radius and theta - radii = self.config.fisheye_crop_radius * torch.sqrt(rand_samples[:, 1]) - theta = 2.0 * torch.pi * rand_samples[:, 2] + # Rejection sampling. + valid: Optional[torch.Tensor] = None + indices = None + while True: + samples_needed = batch_size if valid is None else int(batch_size - torch.sum(valid).item()) - # convert radius and theta to x and y between -radii and radii - x = radii * torch.cos(theta) - y = radii * torch.sin(theta) + # Check if done! + if samples_needed == 0: + break - # Multiply by the batch size and height/width to get pixel indices. - indices = torch.floor( - torch.stack([rand_samples[:, 0], y, x], dim=1) - * torch.tensor([num_images, image_height // 2, image_width // 2], device=device) - + torch.tensor([0, image_height // 2, image_width // 2], device=device) - ).long() + rand_samples = torch.rand((samples_needed, 2), device=device) + # Convert random samples to radius and theta. + radii = self.config.fisheye_crop_radius * torch.sqrt(rand_samples[:, 0]) + theta = 2.0 * torch.pi * rand_samples[:, 1] + # Convert radius and theta to x and y. + x = (radii * torch.cos(theta) + image_width // 2).long() + y = (radii * torch.sin(theta) + image_height // 2).long() + sampled_indices = torch.stack( + [torch.randint(0, num_images, size=(samples_needed,), device=device), y, x], dim=-1 + ) + + # Update indices. + if valid is None: + indices = sampled_indices + valid = ( + (sampled_indices[:, 1] >= 0) + & (sampled_indices[:, 1] < image_height) + & (sampled_indices[:, 2] >= 0) + & (sampled_indices[:, 2] < image_width) + ) + else: + assert indices is not None + not_valid = ~valid + indices[not_valid, :] = sampled_indices + valid[not_valid] = ( + (sampled_indices[:, 1] >= 0) + & (sampled_indices[:, 1] < image_height) + & (sampled_indices[:, 2] >= 0) + & (sampled_indices[:, 2] < image_width) + ) + assert indices is not None + + assert indices.shape == (batch_size, 3) return indices def collate_image_dataset_batch(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): diff --git a/nerfstudio/viewer/render_state_machine.py b/nerfstudio/viewer/render_state_machine.py index b0afb490c..1fa7c07d9 100644 --- a/nerfstudio/viewer/render_state_machine.py +++ b/nerfstudio/viewer/render_state_machine.py @@ -306,7 +306,7 @@ class RenderStateMachine(threading.Thread): jpeg_quality=jpg_quality, depth=depth, ) - res = f"{selected_output.shape[0]}x{selected_output.shape[1]}px" + res = f"{selected_output.shape[1]}x{selected_output.shape[0]}px" self.viewer.stats_markdown.content = self.viewer.make_stats_markdown(None, res) def _calculate_image_res(self, aspect_ratio: float) -> Tuple[int, int]: