From 1fa5cba2d7773ec28c944480929ad2ea20750a51 Mon Sep 17 00:00:00 2001 From: lefreud Date: Mon, 13 May 2024 02:09:24 +0000 Subject: [PATCH] Clipped HDR-nerfacto --- nerfstudio/configs/method_configs.py | 5 +++-- nerfstudio/data/pixel_samplers.py | 2 +- nerfstudio/field_components/field_heads.py | 3 ++- nerfstudio/hdr_nerfacto/config.py | 9 +++++++-- nerfstudio/hdr_nerfacto/hdr_nerf_field.py | 2 +- nerfstudio/hdr_nerfacto/hdr_nerf_model.py | 19 ++++++++++++++++--- nerfstudio/model_components/renderers.py | 4 ++-- 7 files changed, 32 insertions(+), 12 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 38776f3da..5b092a56f 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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", diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index c2fe2407d..1b8533e9f 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/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(): diff --git a/nerfstudio/field_components/field_heads.py b/nerfstudio/field_components/field_heads.py index 5ed438b15..8f238912a 100644 --- a/nerfstudio/field_components/field_heads.py +++ b/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 diff --git a/nerfstudio/hdr_nerfacto/config.py b/nerfstudio/hdr_nerfacto/config.py index ecb9168dc..b9a4876bd 100644 --- a/nerfstudio/hdr_nerfacto/config.py +++ b/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={ diff --git a/nerfstudio/hdr_nerfacto/hdr_nerf_field.py b/nerfstudio/hdr_nerfacto/hdr_nerf_field.py index 9143339bc..e5af2dff5 100644 --- a/nerfstudio/hdr_nerfacto/hdr_nerf_field.py +++ b/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: diff --git a/nerfstudio/hdr_nerfacto/hdr_nerf_model.py b/nerfstudio/hdr_nerfacto/hdr_nerf_model.py index 991ca55df..5b14835b4 100644 --- a/nerfstudio/hdr_nerfacto/hdr_nerf_model.py +++ b/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( diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index fbf2fc5a9..24090872e 100644 --- a/nerfstudio/model_components/renderers.py +++ b/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