Fix: bug

nerfstudio/configs/method_configs.py

@@ -19,44 +19,38 @@ Put all the method implementations in one location.
 from __future__ import annotations
 
 from collections import OrderedDict
-from dataclasses import dataclass, field
 from typing import Dict
+from dataclasses import dataclass, field
 
 import tyro
-from lantern.config import Get_lantern_config
 
 from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
 from nerfstudio.configs.base_config import ViewerConfig
 from nerfstudio.configs.external_methods import get_external_methods
-from nerfstudio.data.datamanagers.base_datamanager import (
-    HDRVanillaDataManagerConfig, VanillaDataManager, VanillaDataManagerConfig)
-from nerfstudio.data.datamanagers.random_cameras_datamanager import \
-    RandomCamerasDataManagerConfig
-from nerfstudio.data.dataparsers.blender_dataparser import \
-    BlenderDataParserConfig
+
+from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManagerConfig
+from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager, VanillaDataManagerConfig, HDRVanillaDataManagerConfig
+
+from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
 from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig
-from nerfstudio.data.dataparsers.instant_ngp_dataparser import \
-    InstantNGPDataParserConfig
-from nerfstudio.data.dataparsers.nerfstudio_dataparser import \
-    NerfstudioDataParserConfig
-from nerfstudio.data.dataparsers.phototourism_dataparser import \
-    PhototourismDataParserConfig
-from nerfstudio.data.dataparsers.sdfstudio_dataparser import \
-    SDFStudioDataParserConfig
-from nerfstudio.data.dataparsers.sitcoms3d_dataparser import \
-    Sitcoms3DDataParserConfig
+from nerfstudio.data.dataparsers.instant_ngp_dataparser import InstantNGPDataParserConfig
+from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
+from nerfstudio.data.dataparsers.phototourism_dataparser import PhototourismDataParserConfig
+from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig
+from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParserConfig
 from nerfstudio.data.datasets.depth_dataset import DepthDataset
-from nerfstudio.data.datasets.hdr_dataset import HDRInputDataset
 from nerfstudio.data.datasets.sdf_dataset import SDFDataset
 from nerfstudio.data.datasets.semantic_dataset import SemanticDataset
-from nerfstudio.engine.optimizers import (AdamOptimizerConfig,
-                                          RAdamOptimizerConfig)
-from nerfstudio.engine.schedulers import (CosineDecaySchedulerConfig,
-                                          ExponentialDecaySchedulerConfig,
-                                          MultiStepSchedulerConfig)
+from nerfstudio.data.datasets.hdr_dataset import HDRInputDataset
+
+from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
+from nerfstudio.engine.schedulers import (
+    CosineDecaySchedulerConfig,
+    ExponentialDecaySchedulerConfig,
+    MultiStepSchedulerConfig,
+)
 from nerfstudio.engine.trainer import TrainerConfig
-from nerfstudio.field_components.temporal_distortions import \
-    TemporalDistortionKind
+from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
 from nerfstudio.fields.sdf_field import SDFFieldConfig
 from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig
 from nerfstudio.models.generfacto import GenerfactoModelConfig
@@ -72,6 +66,8 @@ from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig
 from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig
 from nerfstudio.plugins.registry import discover_methods
 
+from nerfstudio.lantern.config import Get_lantern_config
+
 method_configs: Dict[str, TrainerConfig] = {}
 descriptions = {
     "nerfacto": "Recommended real-time model tuned for real captures. This model will be continually updated.",
@@ -125,6 +121,39 @@ method_configs["nerfacto"] = TrainerConfig(
     vis="viewer",
 )
 
+method_configs["nerfacto-HDR"] = TrainerConfig(
+    method_name="nerfacto-HDR",
+    steps_per_eval_batch=500,
+    steps_per_save=2000,
+    max_num_iterations=30000,
+    mixed_precision=True,
+    pipeline=VanillaPipelineConfig(
+            datamanager=HDRVanillaDataManagerConfig(
+            dataparser=NerfstudioDataParserConfig(),
+            train_num_rays_per_batch=4096,
+            eval_num_rays_per_batch=4096,
+            camera_optimizer=CameraOptimizerConfig(
+                mode="SO3xR3",
+                optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2),
+                scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000),
+            ),
+        ),
+        model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
+    ),
+    optimizers={
+        "proposal_networks": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
+        },
+        "fields": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
 method_configs["lantern-nerfacto"] = Get_lantern_config()
 
 method_configs["nerfacto-big"] = TrainerConfig(

nerfstudio/lantern/__init__.py

@@ -0,0 +1 @@
+# init

nerfstudio/lantern/config.py

@@ -6,6 +6,8 @@ Define custom method here that registers with Nerfstudio CLI.
 
 from __future__ import annotations
 
+from nerfstudio.lantern.datamanager import HDRVanillaDataManagerConfig
+from nerfstudio.lantern.model import LanternModelConfig
 from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
 from nerfstudio.configs.base_config import ViewerConfig
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import \
@@ -14,15 +16,12 @@ from nerfstudio.engine.optimizers import (AdamOptimizerConfig,
                                           RAdamOptimizerConfig)
 from nerfstudio.engine.schedulers import ExponentialDecaySchedulerConfig
 from nerfstudio.engine.trainer import TrainerConfig
-from nerfstudio.lantern.datamanager import HDRVanillaDataManagerConfig
-from nerfstudio.lantern.model import LanternModelConfig
 from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig
 from nerfstudio.plugins.types import MethodSpecification
 
-
 def Get_lantern_config():
     return TrainerConfig(
-    method_name="Lantern-nerfacto",
+    method_name="lantern-nerfacto",
     steps_per_eval_batch=500,
     steps_per_save=2000,
     max_num_iterations=30000,

nerfstudio/lantern/datamanager.py

@@ -10,7 +10,7 @@ import torch
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.data.datamanagers.base_datamanager import (
     VanillaDataManager, VanillaDataManagerConfig)
-from nerfstudio.data.datasets.hdr_dataset import HDRInputDataset
+from nerfstudio.lantern.dataset import HDRInputDataset
 
 
 @dataclass

nerfstudio/lantern/dataset.py

@@ -87,9 +87,6 @@ class HDRInputDataset(InputDataset):
         # Compress HDR using: pixel = log(pixel + 1.)
             
         image = torch.from_numpy(self.get_numpy_image(image_idx).astype("float32"))
-        # image = torch.log(1. + image)
-        # image = torch.log(1. + image)
-        # image = torch.log(1. + image)
         
         # Compress HDR using u-law: pixel = log(u * pixel + 1.) / log(u + 1.)
         u = 5000.

nerfstudio/lantern/model.py

@@ -12,9 +12,11 @@ from torchmetrics.functional import structural_similarity_index_measure
 from torchmetrics.image import PeakSignalNoiseRatio
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
 
+from nerfstudio.lantern.field import LanternNerfactoField
+from nerfstudio.lantern.renderer import RGBRenderer_HDR
+
 from nerfstudio.field_components.spatial_distortions import SceneContraction
 from nerfstudio.fields.density_fields import HashMLPDensityField
-from nerfstudio.lantern.field import LanternNerfactoField
 from nerfstudio.model_components.losses import (
     MSELoss, distortion_loss, interlevel_loss, orientation_loss,
     pred_normal_loss, scale_gradients_by_distance_squared)
@@ -127,7 +129,7 @@ class LanternModel(NerfactoModel):
         self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)
 
         # renderers
-        self.renderer_rgb = RGBRenderer(background_color=self.config.background_color)
+        self.renderer_rgb = RGBRenderer_HDR(background_color=self.config.background_color)
         self.renderer_accumulation = AccumulationRenderer()
         self.renderer_depth = DepthRenderer()
         self.renderer_normals = NormalsRenderer()
@@ -138,7 +140,7 @@ class LanternModel(NerfactoModel):
         # TODO losses
         self.rgb_loss = MSELoss()
 
-        # metrics
+        # metrics, psnr
         self.psnr = PeakSignalNoiseRatio(data_range=1.0)
         self.ssim = structural_similarity_index_measure
         self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True)

nerfstudio/lantern/renderer.py

@@ -0,0 +1,47 @@
+from nerfstudio.model_components.renderers import RGBRenderer
+import torch
+from jaxtyping import Float, Int
+from torch import Tensor, nn
+
+from typing import Generator, Literal, Optional, Tuple, Union
+from nerfstudio.cameras.rays import RaySamples
+from nerfstudio.utils import colors
+from nerfstudio.utils.math import components_from_spherical_harmonics, safe_normalize
+
+BackgroundColor = Union[Literal["random", "last_sample", "black", "white"], Float[Tensor, "3"], Float[Tensor, "*bs 3"]]
+BACKGROUND_COLOR_OVERRIDE: Optional[Float[Tensor, "3"]] = None
+
+class RGBRenderer_HDR(RGBRenderer):
+    
+    def forward(
+        self,
+        rgb: Float[Tensor, "*bs num_samples 3"],
+        weights: Float[Tensor, "*bs num_samples 1"],
+        ray_indices: Optional[Int[Tensor, "num_samples"]] = None,
+        num_rays: Optional[int] = None,
+        background_color: Optional[BackgroundColor] = None,
+    ) -> Float[Tensor, "*bs 3"]:
+        """Composite samples along ray and render color image
+
+        Args:
+            rgb: RGB for each sample
+            weights: Weights for each sample
+            ray_indices: Ray index for each sample, used when samples are packed.
+            num_rays: Number of rays, used when samples are packed.
+            background_color: The background color to use for rendering.
+
+        Returns:
+            Outputs of rgb values.
+        """
+        if background_color is None:
+            background_color = self.background_color
+
+        if not self.training:
+            rgb = torch.nan_to_num(rgb)
+        rgb = self.combine_rgb(
+            rgb, weights, background_color=background_color, ray_indices=ray_indices, num_rays=num_rays
+        )
+        '''annote code below to enable HDR radiance'''
+        # if not self.training:
+        #     torch.clamp_(rgb, min=0.0, max=1.0)
+        return rgb 

nerfstudio/process_data/aligned_pano2plane.py

@@ -17,16 +17,16 @@
 import json
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Literal, Optional, Tuple
+from typing import Optional, Literal, Tuple
 
 import numpy as np
 
 from nerfstudio.process_data import equirect_utils, process_data_utils
-from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import \
-    BaseConverterToNerfstudioDataset
-from nerfstudio.utils import io
+from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import (
+    BaseConverterToNerfstudioDataset,
+)
 from nerfstudio.utils.rich_utils import CONSOLE
-
+from nerfstudio.utils import io
 
 @dataclass
 class ProcessAlignedPano(BaseConverterToNerfstudioDataset):

nerfstudio/process_data/equirect_utils.py

@@ -15,27 +15,32 @@
 """Helper utils for processing equirectangular data."""
 
 import os
-
 os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
 
 import json
-import math
 import sys
 from pathlib import Path
 from typing import List, Tuple
+import math
 
 import cv2
 import numpy as np
-import torch
 from numpy.linalg import inv
-from rich.progress import (BarColumn, Progress, TaskProgressColumn, TextColumn,
-                           TimeRemainingColumn)
-from scipy.spatial.transform import Rotation
 
+from scipy.spatial.transform import Rotation
+import torch
 from equilib import Equi2Pers
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
+
+from nerfstudio.utils.rich_utils import CONSOLE, ItersPerSecColumn
 from nerfstudio.process_data.process_data_utils import CAMERA_MODELS
 from nerfstudio.utils import io
-from nerfstudio.utils.rich_utils import CONSOLE, ItersPerSecColumn
 
 
 def _crop_bottom(bound_arr: list, fov: int, crop_factor: float) -> List[float]:

nerfstudio/scripts/render.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 
 import json
 import os
+os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
 import struct
 import shutil
 import sys
@@ -28,6 +29,7 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Dict, List, Literal, Optional, Union
 
+import cv2
 import mediapy as media
 import numpy as np
 import torch
@@ -72,7 +74,7 @@ def _render_trajectory_video(
     rendered_resolution_scaling_factor: float = 1.0,
     seconds: float = 5.0,
     output_format: Literal["images", "video"] = "video",
-    image_format: Literal["jpeg", "png"] = "jpeg",
+    image_format: Literal["jpeg", "png", "exr"] = "jpeg",
     jpeg_quality: int = 100,
     depth_near_plane: Optional[float] = None,
     depth_far_plane: Optional[float] = None,
@@ -97,7 +99,7 @@ def _render_trajectory_video(
     cameras.rescale_output_resolution(rendered_resolution_scaling_factor)
     cameras = cameras.to(pipeline.device)
     fps = len(cameras) / seconds
-
+    print(len(cameras))
     progress = Progress(
         TextColumn(":movie_camera: Rendering :movie_camera:"),
         BarColumn(),
@@ -166,6 +168,13 @@ def _render_trajectory_video(
                             .cpu()
                             .numpy()
                         )
+                    elif image_format == "exr":
+                        u = 5000.
+                        output_image = torch.exp(output_image * torch.log(torch.tensor(u+1.))) - 1.
+                        output_image /= u
+                        output_image = (
+                            output_image.cpu().numpy()
+                        )
                     else:
                         output_image = (
                             colormaps.apply_colormap(
@@ -184,6 +193,12 @@ def _render_trajectory_video(
                         media.write_image(
                             output_image_dir / f"{camera_idx:05d}.jpg", render_image, fmt="jpeg", quality=jpeg_quality
                         )
+                    if image_format == "exr":
+                        output_path_str = str(output_image_dir / f"{camera_idx:05d}.exr")
+                        CONSOLE.print(f"HDR Img {camera_idx}, max value: {render_image.max()}", justify="center")
+                        render_image = cv2.cvtColor(render_image, cv2.COLOR_BGR2RGB)
+                        cv2.imwrite(output_path_str, render_image)
+                        
                 if output_format == "video":
                     if writer is None:
                         render_width = int(render_image.shape[1])
@@ -323,7 +338,7 @@ class BaseRender:
     """Path to config YAML file."""
     output_path: Path = Path("renders/output.mp4")
     """Path to output video file."""
-    image_format: Literal["jpeg", "png"] = "jpeg"
+    image_format: Literal["jpeg", "png", "exr"] = "jpeg"
     """Image format"""
     jpeg_quality: int = 100
     """JPEG quality"""

nerfstudio/utils/colormaps.py

@@ -65,17 +65,14 @@ def apply_colormap(
 
     # default for rgb images
     if image.shape[-1] == 3:
-        if(colormap_options.exposure_scale != 0 or image.max() > 1.):
+        if(colormap_options.exposure_scale != 0 or image.max() > 1):
             # Log
-            img_gamma_22 = torch.exp(image) - 1.
+            # img_gamma_22 = torch.exp(image) - 1.
             
             # u-law
             u = 5000.
             img_gamma_22 = torch.exp(image * torch.log(torch.tensor(u+1.))) - 1.
             img_gamma_22 /= u
-            
-            # img_gamma_22 = torch.exp(img_gamma_22) - 1.
-            # img_gamma_22 = torch.exp(img_gamma_22) - 1.
             img_gamma_22 *= 2.**(colormap_options.exposure_scale)
             img_gamma_22 = torch.pow(img_gamma_22, 1./2.2)
             return torch.clamp(img_gamma_22, 0., 1.)