initial impl

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"""

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

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

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(

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 (