diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index d032c4d48..8114cf6cb 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -50,7 +50,7 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import ( PhototourismDataParserConfig, ) from nerfstudio.data.datasets.base_dataset import InputDataset -from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PixelSampler +from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PixelSampler,PatchPixelSampler from nerfstudio.data.utils.dataloaders import ( CacheDataloader, FixedIndicesEvalDataloader, @@ -351,7 +351,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method # Otherwise, use the default pixel sampler if is_equirectangular.any(): CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.") - return PixelSampler(*args, **kwargs) + return PatchPixelSampler(*args, **kwargs,patch_size=3) def setup_train(self): """Sets up the data loaders for training""" diff --git a/nerfstudio/data/dataparsers/nuscenes_dataparser.py b/nerfstudio/data/dataparsers/nuscenes_dataparser.py index 9269f555f..eaf9f15c7 100644 --- a/nerfstudio/data/dataparsers/nuscenes_dataparser.py +++ b/nerfstudio/data/dataparsers/nuscenes_dataparser.py @@ -21,7 +21,7 @@ from typing import Optional, Tuple, Type import numpy as np import torch -from nuscenes.nuscenes import NuScenes as NuScenesDatabase +# from nuscenes.nuscenes import NuScenes as NuScenesDatabase from typing_extensions import Literal from nerfstudio.cameras.cameras import Cameras, CameraType diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 91fc6301d..5aaf52644 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -99,7 +99,6 @@ class PixelSampler: # pylint: disable=too-few-public-methods collated_batch = { key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None } - assert collated_batch["image"].shape == (num_rays_per_batch, 3), collated_batch["image"].shape # Needed to correct the random indices to their actual camera idx locations. @@ -243,3 +242,45 @@ class EquirectangularPixelSampler(PixelSampler): # pylint: disable=too-few-publ ).long() return indices + +class PatchPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods + """todo""" + + def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False, **kwargs) -> None: + self.patch_size = kwargs["patch_size"] + num_rays = (num_rays_per_batch//(self.patch_size**2))*(self.patch_size**2) + super().__init__(num_rays, keep_full_image, **kwargs) + + # overrides base method + def sample_method( # pylint: disable=no-self-use + self, + batch_size: int, + num_images: int, + image_height: int, + image_width: int, + mask: Optional[TensorType] = None, + device: Union[torch.device, str] = "cpu", + ) -> TensorType["batch_size", 3]: + + if mask: + # Note: if there is a mask, sampling reduces back to uniform sampling + indices = super().sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device) + else: + sub_bs = batch_size // (self.patch_size**2) + indices = torch.rand((sub_bs, 3), device=device) * torch.tensor( + [num_images, image_height - self.patch_size, image_width - self.patch_size], + device=device, + ) + + indices = indices.view(sub_bs, 1, 1, 3).broadcast_to(sub_bs, self.patch_size, self.patch_size, 3).clone() + + yys, xxs = torch.meshgrid( + torch.arange(self.patch_size, device=device), torch.arange(self.patch_size, device=device) + ) + indices[:, ..., 1] += yys + indices[:, ..., 2] += xxs + + indices = torch.floor(indices).long() + indices = indices.flatten(0, 2) + + return indices \ No newline at end of file diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index c996e20b3..21f4f6c81 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -184,7 +184,8 @@ class NerfactoModel(Model): # renderers self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) self.renderer_accumulation = AccumulationRenderer() - self.renderer_depth = DepthRenderer() + self.renderer_depth = DepthRenderer(method='expected') + self.renderer_truncdepth = DepthRenderer(method='median') self.renderer_normals = NormalsRenderer() # losses @@ -230,8 +231,18 @@ class NerfactoModel(Model): func=self.proposal_sampler.step_cb, ) ) + callbacks.append( + TrainingCallback( + where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], + update_every_num_iters=1, + func=self.set_step + ) + ) return callbacks + def set_step(self,step): + self.step=step + def get_outputs(self, ray_bundle: RayBundle): ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns) field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals) @@ -248,6 +259,8 @@ class NerfactoModel(Model): "accumulation": accumulation, "depth": depth, } + with torch.no_grad(): + outputs['depth_truncated']=self.renderer_truncdepth(weights=weights, ray_samples=ray_samples) if self.config.predict_normals: outputs["normals"] = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) @@ -302,6 +315,12 @@ class NerfactoModel(Model): loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( outputs["rendered_pred_normal_loss"] ) + #depth regularity loss + if self.step>100: + p=3 + patches = outputs['depth'].view(-1,p,p) #Bx3x3 + loss_dict['depth_reg'] = .1*torch.mean(torch.sum((patches.view(-1,p*p) - patches[:,p//2,p//2][:,None])**2,dim=1)) + return loss_dict def get_image_metrics_and_images( diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index f62cf9d8f..6bca93433 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -615,7 +615,8 @@ class ViewerState: if OutputTypes.ACCUMULATION in self.output_list else OutputTypes.ACCUMULATION_FINE ) - return colormaps.apply_depth_colormap(outputs[reformatted_output], accumulation=outputs[accumulation_str]) + return colormaps.apply_depth_colormap(outputs[reformatted_output], + accumulation=outputs[accumulation_str],near_plane=.1,far_plane=2) # rendering accumulation outputs if self.prev_colormap_type == ColormapTypes.TURBO or (