diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 08e820383..1801bd0c1 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -19,38 +19,44 @@ Put all the method implementations in one location. from __future__ import annotations from collections import OrderedDict -from typing import Dict from dataclasses import dataclass, field +from typing import Dict 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.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.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.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.data.datasets.hdr_dataset import HDRInputDataset - -from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig -from nerfstudio.engine.schedulers import ( - CosineDecaySchedulerConfig, - ExponentialDecaySchedulerConfig, - MultiStepSchedulerConfig, -) +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 @@ -70,6 +76,7 @@ method_configs: Dict[str, TrainerConfig] = {} descriptions = { "nerfacto": "Recommended real-time model tuned for real captures. This model will be continually updated.", "nerfacto-HDR": "Nerfacto that take HDR as input, and output HDR field.", + "lantern-nerfacto": "Nerfacto that take HDR as input, and output HDR field.", "depth-nerfacto": "Nerfacto with depth supervision.", "volinga": "Real-time rendering model from Volinga. Directly exportable to NVOL format at https://volinga.ai/", "instant-ngp": "Implementation of Instant-NGP. Recommended real-time model for unbounded scenes.", @@ -118,38 +125,8 @@ 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"] = lantern_config + method_configs["nerfacto-big"] = TrainerConfig( method_name="nerfacto", steps_per_eval_batch=500, diff --git a/nerfstudio/lantern/__init__.py b/nerfstudio/lantern/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/nerfstudio/lantern/config.py b/nerfstudio/lantern/config.py new file mode 100644 index 000000000..3da27e762 --- /dev/null +++ b/nerfstudio/lantern/config.py @@ -0,0 +1,55 @@ +""" +Nerfstudio Config + +Define custom method here that registers with Nerfstudio CLI. +""" + +from __future__ import annotations + +from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig +from nerfstudio.configs.base_config import ViewerConfig +from nerfstudio.data.dataparsers.nerfstudio_dataparser import \ + NerfstudioDataParserConfig +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", + 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=LanternModelConfig(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", + ) \ No newline at end of file diff --git a/nerfstudio/lantern/datamanager.py b/nerfstudio/lantern/datamanager.py new file mode 100644 index 000000000..9af41dccf --- /dev/null +++ b/nerfstudio/lantern/datamanager.py @@ -0,0 +1,22 @@ +""" +Lantern HDR DataManager +""" + +from dataclasses import dataclass, field +from typing import Dict, Literal, Tuple, Type, Union + +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 + + +@dataclass +class HDRVanillaDataManagerConfig(VanillaDataManagerConfig): + """A HDR data manager, based on VanillaDataManager""" + + _target: Type = field(default_factory=lambda: VanillaDataManager[HDRInputDataset]) + """Target class to instantiate.""" + \ No newline at end of file diff --git a/nerfstudio/lantern/dataset.py b/nerfstudio/lantern/dataset.py new file mode 100644 index 000000000..2f203753c --- /dev/null +++ b/nerfstudio/lantern/dataset.py @@ -0,0 +1,140 @@ +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +HDR Dataset of .exr format +""" +from __future__ import annotations + +import os +from copy import deepcopy +from pathlib import Path +from typing import Dict, List + +os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" + +import cv2 +import numpy as np +import numpy.typing as npt +import torch +from jaxtyping import Float +from PIL import Image +from torch import Tensor +from torch.utils.data import Dataset + +from nerfstudio.cameras.cameras import Cameras +from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs +from nerfstudio.data.datasets.base_dataset import InputDataset +from nerfstudio.data.utils.data_utils import get_image_mask_tensor_from_path + + +class HDRInputDataset(InputDataset): + """Dataset that returns HDR images, OpenEXR. + + Args: + dataparser_outputs: description of where and how to read input images. + scale_factor: The scaling factor for the dataparser outputs + """ + + # exclude_batch_keys_from_device: List[str] = ["image", "mask"] + # cameras: Cameras + + def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0): + super().__init__(dataparser_outputs, scale_factor) + + def __len__(self): + return len(self._dataparser_outputs.image_filenames) + + def get_numpy_image(self, image_idx: int) -> npt.NDArray[np.uint8]: + """Returns the image of shape (H, W, 3 or 4), with np.float32 values. + + Args: + image_idx: The image index in the dataset. + """ + image_filename = self._dataparser_outputs.image_filenames[image_idx] + hdr_image = cv2.imread(str(image_filename), cv2.IMREAD_UNCHANGED) + if self.scale_factor != 1.0: + # Here cv2 img shape is (height, width, BGR) + height, width = hdr_image.shape[:2] + newsize = (int(width * self.scale_factor), int(height * self.scale_factor)) + hdr_image = cv2.resize(hdr_image, newsize, interpolation = cv2.INTER_LINEAR) + hdr_image = cv2.cvtColor(hdr_image, cv2.COLOR_BGR2RGB) + hdr_image = hdr_image.astype("float32") # shape is (h, w) or (h, w, 3 or 4) + if len(hdr_image.shape) == 2: + hdr_image = hdr_image[:, :, None].repeat(3, axis=2) + assert len(hdr_image.shape) == 3 + assert hdr_image.dtype == np.float32 + assert hdr_image.shape[2] in [3, 4], f"Image shape of {hdr_image.shape} is in correct." + return hdr_image + + def get_image(self, image_idx: int) -> Float[Tensor, "image_height image_width num_channels"]: + """Returns a 3 channel image. + + Args: + image_idx: The image index in the dataset. + """ + # 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. + image = torch.log(1. + u * image) / torch.log(torch.tensor(1.+u)) + + if self._dataparser_outputs.alpha_color is not None and image.shape[-1] == 4: + image = image[:, :, :3] * image[:, :, -1:] + self._dataparser_outputs.alpha_color * (1.0 - image[:, :, -1:]) + return image + + def get_data(self, image_idx: int) -> Dict: + """Returns the ImageDataset data as a dictionary. + + Args: + image_idx: The image index in the dataset. + """ + image = self.get_image(image_idx) + data = {"image_idx": image_idx, "image": image} + if self._dataparser_outputs.mask_filenames is not None: + mask_filepath = self._dataparser_outputs.mask_filenames[image_idx] + data["mask"] = get_image_mask_tensor_from_path(filepath=mask_filepath, scale_factor=self.scale_factor) + assert ( + data["mask"].shape[:2] == data["image"].shape[:2] + ), f"Mask and image have different shapes. Got {data['mask'].shape[:2]} and {data['image'].shape[:2]}" + metadata = self.get_metadata(data) + data.update(metadata) + return data + + def get_metadata(self, data: Dict) -> Dict: + """Method that can be used to process any additional metadata that may be part of the model inputs. + + Args: + image_idx: The image index in the dataset. + """ + del data + return {} + + def __getitem__(self, image_idx: int) -> Dict: + data = self.get_data(image_idx) + return data + + @property + def image_filenames(self) -> List[Path]: + """ + Returns image filenames for this dataset. + The order of filenames is the same as in the Cameras object for easy mapping. + """ + + return self._dataparser_outputs.image_filenames diff --git a/nerfstudio/lantern/field.py b/nerfstudio/lantern/field.py new file mode 100644 index 000000000..ea1be7334 --- /dev/null +++ b/nerfstudio/lantern/field.py @@ -0,0 +1,71 @@ +""" +HDR Nerfstudio Field + +Currently this subclasses the NerfactoField. Consider subclassing the base Field. +""" + +from typing import Literal, Optional + +from torch import Tensor, nn + +from nerfstudio.field_components.mlp import MLP +from nerfstudio.field_components.spatial_distortions import SpatialDistortion +from nerfstudio.fields.nerfacto_field import \ + NerfactoField # for subclassing NerfactoField + + +class LanternNerfactoField(NerfactoField): + """Template Field + + Args: + aabb: parameters of scene aabb bounds + num_images: number of images in the dataset + """ + + aabb: Tensor + + def __init__( + self, + aabb: Tensor, + num_images: int, + num_layers: int = 2, + hidden_dim: int = 64, + geo_feat_dim: int = 15, + num_levels: int = 16, + base_res: int = 16, + max_res: int = 2048, + log2_hashmap_size: int = 19, + num_layers_color: int = 3, + num_layers_transient: int = 2, + features_per_level: int = 2, + hidden_dim_color: int = 64, + hidden_dim_transient: int = 64, + appearance_embedding_dim: int = 32, + transient_embedding_dim: int = 16, + use_transient_embedding: bool = False, + use_semantics: bool = False, + num_semantic_classes: int = 100, + pass_semantic_gradients: bool = False, + use_pred_normals: bool = False, + use_average_appearance_embedding: bool = False, + spatial_distortion: Optional[SpatialDistortion] = None, + implementation: Literal["tcnn", "torch"] = "tcnn", + ) -> None: + super().__init__(aabb, num_images, num_layers, hidden_dim, + geo_feat_dim, num_levels, base_res, + max_res, log2_hashmap_size, num_layers_color, + num_layers_transient, features_per_level, hidden_dim_color, + hidden_dim_transient, appearance_embedding_dim, transient_embedding_dim, + use_transient_embedding, use_semantics, num_semantic_classes, + pass_semantic_gradients, use_pred_normals, use_average_appearance_embedding, + spatial_distortion, implementation) + # Define the last part MLP for RGB outputs + self.mlp_head = MLP( + in_dim=self.direction_encoding.get_out_dim() + self.geo_feat_dim + self.appearance_embedding_dim, + num_layers=num_layers_color, + layer_width=hidden_dim_color, + out_dim=3, + activation=nn.ReLU(), + out_activation=nn.ReLU(), + implementation=implementation, + ) \ No newline at end of file diff --git a/nerfstudio/lantern/model.py b/nerfstudio/lantern/model.py new file mode 100644 index 000000000..e11f17101 --- /dev/null +++ b/nerfstudio/lantern/model.py @@ -0,0 +1,190 @@ +""" +Lantern Model File + +Currently this subclasses the Nerfacto model. Consider subclassing from the base Model. +""" +from dataclasses import dataclass, field +from typing import Dict, List, Literal, Tuple, Type + +import numpy as np +import torch +from torchmetrics.functional import structural_similarity_index_measure +from torchmetrics.image import PeakSignalNoiseRatio +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity + +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) +from nerfstudio.model_components.ray_samplers import (ProposalNetworkSampler, + UniformSampler) +from nerfstudio.model_components.renderers import (AccumulationRenderer, + DepthRenderer, + NormalsRenderer, + RGBRenderer) +from nerfstudio.model_components.scene_colliders import NearFarCollider +from nerfstudio.model_components.shaders import NormalsShader +from nerfstudio.models.base_model import Model, ModelConfig # for custom Model +from nerfstudio.models.nerfacto import ( # for subclassing Nerfacto model + NerfactoModel, NerfactoModelConfig) +from nerfstudio.utils import colormaps + + +@dataclass +class LanternModelConfig(NerfactoModelConfig): + """Template Model Configuration. + + Add your custom model config parameters here. + """ + + _target: Type = field(default_factory=lambda: LanternModel) + + +class LanternModel(NerfactoModel): + """Template Model.""" + + config: LanternModelConfig + + def populate_modules(self): + super().populate_modules() + + if self.config.disable_scene_contraction: + scene_contraction = None + else: + scene_contraction = SceneContraction(order=float("inf")) + + # Fields + self.field = LanternNerfactoField( + self.scene_box.aabb, + hidden_dim=self.config.hidden_dim, + num_levels=self.config.num_levels, + max_res=self.config.max_res, + base_res=self.config.base_res, + features_per_level=self.config.features_per_level, + log2_hashmap_size=self.config.log2_hashmap_size, + hidden_dim_color=self.config.hidden_dim_color, + hidden_dim_transient=self.config.hidden_dim_transient, + spatial_distortion=scene_contraction, + num_images=self.num_train_data, + use_pred_normals=self.config.predict_normals, + use_average_appearance_embedding=self.config.use_average_appearance_embedding, + appearance_embedding_dim=self.config.appearance_embed_dim, + implementation=self.config.implementation, + ) + + self.density_fns = [] + num_prop_nets = self.config.num_proposal_iterations + # Build the proposal network(s) + self.proposal_networks = torch.nn.ModuleList() + if self.config.use_same_proposal_network: + assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed." + prop_net_args = self.config.proposal_net_args_list[0] + network = HashMLPDensityField( + self.scene_box.aabb, + spatial_distortion=scene_contraction, + **prop_net_args, + implementation=self.config.implementation, + ) + self.proposal_networks.append(network) + self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)]) + else: + for i in range(num_prop_nets): + prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)] + network = HashMLPDensityField( + self.scene_box.aabb, + spatial_distortion=scene_contraction, + **prop_net_args, + implementation=self.config.implementation, + ) + self.proposal_networks.append(network) + self.density_fns.extend([network.density_fn for network in self.proposal_networks]) + + # Samplers + def update_schedule(step): + return np.clip( + np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]), + 1, + self.config.proposal_update_every, + ) + + # Change proposal network initial sampler if uniform + initial_sampler = None # None is for piecewise as default (see ProposalNetworkSampler) + if self.config.proposal_initial_sampler == "uniform": + initial_sampler = UniformSampler(single_jitter=self.config.use_single_jitter) + + self.proposal_sampler = ProposalNetworkSampler( + num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray, + num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray, + num_proposal_network_iterations=self.config.num_proposal_iterations, + single_jitter=self.config.use_single_jitter, + update_sched=update_schedule, + initial_sampler=initial_sampler, + ) + + # Collider + 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_accumulation = AccumulationRenderer() + self.renderer_depth = DepthRenderer() + self.renderer_normals = NormalsRenderer() + + # shaders + self.normals_shader = NormalsShader() + + # TODO losses + self.rgb_loss = MSELoss() + + # metrics + self.psnr = PeakSignalNoiseRatio(data_range=1.0) + self.ssim = structural_similarity_index_measure + self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) + + def get_image_metrics_and_images( + self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor] + ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: + gt_rgb = batch["image"].to(self.device) + predicted_rgb = outputs["rgb"] # Blended with background (black if random background) + gt_rgb = self.renderer_rgb.blend_background(gt_rgb) + acc = colormaps.apply_colormap(outputs["accumulation"]) + depth = colormaps.apply_depth_colormap( + outputs["depth"], + accumulation=outputs["accumulation"], + ) + + combined_rgb = torch.cat([gt_rgb, predicted_rgb], dim=1) + combined_acc = torch.cat([acc], dim=1) + combined_depth = torch.cat([depth], dim=1) + + # Switch images from [H, W, C] to [1, C, H, W] for metrics computations + gt_rgb_tonemapped = torch.clamp(colormaps.apply_colormap(gt_rgb), 0, 1) + predicted_rgb_tonemapped = torch.clamp(colormaps.apply_colormap(predicted_rgb), 0, 1) + + gt_rgb_tonemapped = torch.moveaxis(gt_rgb_tonemapped, -1, 0)[None, ...] + predicted_rgb_tonemapped = torch.moveaxis(predicted_rgb_tonemapped, -1, 0)[None, ...] + gt_rgb = torch.moveaxis(gt_rgb, -1, 0)[None, ...] + predicted_rgb = torch.moveaxis(predicted_rgb, -1, 0)[None, ...] + + psnr = self.psnr(gt_rgb, predicted_rgb) + ssim = self.ssim(gt_rgb, predicted_rgb) + # tonemapping before computing lpips + lpips = self.lpips(gt_rgb_tonemapped, predicted_rgb_tonemapped) + + # all of these metrics will be logged as scalars + metrics_dict = {"psnr": float(psnr.item()), "ssim": float(ssim)} # type: ignore + metrics_dict["lpips"] = float(lpips) + + images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth} + + for i in range(self.config.num_proposal_iterations): + key = f"prop_depth_{i}" + prop_depth_i = colormaps.apply_depth_colormap( + outputs[key], + accumulation=outputs["accumulation"], + ) + images_dict[key] = prop_depth_i + + return metrics_dict, images_dict \ No newline at end of file diff --git a/nerfstudio/process_data/aligned_pano2plane.py b/nerfstudio/process_data/aligned_pano2plane.py index 3400801ce..429568447 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 Optional, Literal, Tuple +from typing import Literal, Optional, 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.rich_utils import CONSOLE +from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import \ + BaseConverterToNerfstudioDataset from nerfstudio.utils import io +from nerfstudio.utils.rich_utils import CONSOLE + @dataclass class ProcessAlignedPano(BaseConverterToNerfstudioDataset): @@ -69,11 +69,10 @@ class ProcessAlignedPano(BaseConverterToNerfstudioDataset): summary_log = [] pers_size = equirect_utils.compute_resolution_from_equirect(self.data, self.images_per_equirect) CONSOLE.log(f"Generating {self.images_per_equirect} {pers_size} sized images per equirectangular image") - # self.data = equirect_utils.generate_planar_projections_from_equirectangular( - # self.metadata, self.data, pers_size, self.images_per_equirect, crop_factor=self.crop_factor - # ) + self.data = equirect_utils.generate_planar_projections_from_equirectangular_GT( + self.metadata, self.data, pers_size, self.images_per_equirect, crop_factor=self.crop_factor, clip_output = False + ) self.camera_type = "perspective" - self.data = self.data / "planar_projections" metadata_dict = io.load_from_json(self.data / "transforms.json") # Copy images to output directory cropped_images_filename = [] diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py index 9cf1ae487..3c9cd07b2 100644 --- a/nerfstudio/process_data/equirect_utils.py +++ b/nerfstudio/process_data/equirect_utils.py @@ -15,32 +15,27 @@ """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 -from numpy.linalg import inv - -from scipy.spatial.transform import Rotation import torch -from equilib import Equi2Pers -from rich.progress import ( - BarColumn, - Progress, - TaskProgressColumn, - TextColumn, - TimeRemainingColumn, -) +from numpy.linalg import inv +from rich.progress import (BarColumn, Progress, TaskProgressColumn, TextColumn, + TimeRemainingColumn) +from scipy.spatial.transform import Rotation -from nerfstudio.utils.rich_utils import CONSOLE, ItersPerSecColumn +from equilib import Equi2Pers 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]: @@ -119,8 +114,11 @@ 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), @@ -184,6 +182,7 @@ 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: @@ -197,6 +196,8 @@ 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"), @@ -208,6 +209,7 @@ 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") @@ -241,10 +243,17 @@ 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 @@ -278,12 +287,24 @@ 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], samples_per_im: int, mask_dir: Path = None, crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), + clip_output: bool = False, ) -> Path: """Given camera pose, generate planar projections from an equirectangular image. And output corresponding camera pose. @@ -381,7 +402,7 @@ def generate_planar_projections_from_equirectangular_2( 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 - pers_image = equi2pers(im, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) + pers_image = equi2pers(im, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}, clip_output=clip_output) if mask_dir is not None: pers_mask = equi2pers(mask, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) * 255.0 # transform matrix for blender: object.matrix_world diff --git a/nerfstudio/utils/colormaps.py b/nerfstudio/utils/colormaps.py index 4aab6eb2a..a63d942a7 100644 --- a/nerfstudio/utils/colormaps.py +++ b/nerfstudio/utils/colormaps.py @@ -43,8 +43,11 @@ 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"], @@ -67,6 +70,7 @@ 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. @@ -76,6 +80,24 @@ def apply_colormap( # else: # return image return image +======= + if(colormap_options.exposure_scale != 0 or image.max() > 1.): + # Log + 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):