Prepare zipnerf

nerfstudio/model_components/losses.py

@@ -120,6 +120,69 @@ def interlevel_loss(weights_list, ray_samples_list):
     return loss_interlevel
 
 
+def blur_stepfun(x, y, r):
+    xr, xr_idx = torch.sort(torch.cat([x - r, x + r], dim=-1))
+    y1 = (
+        torch.cat([y, torch.zeros_like(y[..., :1])], dim=-1) - torch.cat([torch.zeros_like(y[..., :1]), y], dim=-1)
+    ) / (2 * r)
+    y2 = torch.cat([y1, -y1], dim=-1).take_along_dim(xr_idx[..., :-1], dim=-1)
+    yr = torch.cumsum((xr[..., 1:] - xr[..., :-1]) * torch.cumsum(y2, dim=-1), dim=-1).clamp_min(0)
+    yr = torch.cat([torch.zeros_like(yr[..., :1]), yr], dim=-1)
+    return xr, yr
+
+
+def sorted_interp_quad(x, xp, fpdf, fcdf):
+    """interp in quadratic"""
+
+    # Identify the location in `xp` that corresponds to each `x`.
+    # The final `True` index in `mask` is the start of the matching interval.
+    mask = x[..., None, :] >= xp[..., :, None]
+
+    def find_interval(x):
+        # Grab the value where `mask` switches from True to False, and vice versa.
+        # This approach takes advantage of the fact that `x` is sorted.
+        x0 = torch.max(torch.where(mask, x[..., None], x[..., :1, None]), -2).values
+        x1 = torch.min(torch.where(~mask, x[..., None], x[..., -1:, None]), -2).values
+        return x0, x1
+
+    fpdf0, fpdf1 = find_interval(fpdf)
+    fcdf0, fcdf1 = find_interval(fcdf)
+    xp0, xp1 = find_interval(xp)
+
+    offset = torch.clip(torch.nan_to_num((x - xp0) / (xp1 - xp0), 0), 0, 1)
+    ret = fcdf0 + (x - xp0) * (fpdf0 + fpdf1 * offset + fpdf0 * (1 - offset)) / 2
+    return ret
+
+
+def zipnerf_loss(weights_list, ray_samples_list):
+    # ground truth s and w (real nerf samples)
+    # This implementation matches ZipNeRF up to the scale.
+    # In the paper the loss is computed as the sum over the ray samples.
+    # Here we take the mean and the multiplier for this loss should be changed accordingly.
+    pulse_widths = [0.003, 0.03]
+    c = ray_samples_to_sdist(ray_samples_list[-1]).detach()
+    w = weights_list[-1][..., 0].detach()
+
+    w_norm = w / (c[..., 1:] - c[..., :-1])
+    loss = 0
+    for i, (ray_samples, weights) in enumerate(zip(ray_samples_list[:-1], weights_list[:-1])):
+        cp = ray_samples_to_sdist(ray_samples)
+        wp = weights[..., 0]  # (num_rays, num_samples)
+        c_, w_ = blur_stepfun(c, w_norm, pulse_widths[i])
+
+        # piecewise linear pdf to piecewise quadratic cdf
+        area = 0.5 * (w_[..., 1:] + w_[..., :-1]) * (c_[..., 1:] - c_[..., :-1])
+        cdf = torch.cat([torch.zeros_like(area[..., :1]), torch.cumsum(area, dim=-1)], dim=-1)
+
+        # query piecewise quadratic interpolation
+        cdf_interp = sorted_interp_quad(cp, c_, w_, cdf)
+
+        # difference between adjacent interpolated values
+        w_s = torch.diff(cdf_interp, dim=-1)
+        loss += ((w_s - wp).clamp_min(0) ** 2 / (wp + 1e-5)).mean()
+    return loss
+
+
 # Verified
 def lossfun_distortion(t, w):
     """

nerfstudio/model_components/ray_samplers.py

@@ -262,7 +262,7 @@ class PDFSampler(Sampler):
         train_stratified: bool = True,
         single_jitter: bool = False,
         include_original: bool = True,
-        histogram_padding: float = 0.01,
+        histogram_padding: float = 0.00,
     ) -> None:
         super().__init__(num_samples=num_samples)
         self.train_stratified = train_stratified
@@ -384,7 +384,6 @@ class VolumetricSampler(Sampler):
         occupancy_grid: OccGridEstimator,
         density_fn: Optional[Callable[[TensorType[..., 3]], TensorType[..., 1]]] = None,
     ) -> None:
-
         super().__init__()
         assert occupancy_grid is not None
         self.density_fn = density_fn
@@ -651,7 +650,6 @@ class NeuSSampler(Sampler):
         base_variance = self.base_variance
 
         while total_iters < self.num_upsample_steps:
-
             with torch.no_grad():
                 new_sdf = sdf_fn(new_samples)
 

nerfstudio/models/nerfacto.py

@@ -45,6 +45,7 @@ from nerfstudio.model_components.losses import (
     interlevel_loss,
     orientation_loss,
     pred_normal_loss,
+    zipnerf_loss,
 )
 from nerfstudio.model_components.ray_samplers import (
     ProposalNetworkSampler,
@@ -85,9 +86,9 @@ class NerfactoModelConfig(ModelConfig):
     """Maximum resolution of the hashmap for the base mlp."""
     log2_hashmap_size: int = 19
     """Size of the hashmap for the base mlp"""
-    num_proposal_samples_per_ray: Tuple[int, ...] = (256, 96)
+    num_proposal_samples_per_ray: Tuple[int, ...] = (256, 128)
     """Number of samples per ray for each proposal network."""
-    num_nerf_samples_per_ray: int = 48
+    num_nerf_samples_per_ray: int = 64
     """Number of samples per ray for the nerf network."""
     proposal_update_every: int = 5
     """Sample every n steps after the warmup"""
@@ -99,8 +100,8 @@ class NerfactoModelConfig(ModelConfig):
     """Use the same proposal network. Otherwise use different ones."""
     proposal_net_args_list: List[Dict] = field(
         default_factory=lambda: [
-            {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128, "use_linear": False},
             {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256, "use_linear": False},
+            {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 512, "use_linear": False},
         ]
     )
     """Arguments for the proposal density fields."""
@@ -120,7 +121,7 @@ class NerfactoModelConfig(ModelConfig):
     """Whether to use average appearance embedding or zeros for inference."""
     proposal_weights_anneal_slope: float = 10.0
     """Slope of the annealing function for the proposal weights."""
-    proposal_weights_anneal_max_num_iters: int = 1000
+    proposal_weights_anneal_max_num_iters: int = 1
     """Max num iterations for the annealing function."""
     use_single_jitter: bool = True
     """Whether use single jitter or not for the proposal networks."""
@@ -128,6 +129,8 @@ class NerfactoModelConfig(ModelConfig):
     """Whether to predict normals or not."""
     disable_scene_contraction: bool = False
     """Whether to disable scene contraction or not."""
+    use_zipnerf_loss: bool = False
+    """Whether to use zipnerf loss or not."""
 
 
 class NerfactoModel(Model):
@@ -318,9 +321,14 @@ class NerfactoModel(Model):
         image = batch["image"].to(self.device)
         loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
         if self.training:
-            loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
-                outputs["weights_list"], outputs["ray_samples_list"]
-            )
+            if self.config.use_zipnerf_loss:
+                loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * zipnerf_loss(
+                    outputs["weights_list"], outputs["ray_samples_list"]
+                )
+            else:
+                loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
+                    outputs["weights_list"], outputs["ray_samples_list"]
+                )
             assert metrics_dict is not None and "distortion" in metrics_dict
             loss_dict["distortion_loss"] = self.config.distortion_loss_mult * metrics_dict["distortion"]
             if self.config.predict_normals: