diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index fc4425894..92e503b4e 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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( diff --git a/nerfstudio/lantern/__init__.py b/nerfstudio/lantern/__init__.py index e69de29bb..84952a816 100644 --- a/nerfstudio/lantern/__init__.py +++ b/nerfstudio/lantern/__init__.py @@ -0,0 +1 @@ +# init \ No newline at end of file diff --git a/nerfstudio/lantern/config.py b/nerfstudio/lantern/config.py index 3da27e762..cf1a86d02 100644 --- a/nerfstudio/lantern/config.py +++ b/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, diff --git a/nerfstudio/lantern/datamanager.py b/nerfstudio/lantern/datamanager.py index 9af41dccf..90d8885aa 100644 --- a/nerfstudio/lantern/datamanager.py +++ b/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 diff --git a/nerfstudio/lantern/dataset.py b/nerfstudio/lantern/dataset.py index 2f203753c..d0c842408 100644 --- a/nerfstudio/lantern/dataset.py +++ b/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. diff --git a/nerfstudio/lantern/model.py b/nerfstudio/lantern/model.py index e11f17101..96eafcc6c 100644 --- a/nerfstudio/lantern/model.py +++ b/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) diff --git a/nerfstudio/lantern/renderer.py b/nerfstudio/lantern/renderer.py new file mode 100644 index 000000000..04943d50b --- /dev/null +++ b/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 diff --git a/nerfstudio/process_data/aligned_pano2plane.py b/nerfstudio/process_data/aligned_pano2plane.py index 429568447..6d2b647ad 100644 --- a/nerfstudio/process_data/aligned_pano2plane.py +++ b/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): diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py index 17b5d8da4..7d5651d00 100644 --- a/nerfstudio/process_data/equirect_utils.py +++ b/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]: diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index 5a3ca78bd..4f8c18627 100644 --- a/nerfstudio/scripts/render.py +++ b/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""" diff --git a/nerfstudio/utils/colormaps.py b/nerfstudio/utils/colormaps.py index b2477ad85..89ddba073 100644 --- a/nerfstudio/utils/colormaps.py +++ b/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.)