Clipped HDR-nerfacto

nerfstudio/configs/method_configs.py

@@ -129,8 +129,9 @@ method_configs["nerfacto"] = TrainerConfig(
 
 method_configs["lantern-nerfacto"] = Get_lantern_config()
 
-method_configs["hdr-nerfacto"] = get_hdr_nerfacto_config(use_crf=True)
-method_configs["hdr-nerfacto-wo-crf"] = get_hdr_nerfacto_config(use_crf=False)
+method_configs["hdr-nerfacto"] = get_hdr_nerfacto_config(use_crf=True, clip_before_accumulation=False)
+method_configs["hdr-nerfacto-wo-crf"] = get_hdr_nerfacto_config(use_crf=False, clip_before_accumulation=False)
+method_configs["hdr-nerfacto-wo-crf-clip-bf-acc"] = get_hdr_nerfacto_config(use_crf=False, clip_before_accumulation=True)
 
 method_configs["nerfacto-big"] = TrainerConfig(
     method_name="nerfacto",

nerfstudio/data/pixel_samplers.py

@@ -213,7 +213,7 @@ class PixelSampler:
         c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1))
         c, y, x = c.cpu(), y.cpu(), x.cpu()
         collated_batch = {
-            key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and key != "exposure" and value is not None
+            key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and key != "exposure" and key != "image_filename" and value is not None
         }
 
         for key, value in batch.items():

nerfstudio/field_components/field_heads.py

@@ -42,7 +42,8 @@ class FieldHeadNames(Enum):
     ALPHA = "alpha"
     GRADIENT = "gradient"
     
-    RGB_HDR = "rgb_hdr"                        # HDR-Nerfacto
+    RGB_HDR = "rgb_hdr"                        # HDR-Nerfacto`
+    RGB_LINEAR_CLIPPED = "rgb_linear_clipped"  # HDR-Nerfacto
     RGB_FAST = "rgb_fast"                      # HDR-Nerfacto
     PRED_RGB_LDR = "pred_rgb_ldr"              # HDR-Nerfacto
     ZERO_RADIANCE_CRF = "zero_radiance_crf"    # HDR-Nerfacto

nerfstudio/hdr_nerfacto/config.py

@@ -10,11 +10,15 @@ from nerfstudio.lantern.datamanager import HDRNerfactoDataManagerConfig, HDRNerf
 from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig
 
 
-def get_hdr_nerfacto_config(*, use_crf=True):
+def get_hdr_nerfacto_config(*, use_crf=True, clip_before_accumulation=False):
     if use_crf:
         method_name = "hdr-nerfacto"
+        assert not clip_before_accumulation
     else:
-        method_name = "hdr-nerfacto-wo-crf"
+        if clip_before_accumulation:
+            method_name = "hdr-nerfacto-wo-crf-clip-bf-acc"
+        else:
+            method_name = "hdr-nerfacto-wo-crf"
     
     datamanager_config_class = HDRNerfactoDataManagerConfig if use_crf else HDRNerfactoWoCrfDataManagerConfig
         
@@ -37,6 +41,7 @@ def get_hdr_nerfacto_config(*, use_crf=True):
             model=HdrNerfactoModelConfig(
                 eval_num_rays_per_chunk=1 << 15,
                 use_crf=use_crf,
+                clip_before_accumulation=clip_before_accumulation,
             ),
         ),
         optimizers={

nerfstudio/hdr_nerfacto/hdr_nerf_field.py

@@ -326,7 +326,7 @@ class HdrNerfactoField(Field):
         rgb_hdr = torch.exp(log_rgb_hdr_exposed).view(*outputs_shape, -1).to(directions)
         #print('rgb_hdr_min', rgb_hdr.min(), 'rgb_hdr_max', rgb_hdr.max())
 
-        # outputs[FieldHeadNames.RGB_HDR] = torch.clamp(rgb_hdr, 0, 1)
+        outputs[FieldHeadNames.RGB_LINEAR_CLIPPED] = torch.clamp(rgb_hdr, 0, 1)
         outputs[FieldHeadNames.RGB_HDR] = rgb_hdr
 
         if self.use_crf:

nerfstudio/hdr_nerfacto/hdr_nerf_model.py

@@ -38,6 +38,7 @@ from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.field_components.spatial_distortions import SceneContraction
 from nerfstudio.fields.density_fields import HashMLPDensityField
 from nerfstudio.hdr_nerfacto.hdr_nerf_field import HdrNerfactoField
+from nerfstudio.lantern.renderer import RGBRenderer_HDR
 from nerfstudio.model_components.losses import (
     MSELoss,
     distortion_loss,
@@ -143,6 +144,7 @@ class HdrNerfactoModelConfig(ModelConfig):
     unit_exposure_expected_value: float = 0.5
 
     use_crf: bool = True
+    clip_before_accumulation: bool = False
 
 
 class HdrNerfactoModel(Model):
@@ -237,6 +239,7 @@ class HdrNerfactoModel(Model):
 
         # renderers
         self.renderer_rgb = RGBRenderer(background_color=self.config.background_color)
+        self.renderer_rgb_hdr = RGBRenderer_HDR(background_color=self.config.background_color)
         self.renderer_accumulation = AccumulationRenderer()
         self.renderer_depth = DepthRenderer()
         self.renderer_normals = NormalsRenderer()
@@ -305,9 +308,18 @@ class HdrNerfactoModel(Model):
         ray_samples_list.append(ray_samples)
 
         # HDR-Nerfacto
-        rgb_hdr = self.renderer_rgb(
+        rgb_hdr = self.renderer_rgb_hdr(
             rgb=field_outputs[FieldHeadNames.RGB_HDR], weights=weights
         )
+        if self.config.clip_before_accumulation:
+            rgb_linear_clipped = self.renderer_rgb(
+                rgb=field_outputs[FieldHeadNames.RGB_LINEAR_CLIPPED], weights=weights
+            )
+        else:
+            # clipping only after accumulation (in the renderer)
+            rgb_linear_clipped = self.renderer_rgb(
+                rgb=field_outputs[FieldHeadNames.RGB_HDR], weights=weights
+            )
 
         depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
         accumulation = self.renderer_accumulation(weights=weights)
@@ -316,6 +328,7 @@ class HdrNerfactoModel(Model):
             "accumulation": accumulation,
             "depth": depth,
             "rgb_hdr": rgb_hdr,
+            "rgb_linear_clipped": rgb_linear_clipped,
         }
         if self.config.use_crf:
             rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
@@ -379,7 +392,7 @@ class HdrNerfactoModel(Model):
             )
         else:
             pred_rgb, gt_rgb = self.renderer_rgb.blend_background_for_loss_computation(
-                pred_image=outputs["rgb_hdr"],
+                pred_image=outputs["rgb_linear_clipped"],
                 pred_accumulation=outputs["accumulation"],
                 gt_image=image,
             )
@@ -412,7 +425,7 @@ class HdrNerfactoModel(Model):
         if self.config.use_crf:
             predicted_rgb = outputs["rgb"]
         else:
-            predicted_rgb = outputs["rgb_hdr"]
+            predicted_rgb = outputs["rgb_linear_clipped"]
         gt_rgb = self.renderer_rgb.blend_background(gt_rgb)
         acc = colormaps.apply_colormap(outputs["accumulation"])
         depth = colormaps.apply_depth_colormap(

nerfstudio/model_components/renderers.py

@@ -224,8 +224,8 @@ class RGBRenderer(nn.Module):
         rgb = self.combine_rgb(
             rgb, weights, background_color=background_color, ray_indices=ray_indices, num_rays=num_rays
         )
-        # if not self.training:
-        #     torch.clamp_(rgb, min=0.0, max=1.0)
+        if not self.training:
+            torch.clamp_(rgb, min=0.0, max=1.0)
         return rgb