Merge remote-tracking branch 'remotes/origin/jmchen' into main

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 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,7 +121,40 @@ method_configs["nerfacto"] = TrainerConfig(
     vis="viewer",
 )
 
-method_configs["lantern-nerfacto"] = lantern_config
+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(
     method_name="nerfacto",

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

@@ -114,11 +114,8 @@ def generate_planar_projections_from_equirectangular(
     image_dir: Path,
     planar_image_size: Tuple[int, int],
     samples_per_im: int,
-<<<<<<< HEAD
     mask_dir: Path,
     hdr_dir: Path,
-=======
->>>>>>> origin/jmchen
     crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
     is_HDR: bool = False,
     HDR_planar_image_size: Tuple[int, int] = (None, None),
@@ -182,7 +179,6 @@ def generate_planar_projections_from_equirectangular(
     frame_dir = image_dir
     output_dir = image_dir / "planar_projections"
     output_dir.mkdir(exist_ok=True)
-<<<<<<< HEAD
     
     output_mask_dir = None
     if mask_dir is not None:
@@ -196,8 +192,6 @@ def generate_planar_projections_from_equirectangular(
     else:
         equi2pers_hdr = equi2pers
     
-=======
->>>>>>> origin/jmchen
     num_ims = len(os.listdir(frame_dir))
     progress = Progress(
         TextColumn("[bold blue]Generating Planar Images", justify="right"),
@@ -209,7 +203,6 @@ def generate_planar_projections_from_equirectangular(
 
     with progress:
         for i in progress.track(os.listdir(frame_dir), description="", total=num_ims):
-<<<<<<< HEAD
             if i.lower().endswith((".jpg", ".png", ".jpeg", ".exr")):
                 if i.lower().endswith((".exr")):
                     im = np.array(cv2.imread(os.path.join(frame_dir, i), cv2.IMREAD_UNCHANGED)).astype("float32")
@@ -243,17 +236,10 @@ def generate_planar_projections_from_equirectangular(
                     hdr = torch.tensor(hdr, dtype=torch.float32, device=device)
                     hdr = torch.permute(hdr, (2, 0, 1))
 
-=======
-            if i.lower().endswith((".jpg", ".png", ".jpeg")):
-                im = np.array(cv2.imread(os.path.join(frame_dir, i)))
-                im = torch.tensor(im, dtype=torch.float32, device=device)
-                im = torch.permute(im, (2, 0, 1)) / 255.0
->>>>>>> origin/jmchen
                 count = 0
                 for u_deg, v_deg in yaw_pitch_pairs:
                     v_rad = torch.pi * v_deg / 180.0
                     u_rad = torch.pi * u_deg / 180.0
-<<<<<<< HEAD
                     pers_image = equi2pers(im, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad})
                     if mask_dir is not None:
                         pers_mask = equi2pers_hdr(mask, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) * 255.0
@@ -286,18 +272,7 @@ def generate_planar_projections_from_equirectangular(
 
     return output_dir, output_mask_dir, output_hdr_dir
 
-def generate_planar_projections_from_equirectangular_2(
-=======
-                    pers_image = equi2pers(im, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) * 255.0
-                    assert isinstance(pers_image, torch.Tensor)
-                    pers_image = (pers_image.permute(1, 2, 0)).type(torch.uint8).to("cpu").numpy()
-                    cv2.imwrite(f"{output_dir}/{i[:-4]}_{count}.png", pers_image)
-                    count += 1
-
-    return output_dir
-
 def generate_planar_projections_from_equirectangular_GT(
->>>>>>> origin/jmchen
     metadata_path: Path,
     image_dir: Path,
     planar_image_size: Tuple[int, int],

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/scripts/viewer/run_viewer.py

@@ -21,6 +21,7 @@ from __future__ import annotations
 import time
 from dataclasses import dataclass, field, fields
 from pathlib import Path
+
 from typing import Literal
 
 import tyro
@@ -112,9 +113,8 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int):
 
     assert viewer_state and pipeline.datamanager.train_dataset
     viewer_state.init_scene(
-        train_dataset=pipeline.datamanager.train_dataset,
+        dataset=pipeline.datamanager.train_dataset,
         train_state="completed",
-        eval_dataset=pipeline.datamanager.eval_dataset,
     )
     if isinstance(viewer_state, ViewerState):
         viewer_state.viser_server.set_training_state("completed")
@@ -133,4 +133,4 @@ if __name__ == "__main__":
     entrypoint()
 
 # For sphinx docs
-get_parser_fn = lambda: tyro.extras.get_parser(RunViewer)  # noqa
+get_parser_fn = lambda: tyro.extras.get_parser(RunViewer)  # noqa

nerfstudio/utils/colormaps.py

@@ -43,11 +43,6 @@ class ColormapOptions:
     """ Whether to invert the output colormap """
     exposure_scale: float = 0
     """Scale the brightness prior to tonemapping by 2 ^ Exposure"""
-<<<<<<< HEAD
-
-
-=======
->>>>>>> origin/jmchen
 
 def apply_colormap(
     image: Float[Tensor, "*bs channels"],
@@ -70,34 +65,19 @@ def apply_colormap(
 
     # default for rgb images
     if image.shape[-1] == 3:
-<<<<<<< HEAD
-        # print('colormap_options.exposure_scale', colormap_options.exposure_scale)
-        # if(colormap_options.exposure_scale != 0):
-        #     img_gamma_22 = image ** (2.2) - 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.)
-        # else:
-        #     return image
-        return image
-=======
-        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.)
         else:
             return image
->>>>>>> origin/jmchen
 
     # rendering depth outputs
     if image.shape[-1] == 1 and torch.is_floating_point(image):

nerfstudio/viewer/server/viewer_state.py

@@ -440,4 +440,4 @@ class ViewerState:
     def training_complete(self) -> None:
         """Called when training is complete."""
         self.training_state = "completed"
-        self.viser_server.set_training_state("completed")
+        self.viser_server.set_training_state("completed")