diff --git a/lantern_scripts/combine_LDR_HDR.py b/lantern_scripts/combine_LDR_HDR.py index 1c29e0cd2..577efe288 100644 --- a/lantern_scripts/combine_LDR_HDR.py +++ b/lantern_scripts/combine_LDR_HDR.py @@ -5,9 +5,9 @@ from envmap import EnvironmentMap, rotation_matrix from hdrio import imwrite # Outdoor Table -LDR_dir = '/mnt/data/arch_level1/panoramas_first_100/' -HDR_dir = '/mnt/data/arch_level1/panoramas_256x512_hdr/' -output_dir = '/mnt/data/arch_level1/panoramas_LDR_HDR_small/' +LDR_dir = '/mnt/data/nerfstudio_ds/real_data/lab_downstairs_data/left_e1_linearized/' +HDR_dir = '/mnt/data/nerfstudio_ds/real_data/lab_downstairs_data/left_e1_ldr2hdr_network/' +output_dir = '/mnt/data/nerfstudio_ds/real_data/lab_downstairs_data/left_e1_ldr2hdr_network_high_res/' imgFiles = [os.path.join(LDR_dir, f) for f in os.listdir(LDR_dir) if os.path.isfile(os.path.join(LDR_dir, f)) and f.endswith('.png')] imgFiles.sort() @@ -21,7 +21,7 @@ for ldr_addr in tqdm(imgFiles): hdr_pano = EnvironmentMap(hdr_addr, 'latlong') hdr_pano.resize(1920) - out_pano = ldr_pano.data ** 2.2 + out_pano = ldr_pano.data # ** 2.2 list_hdr = hdr_pano.data > 1.0 print("list_hdr.shape: ", list_hdr.shape, "list_hdr.shape: ", ldr_pano.data.shape) out_pano[list_hdr] = hdr_pano.data[list_hdr] diff --git a/lantern_scripts/train_test_split.py b/lantern_scripts/train_test_split.py new file mode 100644 index 000000000..34d36e285 --- /dev/null +++ b/lantern_scripts/train_test_split.py @@ -0,0 +1,80 @@ +import argparse +import os +from tqdm import tqdm +from pathlib import Path +import json +import numpy as np +import math + + +def load_from_json(filename: Path): + """Load a dictionary from a JSON filename. + + Args: + filename: The filename to load from. + """ + assert filename.suffix == ".json" + with open(filename, encoding="UTF-8") as file: + return json.load(file) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument('--input_json', type=str) + parser.add_argument('--output_json', type=str) + parser.add_argument('--train_split_fraction', type=float, default= 0.8) + parser.add_argument('--eval_split_fraction', type=float, default= 0.1) + args = parser.parse_args() + + assert os.path.isfile(args.input_json) and args.input_json.endswith('.json'), 'The input JSON file is missing!' + + json_data = load_from_json( Path(args.input_json) ) + + + new_json_data = {} + all_keys = json_data.keys() + for key in all_keys: + new_json_data[key] = json_data[key] + + + num_images = len(json_data['frames']) + num_train_images = math.ceil(num_images * args.train_split_fraction) + num_eval_test_images = num_images - num_train_images + i_all = np.arange(num_images) + i_train = np.linspace( + 0, num_images - 1, num_train_images, dtype=int + ) # equally spaced training images starting and ending at 0 and num_images-1 + i_eval_test = np.setdiff1d(i_all, i_train) # eval images are the remaining images + assert len(i_eval_test) == num_eval_test_images + + num_eval_images = math.ceil(num_images * args.eval_split_fraction) + num_test_images = num_eval_test_images - num_train_images + i_eval = np.linspace( + 0, num_eval_test_images - 1, num_eval_images, dtype=int + ) # equally spaced training images starting and ending at 0 and num_images-1 + i_eval = i_eval_test[i_eval] + i_test = np.setdiff1d(i_eval_test, i_eval) # eval images are the remaining images + assert len(i_eval) == num_eval_images + + + # has_split_files_spec = any(f"{split}_filenames" in meta for split in ("train", "val", "test")) + + train_list = [] + for i in i_train: + train_list.append(json_data['frames'][i]["file_path"]) + new_json_data["train_filenames"] = train_list + + test_list = [] + for i in i_test: + test_list.append(json_data['frames'][i]["file_path"]) + new_json_data["test_filenames"] = test_list + + eval_list = [] + for i in i_eval: + eval_list.append(json_data['frames'][i]["file_path"]) + new_json_data["val_filenames"] = eval_list + + with open(args.output_json, "w") as outfile: + json.dump(new_json_data, outfile, indent=4) + + print("All done!") \ No newline at end of file diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 78be90da0..b00d29d8a 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -50,6 +50,9 @@ class PixelSamplerConfig(InstantiateConfig): is_equirectangular: bool = False """List of whether or not camera i is equirectangular.""" + lantern_steps: int = 1 + """ Define which steps of lantern is going on!! """ + class PixelSampler: """Samples 'pixel_batch's from 'image_batch's. @@ -155,7 +158,7 @@ class PixelSampler: chosen_indices = torch.randint(0, self.light_sources_indices.shape[0], (batch_size,)) return self.light_sources_indices[chosen_indices].to(device) else: - assert validity is not None, "Validy needs to be defined for the lught sources samples" + assert validity is not None, "Validy needs to be defined for the light sources samples" validity = validity.squeeze(-1).cpu() mask_combine = torch.logical_and(mask, validity) indices, nonzero_indices_whole = self.sample_considering_mask(mask_combine, device, batch_size) @@ -239,11 +242,19 @@ class PixelSampler: num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device ) else: - indices = self.sample_method( - num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device, - is_fast_expo=is_fast_expo, - only_light_sources=only_light_sources, - validity=batch["saturation_mask"] + if "saturation_mask" in batch: + indices = self.sample_method( + num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device, + is_fast_expo=is_fast_expo, + only_light_sources=only_light_sources, + validity=batch["saturation_mask"] + ) + else: + indices = self.sample_method( + num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device, + is_fast_expo=is_fast_expo, + only_light_sources=only_light_sources, + validity=None ) else: if self.config.is_equirectangular: @@ -367,19 +378,33 @@ class PixelSampler: for key in image_batch.keys(): self.well_image_batch[key] = image_batch[key][well_ones] self.fast_image_batch[key] = image_batch[key][fast_ones] - well_pixel_batch = self.collate_image_dataset_batch( - self.well_image_batch, int(0.2 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image - ) + + assert self.config.lantern_steps == 1 or self.config.lantern_steps == 2, "The step of Lantern should be either 1 or 2!! " + + if self.config.lantern_steps == 1: + well_pixel_batch = self.collate_image_dataset_batch( + self.well_image_batch, int(1.0 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image + ) + fast_pixel_batch = self.collate_image_dataset_batch( + self.fast_image_batch, int(0.0 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image, is_fast_expo=True + ) + light_pixel_batch = self.collate_image_dataset_batch( + self.fast_image_batch, int(0.0 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image, is_fast_expo=True, only_light_sources=True + ) + else: + well_pixel_batch = self.collate_image_dataset_batch( + self.well_image_batch, int(0.2 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image + ) + fast_pixel_batch = self.collate_image_dataset_batch( + self.fast_image_batch, int(0.3 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image, is_fast_expo=True + ) + light_pixel_batch = self.collate_image_dataset_batch( + self.fast_image_batch, int(0.5 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image, is_fast_expo=True, only_light_sources=True + ) + pixel_batch = well_pixel_batch #.copy() - fast_pixel_batch = self.collate_image_dataset_batch( - self.fast_image_batch, int(0.3 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image, is_fast_expo=True - ) for key, value in fast_pixel_batch.items(): pixel_batch[key] = torch.cat((pixel_batch[key], value), 0) - - light_pixel_batch = self.collate_image_dataset_batch( - self.fast_image_batch, int(0.5 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image, is_fast_expo=True, only_light_sources=True - ) for key, value in light_pixel_batch.items(): pixel_batch[key] = torch.cat((pixel_batch[key], value), 0) else: diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 5a549369b..580d15f7a 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -349,12 +349,4 @@ class NerfactoField(Field): rgb = self.mlp_head(h).view(*outputs_shape, -1).to(directions) outputs.update({FieldHeadNames.RGB: rgb}) - rgb_fast = self.mlp_head_fast(h).view(*outputs_shape, -1).to(directions) - outputs.update({FieldHeadNames.RGB_FAST: rgb_fast}) - - if self.predicts_validity: - h_no_grad = h.clone() - validity = self.mlp_validity(h_no_grad).view(*outputs_shape, -1).to(directions) - outputs.update({FieldHeadNames.VALIDITY: validity}) - return outputs diff --git a/nerfstudio/lantern/field.py b/nerfstudio/lantern/field.py index 20f26bb2f..9dbfe31b8 100644 --- a/nerfstudio/lantern/field.py +++ b/nerfstudio/lantern/field.py @@ -4,12 +4,28 @@ HDR Nerfstudio Field Currently this subclasses the NerfactoField. Consider subclassing the base Field. """ -from typing import Literal, Optional +from typing import Dict, Literal, Optional, Tuple +import torch from torch import Tensor, nn +from nerfstudio.cameras.rays import RaySamples +from nerfstudio.data.scene_box import SceneBox +from nerfstudio.field_components.activations import trunc_exp +from nerfstudio.field_components.embedding import Embedding +from nerfstudio.field_components.encodings import HashEncoding, NeRFEncoding, SHEncoding +from nerfstudio.field_components.encodings import Encoding, Identity +from nerfstudio.field_components.field_heads import ( + FieldHeadNames, + PredNormalsFieldHead, + SemanticFieldHead, + TransientDensityFieldHead, + TransientRGBFieldHead, + UncertaintyFieldHead, +) from nerfstudio.field_components.mlp import MLP from nerfstudio.field_components.spatial_distortions import SpatialDistortion +from nerfstudio.fields.base_field import Field, get_normalized_directions from nerfstudio.fields.nerfacto_field import \ NerfactoField # for subclassing NerfactoField @@ -107,4 +123,97 @@ class LanternNerfactoField(NerfactoField): # out_activation=nn.ReLU(), out_activation=nn.Sigmoid(), implementation=implementation, - ) \ No newline at end of file + ) + + + def get_outputs( + self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None + ) -> Dict[FieldHeadNames, Tensor]: + assert density_embedding is not None + outputs = {} + if ray_samples.camera_indices is None: + raise AttributeError("Camera indices are not provided.") + camera_indices = ray_samples.camera_indices.squeeze() + directions = get_normalized_directions(ray_samples.frustums.directions) + directions_flat = directions.view(-1, 3) + d = self.direction_encoding(directions_flat) + + outputs_shape = ray_samples.frustums.directions.shape[:-1] + + # appearance + if self.use_appearance_embedding: + if self.training: + embedded_appearance = self.embedding_appearance(camera_indices) + else: + if self.use_average_appearance_embedding: + embedded_appearance = torch.ones( + (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device + ) * self.embedding_appearance.mean(dim=0) + else: + embedded_appearance = torch.zeros( + (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device + ) + + # transients + if self.use_transient_embedding and self.training: + embedded_transient = self.embedding_transient(camera_indices) + transient_input = torch.cat( + [ + density_embedding.view(-1, self.geo_feat_dim), + embedded_transient.view(-1, self.transient_embedding_dim), + ], + dim=-1, + ) + x = self.mlp_transient(transient_input).view(*outputs_shape, -1).to(directions) + outputs[FieldHeadNames.UNCERTAINTY] = self.field_head_transient_uncertainty(x) + outputs[FieldHeadNames.TRANSIENT_RGB] = self.field_head_transient_rgb(x) + outputs[FieldHeadNames.TRANSIENT_DENSITY] = self.field_head_transient_density(x) + + # semantics + if self.use_semantics: + semantics_input = density_embedding.view(-1, self.geo_feat_dim) + if not self.pass_semantic_gradients: + semantics_input = semantics_input.detach() + + x = self.mlp_semantics(semantics_input).view(*outputs_shape, -1).to(directions) + outputs[FieldHeadNames.SEMANTICS] = self.field_head_semantics(x) + + # predicted normals + if self.use_pred_normals: + positions = ray_samples.frustums.get_positions() + + positions_flat = self.position_encoding(positions.view(-1, 3)) + pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) + + x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions) + outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x) + + if self.use_appearance_embedding: + h = torch.cat( + [ + d, + density_embedding.view(-1, self.geo_feat_dim), + embedded_appearance.view(-1, self.appearance_embedding_dim), + ], + dim=-1, + ) + else: + h = torch.cat( + [ + d, + density_embedding.view(-1, self.geo_feat_dim), + ], + dim=-1, + ) + rgb = self.mlp_head(h).view(*outputs_shape, -1).to(directions) + outputs.update({FieldHeadNames.RGB: rgb}) + + rgb_fast = self.mlp_head_fast(h).view(*outputs_shape, -1).to(directions) + outputs.update({FieldHeadNames.RGB_FAST: rgb_fast}) + + if self.predicts_validity: + h_no_grad = h.clone() + validity = self.mlp_validity(h_no_grad).view(*outputs_shape, -1).to(directions) + outputs.update({FieldHeadNames.VALIDITY: validity}) + + return outputs diff --git a/nerfstudio/lantern/model.py b/nerfstudio/lantern/model.py index 283e7a013..1fe6e82e1 100644 --- a/nerfstudio/lantern/model.py +++ b/nerfstudio/lantern/model.py @@ -60,6 +60,12 @@ class LanternModelConfig(NerfactoModelConfig): appearance_embed_dim: int = 32 + lantern_steps: int = -1 + """ Define which steps of lantern is going on!! """ + + apply_mu_law: bool = True + """ Define which steps of lantern is going on!! """ + class LanternModel(NerfactoModel): """Template Model.""" @@ -182,6 +188,11 @@ class LanternModel(NerfactoModel): validity = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.VALIDITY], weights=weights) rgb_fast = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB_FAST], weights=weights) + if self.config.apply_mu_law: + u = 5000. + rgb_fast = torch.exp(rgb_fast * torch.log(torch.tensor(u+1.))) - 1. + rgb_fast /= u + outputs = { "rgb": rgb, "rgb_fast": rgb_fast, @@ -222,21 +233,37 @@ class LanternModel(NerfactoModel): 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) + + if batch["exposure"] == 1.0: + predicted_rgb = outputs["rgb"] # Blended with background (black if random background) + else: + # if self.config.lantern_steps == 1: + return None, None + predicted_rgb = outputs["rgb_fast"] # Blended with background (black if random background) + # if self.config.apply_mu_law: + # from nerfstudio.utils.colormaps import ColormapOptions, Colormaps + # cm = ColormapOptions('inverse-mu') + # predicted_rgb = colormaps.apply_colormap(predicted_rgb, cm) + 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"], ) + gt_rgb = batch["mask"] * gt_rgb + predicted_rgb = batch["mask"] * predicted_rgb + print("Exposure: ", batch['exposure']) 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.clamp(colormaps.apply_colormap(gt_rgb), 0, 1) + gt_rgb_tonemapped = torch.clamp(gt_rgb, 0, 1) + # predicted_rgb_tonemapped = torch.clamp(colormaps.apply_colormap(predicted_rgb), 0, 1) + predicted_rgb_tonemapped = torch.clamp(predicted_rgb, 0, 1) predicted_rgb_tonemapped = torch.nan_to_num(predicted_rgb_tonemapped, nan=0.0) gt_rgb_tonemapped = torch.moveaxis(gt_rgb_tonemapped, -1, 0)[None, ...] predicted_rgb_tonemapped = torch.moveaxis(predicted_rgb_tonemapped, -1, 0)[None, ...] @@ -313,10 +340,6 @@ class LanternModel(NerfactoModel): image = batch["image"].to(self.device) if "exposure" in batch: - # u = 10. - # img_uncompress = torch.exp(image * torch.log(torch.tensor(u+1.))) - 1. - # img_uncompress /= u - # img_uncompress = torch.pow(image, 2.2) exposures_resized = batch["exposure"].view(batch["exposure"].shape[0], 1) # img_uncompress_re_exposed = exposures_resized * img_uncompress @@ -339,8 +362,6 @@ class LanternModel(NerfactoModel): gt_image=image, ) - # import pdb; pdb.set_trace() - # MSE loss for mask with weights # if "saturation_mask" in batch: # mask_in_float = batch['saturation_mask'].type(torch.float).to(self.device) @@ -351,37 +372,21 @@ class LanternModel(NerfactoModel): # loss_dict["validity_loss_well_exposed"] = loss_mask_well_expo # 0.01 # loss_dict["validity_loss_fast_exposure"] = loss_mask_fast_expo - gt_rgb_w = gt_rgb.clone() - # u = 5000. - # gt_rgb_w = torch.log(1. + u * gt_rgb_w) / torch.log(torch.tensor(1.+u)) - - # gt_rgb_w[(exposures_resized != 1.0).repeat(1, 3)] = 1.0# torch.clamp(gt_rgb_w[(exposures_resized != 1.0).repeat(1, 3)] / 0.004, 0, 1) - # import pdb; pdb.set_trace() - gt_rgb_f = gt_rgb.clone() - u = 5000. - gt_rgb_f = torch.log(1. + u * gt_rgb_f) / torch.log(torch.tensor(1.+u)) - - # gt_rgb_f[(exposures_resized == 1.0).repeat(1, 3)] = gt_rgb_f[(exposures_resized == 1.0).repeat(1, 3)] * 0.004 - # loss_dict["rgb_loss"] = self.rgb_loss(gt_rgb_w, pred_rgb) - # loss_dict["rgb_loss_fast"] = self.rgb_loss(gt_rgb_f, pred_rgb_f) - # loss_dict["rgb_loss"] = self.rgb_loss(gt_rgb_f, pred_rgb_f) + assert self.config.lantern_steps == 1 or self.config.lantern_steps == 2, "The step of Lantern should be either 1 or 2!! " - loss_dict["rgb_loss"] = (mask_w_w * ((gt_rgb_w - pred_rgb) ** 2)).mean() - loss_dict["rgb_loss_fast"] = (mask_f_w * ((gt_rgb_f - pred_rgb_f) ** 2)).mean() + if self.config.lantern_steps == 1: + loss_dict["rgb_loss"] = self.rgb_loss(gt_rgb, pred_rgb) + else: + gt_rgb_w = gt_rgb.clone() + + gt_rgb_f = gt_rgb.clone() + if self.config.apply_mu_law: + u = 5000. + gt_rgb_f = torch.log(1. + u * gt_rgb_f) / torch.log(torch.tensor(1.+u)) - # RGB_fast = (pred_rgb * pred_rgb_f) * 0.004 - # loss_dict["rgb_loss_fast"] = 0.0005 * (mask_f_w * ((gt_rgb_f - RGB_fast) ** 2)).mean() - # print("rgb_loss_fast: ", loss_dict["rgb_loss_fast"]) - - # loss_dict["rgb_loss_fast"] = 300.0 * (mask_f_w * ((gt_rgb_f - pred_rgb_f) ** 2)).mean() - # import pdb; pdb.set_trace() - - # if weights_for_RGB_loss is not None: - # loss_dict["rgb_loss"] = (weights_for_RGB_loss * ((gt_rgb - pred_rgb) ** 2)).mean() - # else: - # loss_dict["rgb_loss"] = self.rgb_loss(gt_rgb, pred_rgb) + loss_dict["rgb_loss"] = (mask_w_w * ((gt_rgb_w - pred_rgb) ** 2)).mean() + loss_dict["rgb_loss_fast"] = (mask_f_w * ((gt_rgb_f - pred_rgb_f) ** 2)).mean() - # print("rgb_loss: ", loss_dict["rgb_loss"], ", rgb_fast_expo_loss: ", loss_dict["rgb_fast_expo_loss"]) if self.training: loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( outputs["weights_list"], outputs["ray_samples_list"] diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 5e202b06c..75fc493cf 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -389,20 +389,21 @@ class VanillaPipeline(Pipeline): num_rays = height * width outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) metrics_dict, images_dict = self.model.get_image_metrics_and_images(outputs, batch) - - if output_path is not None: - camera_indices = camera_ray_bundle.camera_indices - assert camera_indices is not None - for key, val in images_dict.items(): - Image.fromarray((val * 255).byte().cpu().numpy()).save( - output_path / "{0:06d}-{1}.jpg".format(int(camera_indices[0, 0, 0]), key) - ) - assert "num_rays_per_sec" not in metrics_dict - metrics_dict["num_rays_per_sec"] = num_rays / (time() - inner_start) - fps_str = "fps" - assert fps_str not in metrics_dict - metrics_dict[fps_str] = metrics_dict["num_rays_per_sec"] / (height * width) - metrics_dict_list.append(metrics_dict) + + if metrics_dict: + if output_path is not None: + camera_indices = camera_ray_bundle.camera_indices + assert camera_indices is not None + for key, val in images_dict.items(): + Image.fromarray((val * 255).byte().cpu().numpy()).save( + output_path / "{0:06d}-{1}.jpg".format(int(camera_indices[0, 0, 0]), key) + ) + assert "num_rays_per_sec" not in metrics_dict + metrics_dict["num_rays_per_sec"] = num_rays / (time() - inner_start) + fps_str = "fps" + assert fps_str not in metrics_dict + metrics_dict[fps_str] = metrics_dict["num_rays_per_sec"] / (height * width) + metrics_dict_list.append(metrics_dict) progress.advance(task) # average the metrics list metrics_dict = {} diff --git a/nerfstudio/process_data/GT_HDR_to_nerfstudio.py b/nerfstudio/process_data/GT_HDR_to_nerfstudio.py new file mode 100644 index 000000000..09484450d --- /dev/null +++ b/nerfstudio/process_data/GT_HDR_to_nerfstudio.py @@ -0,0 +1,145 @@ +# 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. + +"""Processes an panorama sequence to a nerfstudio compatible dataset.""" + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Optional, Literal, Tuple +import os +import numpy as np +from numpy.linalg import inv +import cv2 +from tqdm import tqdm + + +from nerfstudio.process_data import equirect_utils, process_data_utils +from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import ( + BaseConverterToNerfstudioDataset, +) +from rich.progress import ( + BarColumn, + Progress, + TaskProgressColumn, + TextColumn, + TimeRemainingColumn, +) +from nerfstudio.utils.rich_utils import CONSOLE, ItersPerSecColumn +from nerfstudio.utils import io + + +def apply_correction(coefs, img, CORRECTION_CURVE_TYPE) -> np.ndarray: + # check that img is in range 0-255 + assert img.dtype == np.uint8 + assert CORRECTION_CURVE_TYPE in ['gamma'] + img = ((img.astype(np.float32) / 255.0) ** np.array(coefs).T) # * 255.0 + # img = ((img.astype(np.float32) / 255.0) ** coefs[None, None, :]) * 255.0 + return img + + +def distance(point_array, ref): + pow_2 = np.power(point_array - ref, 2) + sum_pow_2 = np.sum(pow_2, axis=-1) + dist = np.sqrt(sum_pow_2) + return dist + + +def opensfm_to_opengl(shot_data): + rvec = np.array(tuple(map(float, shot_data['rotation']))) + tvec = np.array(tuple(map(float, shot_data['translation']))) + + # Converting opencv to opengl coordinates + from scipy.spatial.transform import Rotation as R + rot = R.from_rotvec(rvec).as_matrix() + + W2C = np.eye(4) + W2C[:3, :3] = rot + W2C[:3, 3] = np.array(tvec) + + C2W = inv(W2C) + C2W[0:3, 1:3] *= -1 + C2W = C2W[np.array([1, 0, 2, 3]), :] + C2W[2, :] *= -1 + + return C2W, rvec, tvec + + +@dataclass +class GTHDRoNeRFStudioDataset(BaseConverterToNerfstudioDataset): + """Process panorama data into a nerfstudio dataset. Skip the COLMAP + + This script does the following: + + 1. Crops panoramas to planar images, and restore corresponding perspective camera poses. + 1. Scales images to a specified size. + 2. Converts perpective camera poses into the nerfstudio format. + """ + + metadata: Path = "" + """Path the metadata of the panoramas sequence.""" + + checkpoint: Path = "" + """Path the metadata of the panoramas sequence.""" + + + def main(self) -> None: + """Process images into a nerfstudio dataset.""" + + summary_log = [] + + openSFM_reconstruction = io.load_from_json(self.metadata) + frames_names = openSFM_reconstruction[0]['shots'].keys() + frames_names = sorted(frames_names) + + dataset_transform = io.load_from_json(self.checkpoint / "dataparser_transforms.json") + scale_factor = dataset_transform['scale'] + normalization_mat = np.eye(4) + normalization_mat[:3, :] = np.array(dataset_transform['transform']) + + # Convert OpenSFM coordinates to the NeRFStudio + camera_dict = {} + camera_dict["fps"] = 1 + camera_key = 'v2 unknown unknown -1 -1 perspective 0' + camera_dict["camera_type"] = "equirectangular" + camera_dict["render_width"] = openSFM_reconstruction[0]["cameras"][camera_key]['width'] + camera_dict["render_height"] = openSFM_reconstruction[0]["cameras"][camera_key]['height'] + camera_dict["smoothness_value"] = 0 + camera_dict["is_cycle"] = False + camera_dict["camera_path"] = [] + + frame_length = 0 + for idx, fname in enumerate([x for x in frames_names if x.startswith('GT')]): + print(fname) + shot_data = openSFM_reconstruction[0]['shots'][fname] + C2W, rvec, tvec = opensfm_to_opengl(shot_data) + C2W = normalization_mat @ C2W + C2W[:3, 3] *= scale_factor + camera_dict["camera_path"].append( + { + "camera_to_world": C2W.flatten().tolist(), + "fov": 20.407948880476418, + "aspect": 1, + }) + frame_length += 1 + camera_dict["seconds"] = float(frame_length) + + with open(self.output_dir / 'GT_transforms.json', 'w') as f: + json.dump(camera_dict, f, indent=4) + + CONSOLE.rule("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:") + + for summary in summary_log: + CONSOLE.print(summary, justify="center") + CONSOLE.rule() \ No newline at end of file diff --git a/nerfstudio/process_data/aligned_pano2plane.py b/nerfstudio/process_data/aligned_pano2plane.py index 0795115c0..4250e54af 100644 --- a/nerfstudio/process_data/aligned_pano2plane.py +++ b/nerfstudio/process_data/aligned_pano2plane.py @@ -28,6 +28,8 @@ from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import ( from nerfstudio.utils.rich_utils import CONSOLE from nerfstudio.utils import io +from nerfstudio.process_data.openSFM_to_nerfstudio import opensfm_to_opengl, distance + @dataclass class ProcessAlignedPano(BaseConverterToNerfstudioDataset): """Process panorama data into a nerfstudio dataset. Skip the COLMAP @@ -41,7 +43,7 @@ class ProcessAlignedPano(BaseConverterToNerfstudioDataset): metadata: Path = "" """Path the metadata of the panoramas sequence.""" - num_downscales: int = 3 + num_downscales: int = 0 """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3 will downscale the images by 2x, 4x, and 8x.""" max_dataset_size: int = 300 @@ -55,12 +57,13 @@ class ProcessAlignedPano(BaseConverterToNerfstudioDataset): """Portion of the image to crop. All values should be in [0,1]. (top, bottom, left, right)""" skip_image_processing: bool = False """If True, skips copying and downscaling of images and only runs COLMAP if possible and enabled""" - is_HDR: bool = False + is_HDR: bool = True """If True, process the .exr files as HDR images.""" - use_mask: bool = False + use_mask: bool = True """If True, process the .exr files with mask.""" - use_exposure: bool = False + is_metadata_from_openSFM: bool = False """If True, using different exposures.""" + def main(self) -> None: """Process images into a nerfstudio dataset.""" @@ -74,16 +77,66 @@ class ProcessAlignedPano(BaseConverterToNerfstudioDataset): mask_dir.mkdir(parents=True, exist_ok=True) 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_GT( - self.metadata, self.data, pers_size, - self.images_per_equirect, crop_factor=self.crop_factor, clip_output = False, - use_mask = self.use_mask, - use_exposure = self.use_exposure - ) + + if not self.skip_image_processing: + if self.is_metadata_from_openSFM: + openSFM_reconstruction = io.load_from_json(self.metadata) + frames_names = openSFM_reconstruction[0]['shots'].keys() + frames_names = sorted(frames_names) + + # Remove GT files from the training sets + frames_names = list(set(frames_names) - set( [x for x in frames_names if x.startswith('GT')])) + + # Convert OpenSFM coordinates to the NeRFStudio + translations = [] + for fname in frames_names: + # if fname.startswith("left_sfm"): + if fname.startswith("left_e1"): + shot_data = openSFM_reconstruction[0]['shots'][fname] + _, _, left_tvec = opensfm_to_opengl(shot_data) + translations.append(left_tvec) + + translations_np = np.asarray(translations) + left_trn_median = np.median(translations_np, 0) + dist = distance(translations_np, left_trn_median) + percentile_95 = np.percentile(dist, 99) + valid_idx = dist < percentile_95 + + # Convert OpenSFM coordinates to the NeRFStudio + camera_dict = {} + camera_dict["frames"] = [] + for idx, fname in enumerate([x for x in frames_names if x.startswith('left_e1')]): + shot_data = openSFM_reconstruction[0]['shots'][fname] + left_C2W, left_rvec, left_tvec = opensfm_to_opengl(shot_data) + if valid_idx[idx]: + camera_dict["frames"].append( + { + "file_path": fname.split('.')[0], + "transform_matrix": left_C2W.tolist() + }) + + with open(self.output_dir / 'panoramic_transforms.json', 'w') as f: + json.dump(camera_dict, f, indent=4) + self.metadata = self.output_dir / 'panoramic_transforms.json' + + + metadata_dict = io.load_from_json(self.metadata) + + pers_size = equirect_utils.compute_resolution_from_equirect(self.data / 'left_e1', self.images_per_equirect) + CONSOLE.log(f"Generating {self.images_per_equirect} {pers_size} sized images per equirectangular image") + out_dir = equirect_utils.generate_planar_projections_from_equirectangular_GT( + self.metadata, + self.data / 'left_e1', pers_size, self.images_per_equirect, + crop_factor=self.crop_factor, + clip_output = False, + use_mask = self.use_mask, + prefix = 'left_e1' + ) + else: + out_dir = self.data / 'left_e1' / "planar_projections" self.camera_type = "perspective" - metadata_dict = io.load_from_json(self.data / "transforms.json") + metadata_dict = io.load_from_json(out_dir / "transforms.json") + # Copy images to output directory cropped_images_filename = [] cropped_masks_filename = [] @@ -92,6 +145,7 @@ class ProcessAlignedPano(BaseConverterToNerfstudioDataset): if self.use_mask: cropped_masks_filename.append(Path(frame["mask_path"])) # Copy images to output directory + CONSOLE.log(f"Copying perspective images into the target directory") if self.is_HDR: copied_image_paths = process_data_utils.copy_images_list_EXR( cropped_images_filename, image_dir=image_dir, verbose=self.verbose, num_downscales=self.num_downscales diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py index a9544db5b..24fcd33c8 100644 --- a/nerfstudio/process_data/equirect_utils.py +++ b/nerfstudio/process_data/equirect_utils.py @@ -576,7 +576,7 @@ def generate_planar_projections_from_equirectangular_GT( frames = [] idx = 0 files_list = [file for file in os.listdir(frame_dir) - if os.path.isfile(os.path.join(frame_dir, file)) and (prefix + f"_{file[:-4]}") in camera_to_worlds_panos.keys()] + if os.path.isfile(os.path.join(frame_dir, file)) and (str(prefix) + f"_{file[:-4]}") in camera_to_worlds_panos.keys()] files_list = sorted(files_list) num_ims = len(files_list) @@ -597,8 +597,8 @@ def generate_planar_projections_from_equirectangular_GT( 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}, clip_output=clip_output) - 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) + # pers_image = equi2pers(im, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) 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/scripts/process_data.py b/nerfstudio/scripts/process_data.py index 3ee63f896..31e9aebae 100644 --- a/nerfstudio/scripts/process_data.py +++ b/nerfstudio/scripts/process_data.py @@ -49,6 +49,10 @@ from nerfstudio.process_data.openSFM_to_nerfstudio import ( OpenSFMToNeRFStudioDataset ) +from nerfstudio.process_data.GT_HDR_to_nerfstudio import ( + GTHDRoNeRFStudioDataset +) + from nerfstudio.process_data.video_to_nerfstudio_dataset import VideoToNerfstudioDataset from nerfstudio.utils.rich_utils import CONSOLE @@ -410,6 +414,7 @@ Commands = Union[ Annotated[ProcessAlignedPano, tyro.conf.subcommand(name="aligned_pano2plane")], Annotated[LanternImagesToNerfstudioDataset, tyro.conf.subcommand(name="lantern")], Annotated[OpenSFMToNeRFStudioDataset, tyro.conf.subcommand(name="lantern-openSFM")], + Annotated[GTHDRoNeRFStudioDataset, tyro.conf.subcommand(name="lantern-GT-HDR")], ] diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index 9af596eb7..0bc1a4659 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -171,12 +171,12 @@ def _render_trajectory_video( .cpu() .numpy() ) - elif is_fast_exposure: - u = 5000. - img_uncompress = torch.exp(output_image * torch.log(torch.tensor(u+1.))) - 1. - img_uncompress /= u - output_image = img_uncompress - output_image = (output_image.cpu().numpy()) + # elif is_fast_exposure: + # u = 5000. + # img_uncompress = torch.exp(output_image * torch.log(torch.tensor(u+1.))) - 1. + # img_uncompress /= u + # output_image = img_uncompress + # output_image = (output_image.cpu().numpy()) elif is_validity: mask_f = torch.clip(outputs["validity_f"], 0, 1) # mask_f = torch.ones(mask_f.shape).to(mask_f.get_device()) - mask_f diff --git a/nerfstudio/utils/colormaps.py b/nerfstudio/utils/colormaps.py index 617075b9f..81ddad3ee 100644 --- a/nerfstudio/utils/colormaps.py +++ b/nerfstudio/utils/colormaps.py @@ -66,14 +66,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.colormap == 'inverse-mu': - u = 5000. - img_uncompress = torch.exp(image * torch.log(torch.tensor(u+1.))) - 1. - img_uncompress /= u - image = img_uncompress + # if colormap_options.colormap == 'inverse-mu': + # u = 5000. + # img_uncompress = torch.exp(image * torch.log(torch.tensor(u+1.))) - 1. + # img_uncompress /= u + # image = img_uncompress img_gamma_22 = image img_gamma_22 *= 2.**(colormap_options.exposure_scale) - img_gamma_22 = torch.pow(img_gamma_22, 1./2.2) + # img_gamma_22 = torch.pow(img_gamma_22, 1./2.2) return torch.clamp(img_gamma_22, 0., 1.) # rendering depth outputs