From 8b0668188fc41a2ddcb564ef64131ff71303df22 Mon Sep 17 00:00:00 2001 From: lefreud Date: Fri, 10 May 2024 12:55:31 +0000 Subject: [PATCH] Working hdr-nerfacto --- .gitignore | 1 + README-pandora.md | 5 ++--- nerfstudio/field_components/field_heads.py | 1 + nerfstudio/hdr_nerfacto/hdr_nerf_field.py | 6 ++++++ nerfstudio/hdr_nerfacto/hdr_nerf_model.py | 2 ++ nerfstudio/model_components/renderers.py | 4 ++-- nerfstudio/utils/colormaps.py | 6 ++++-- 7 files changed, 18 insertions(+), 7 deletions(-) diff --git a/.gitignore b/.gitignore index 6ec98f95b..9ea2dafce 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ Dual_Cameras gt_test hdr-nerfacto-subset-out +ablations_out # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/README-pandora.md b/README-pandora.md index 06a749fed..ae849bac3 100644 --- a/README-pandora.md +++ b/README-pandora.md @@ -132,9 +132,8 @@ It will be trained on the Laval indoor HDR dataset to do this task. TODO: validate this part -## 4. Train from scratch or fine-tune the NeRF model using MomoNet's predictions - - +## 4. Evaluation ``` +ns-eval --load-config outputs/LVSN_LAB_OpenSFM/hdr-nerfacto/2024-05-10_032723/config.yml --output-path ablations_out/hdr_nerfacto/output.json --render-output-path ablations_out/hdr-nerfacto ``` \ No newline at end of file diff --git a/nerfstudio/field_components/field_heads.py b/nerfstudio/field_components/field_heads.py index ff9831079..5ed438b15 100644 --- a/nerfstudio/field_components/field_heads.py +++ b/nerfstudio/field_components/field_heads.py @@ -43,6 +43,7 @@ class FieldHeadNames(Enum): GRADIENT = "gradient" RGB_HDR = "rgb_hdr" # 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/hdr_nerf_field.py b/nerfstudio/hdr_nerfacto/hdr_nerf_field.py index 362d5d2c5..aff4aa922 100644 --- a/nerfstudio/hdr_nerfacto/hdr_nerf_field.py +++ b/nerfstudio/hdr_nerfacto/hdr_nerf_field.py @@ -318,6 +318,11 @@ class HdrNerfactoField(Field): log_rgb_hdr = self.mlp_head(h) # .view(*outputs_shape, -1).to(directions) rgb_hdr = torch.exp(log_rgb_hdr).view(*outputs_shape, -1).to(directions) log_rgb_hdr_exposed = log_rgb_hdr + torch.log(exposures) + log_rgb_hdr_exposed_fast = log_rgb_hdr + torch.log(exposures * 0.004) # TODO: adaptive fast exposure + + pred_rgb_fast = torch.cat([ + self.crf_mlps[channel](log_rgb_hdr_exposed_fast[..., channel:channel+1]) for channel in range(3) + ], dim=-1).view(*outputs_shape, -1).to(directions) pred_rgb_ldr = torch.cat([ self.crf_mlps[channel](log_rgb_hdr_exposed[..., channel:channel+1]) for channel in range(3) @@ -331,6 +336,7 @@ class HdrNerfactoField(Field): outputs.update({ FieldHeadNames.RGB: pred_rgb_ldr, + FieldHeadNames.RGB_FAST: pred_rgb_fast, FieldHeadNames.RGB_HDR: rgb_hdr, FieldHeadNames.ZERO_RADIANCE_CRF: zero_radiance_crf, }) diff --git a/nerfstudio/hdr_nerfacto/hdr_nerf_model.py b/nerfstudio/hdr_nerfacto/hdr_nerf_model.py index 662ab9f80..65d6c6291 100644 --- a/nerfstudio/hdr_nerfacto/hdr_nerf_model.py +++ b/nerfstudio/hdr_nerfacto/hdr_nerf_model.py @@ -302,6 +302,7 @@ class HdrNerfactoModel(Model): ray_samples_list.append(ray_samples) rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) + rgb_fast = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB_FAST], weights=weights) depth = self.renderer_depth(weights=weights, ray_samples=ray_samples) accumulation = self.renderer_accumulation(weights=weights) @@ -312,6 +313,7 @@ class HdrNerfactoModel(Model): outputs = { "rgb": rgb, + "rgb_fast": rgb_fast, "accumulation": accumulation, "depth": depth, "rgb_hdr": rgb_hdr, diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 24090872e..fbf2fc5a9 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 diff --git a/nerfstudio/utils/colormaps.py b/nerfstudio/utils/colormaps.py index cf15150cb..bb85a6005 100644 --- a/nerfstudio/utils/colormaps.py +++ b/nerfstudio/utils/colormaps.py @@ -43,9 +43,9 @@ class ColormapOptions: """ Whether to invert the output colormap """ exposure_scale: float = 0 """Scale the brightness prior to tonemapping by 2 ^ Exposure""" - use_mu_law: bool = True + use_mu_law: bool = False """Use mu-law de-compression for HDR images""" - use_gamma: bool = True + use_gamma: bool = False """Use gamma for HDR images""" @@ -74,12 +74,14 @@ def apply_colormap( img_gamma_22 = image if colormap_options.use_mu_law: + raise Exception('this shouldnt be used') u = 5000. img_gamma_22 = torch.exp(image * torch.log(torch.tensor(u+1.))) - 1. img_gamma_22 /= u img_gamma_22 *= 2.**(colormap_options.exposure_scale) if colormap_options.use_gamma: + raise Exception('this shouldnt be used') img_gamma_22 = torch.pow(img_gamma_22, 1./2.2) return torch.clamp(img_gamma_22, 0., 1.) else: