diff --git a/lantern_scripts/combine_LDR_HDR.py b/lantern_scripts/combine_LDR_HDR.py new file mode 100644 index 000000000..1c29e0cd2 --- /dev/null +++ b/lantern_scripts/combine_LDR_HDR.py @@ -0,0 +1,32 @@ +import os +import numpy as np +from tqdm import tqdm +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/' + +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() + +for ldr_addr in tqdm(imgFiles): + ldr_pano = EnvironmentMap(ldr_addr, 'latlong') + + hdr_addr = ldr_addr.replace('.png', '.exr') + hdr_addr = hdr_addr.replace(LDR_dir, HDR_dir) + assert os.path.exists(hdr_addr), "HDR file does not exists!!!" + hdr_pano = EnvironmentMap(hdr_addr, 'latlong') + hdr_pano.resize(1920) + + 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] + + out_addr = hdr_addr.replace(HDR_dir, output_dir) + imwrite(out_pano, out_addr) + +print("Done!") \ No newline at end of file diff --git a/lantern_scripts/combine_left_right.py b/lantern_scripts/combine_left_right.py new file mode 100644 index 000000000..8e4c865a2 --- /dev/null +++ b/lantern_scripts/combine_left_right.py @@ -0,0 +1,166 @@ +import argparse +import numpy as np +import os +from tqdm import tqdm +from pathlib import Path +from PIL import Image +import copy +import cv2 +import json +from pathlib import Path + +os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" + + +def weighting_function(pixels, zmin, zmax): + threshold = 0.5 * (zmin + zmax) + weights = np.zeros(pixels.shape) + np.finfo(np.float64).eps + + # mask_lowerbound = zmin <= pixels <= threshold + mask_lowerbound = np.greater_equal(pixels, zmin) & np.less_equal(pixels, threshold) + weights[mask_lowerbound] = pixels[mask_lowerbound] - zmin + + # mask_upperbound = threshold <= pixels <= zmax + mask_upperbound = np.greater_equal(pixels, threshold) & np.less_equal(pixels, zmax) + weights[mask_upperbound] = zmax - pixels[mask_upperbound] + + # sums_R = np.sum(weights[:,:,0]) + # sums_G = np.sum(weights[:,:,1]) + # sums_B = np.sum(weights[:,:,2]) + + return weights #/ [sums_R, sums_G, sums_B] + + +def apply_the_exposure_get_weights(hdr_data, exposure): + e1_data = np.clip(hdr_data * exposure, 0, 1) + e1_data = (e1_data * 255.0).astype("uint8") + e1_data = e1_data.astype('float32') + e1_data = e1_data / 255.0 + weights = weighting_function(e1_data, 0.05, 0.95) + e1_data_HDR = e1_data / exposure + + return e1_data, e1_data_HDR, weights + + +def make_weights_binary(pixels, weights, thresh_min, thresh_max, threshold_weights = False): + outputs = np.zeros(pixels.shape) + if threshold_weights: + weights_e1_if_case = copy.deepcopy(weights) + else: + weights_e1_if_case = copy.deepcopy(pixels) + # masked = thresh_min <= weights_e1_if_case <= thresh_max + masked = np.greater_equal(weights_e1_if_case, thresh_min) & np.less_equal(weights_e1_if_case, thresh_max) + outputs[masked] = 1.0 + outputs_unit = (255. * outputs).astype('uint8') + + return outputs_unit + + +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) + + +def write_to_json(filename: Path, content: dict): + """Write data to a JSON file. + + Args: + filename: The filename to write to. + content: The dictionary data to write. + """ + assert filename.suffix == ".json" + # with open(filename, "w", encoding="UTF-8") as file: + # json.dump(content, file) + with open(filename, "w", encoding="utf-8") as f: + json.dump(content, f, indent=4) + + + +def process_one_image(pixels, out_dir, name, has_exposure, exposure_value, zmin, zmax): + if args.has_exposure: + e1_ldr, e1_hdr, weights_e1 = apply_the_exposure_get_weights(pixels, exposure_value) + weights_e1_unit = make_weights_binary(e1_ldr, weights_e1, zmin, zmax) + mask_e1 = weights_e1_unit[:, :, 0] & weights_e1_unit[:, :, 1] & weights_e1_unit[:, :, 2] + e1_mask_addr = os.path.join(out_dir, name + '.png') + # e1_mask_addr = os.path.join(args.out_dir, pano_name + '_e1.png') + Image.fromarray(mask_e1).save(e1_mask_addr) + else: + e1_hdr = pixels + # out_addr = os.path.join(args.out_dir, pano_name + '_e1.exr') + out_addr = os.path.join(out_dir, name + '.exr') + cv2.imwrite(out_addr, e1_hdr) + + +def process_all_images(address, out_dir, prefix, has_exposure, exposure_value, zmin, zmax): + images = [os.path.join(address, f) for f in os.listdir(address) if os.path.isfile(os.path.join(address, f)) and f.endswith('.exr')] + + images = sorted(images) + for pano_addr in tqdm(images): + # original_pano = EnvironmentMap(pano_addr, 'latlong') + # pixels = original_pano.data[:, :, 0:3] + original_pano = np.array(cv2.imread(str(pano_addr), cv2.IMREAD_UNCHANGED)).astype("float32") + pixels = original_pano[:, :, 0:3] + + pano_file = os.path.basename(pano_addr) + pano_name = pano_file[:-4] + name = prefix + pano_name + + process_one_image(pixels, out_dir, name, has_exposure, exposure_value, zmin, zmax) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument('--left_dir', type=str, default='/mnt/workspace/share_w_Junming/Dual_Cameras/left/') + parser.add_argument('--right_dir', type=str, default='/mnt/workspace/share_w_Junming/Dual_Cameras/right/') + parser.add_argument('--out_dir', type=str, default='/mnt/workspace/share_w_Junming/Dual_Cameras/lantern-HDR/') + parser.add_argument('--has_exposure', action='store_true') + parser.add_argument('--exposure1', type=float, default=1.0) + parser.add_argument('--exposure2', type=float, default=0.009) + args = parser.parse_args() + + assert os.path.isfile(args.left_dir + '/transforms.json'), 'The JSON files for left set is missing!' + assert os.path.isfile(args.right_dir + '/transforms.json'), 'The JSON files for right set is missing!' + assert os.path.isdir(args.left_dir + '/HDR_Normal/'), 'The data folder for left set is missing!' + assert os.path.isdir(args.right_dir + '/HDR_Normal/'), 'The data folder for right set is missing!' + + assert not os.path.isdir(args.out_dir), 'The output folder exists!! Take care of it first!' + os.mkdir(args.out_dir) + os.mkdir(args.out_dir + 'HDR_Normal/') + + cam_params_lhs = load_from_json( Path(args.left_dir + '/transforms.json') ) + cam_params_combine = {} + cam_params_combine["camera_angle_x"] = cam_params_lhs["camera_angle_x"] + cam_params_combine["frames"] = [] + for frame in tqdm(cam_params_lhs["frames"]): + file_path = 'lhs_' + frame["file_path"].split('/')[-1] + transform_matrix = frame["transform_matrix"] + cam_params_combine["frames"].append( + { + "file_path": file_path, + "transform_matrix": transform_matrix + } + ) + # TODO adding exposures + cam_params_rhs = load_from_json( Path(args.right_dir + '/transforms.json') ) + for frame in tqdm(cam_params_rhs["frames"]): + file_path = 'rhs_' + frame["file_path"].split('/')[-1] + transform_matrix = frame["transform_matrix"] + cam_params_combine["frames"].append( + { + "file_path": file_path, + "transform_matrix": transform_matrix + } + ) + # TODO adding exposures + write_to_json(Path(args.out_dir + '/transforms_all.json'), cam_params_combine) + + process_all_images(args.left_dir + '/HDR_Normal/', args.out_dir + '/HDR_Normal/', 'lhs_', args.has_exposure, args.exposure1, 0.0, 0.95) + process_all_images(args.right_dir + '/HDR_Normal/', args.out_dir + '/HDR_Normal/', 'rhs_', args.has_exposure, args.exposure2, 0.2, 1.0) + +print("Done!") \ No newline at end of file diff --git a/lantern_scripts/divide_into_exposures.py b/lantern_scripts/divide_into_exposures.py new file mode 100644 index 000000000..152eb92eb --- /dev/null +++ b/lantern_scripts/divide_into_exposures.py @@ -0,0 +1,102 @@ +import argparse +from envmap import EnvironmentMap +import numpy as np +import os +from tqdm import tqdm +from pathlib import Path +from hdrio import imwrite +from PIL import Image +import copy + + +def weighting_function(pixels, zmin, zmax): + threshold = 0.5 * (zmin + zmax) + weights = np.zeros(pixels.shape) + np.finfo(np.float64).eps + + # mask_lowerbound = zmin <= pixels <= threshold + mask_lowerbound = np.greater_equal(pixels, zmin) & np.less_equal(pixels, threshold) + weights[mask_lowerbound] = pixels[mask_lowerbound] - zmin + + # mask_upperbound = threshold <= pixels <= zmax + mask_upperbound = np.greater_equal(pixels, threshold) & np.less_equal(pixels, zmax) + weights[mask_upperbound] = zmax - pixels[mask_upperbound] + + # sums_R = np.sum(weights[:,:,0]) + # sums_G = np.sum(weights[:,:,1]) + # sums_B = np.sum(weights[:,:,2]) + + return weights #/ [sums_R, sums_G, sums_B] + + +def apply_the_exposure_get_weights(hdr_data, exposure): + e1_data = np.clip(hdr_data * exposure, 0, 1) + e1_data = (e1_data * 255.0).astype("uint8") + e1_data = e1_data.astype('float32') + e1_data = e1_data / 255.0 + weights = weighting_function(e1_data, 0.05, 0.95) + e1_data_HDR = e1_data / exposure + + return e1_data, e1_data_HDR, weights + + +def make_weights_binary(pixels, weights, thresh_min, thresh_max, threshold_weights = False): + outputs = np.zeros(pixels.shape) + if threshold_weights: + weights_e1_if_case = copy.deepcopy(weights) + else: + weights_e1_if_case = copy.deepcopy(pixels) + # masked = thresh_min <= weights_e1_if_case <= thresh_max + masked = np.greater_equal(weights_e1_if_case, thresh_min) & np.less_equal(weights_e1_if_case, thresh_max) + outputs[masked] = 1.0 + outputs_unit = (255. * outputs).astype('uint8') + + return outputs_unit + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument('--data_dir', type=str, default='/mnt/data/nerfstudio_ds/synthetic_classroom/cameras_w_baseline/left/') + parser.add_argument('--out_dir', type=str, default='/mnt/data/nerfstudio_ds/synthetic_classroom/cameras_w_baseline/all/') + args = parser.parse_args() + + e1 = 1.0 # None # Left + e2 = None # 0.009 # Right + + test_images = [] + for path in Path(args.data_dir).rglob('*.exr'): + test_images.append(path) + + for pano_addr in tqdm(test_images): + original_pano = EnvironmentMap(pano_addr, 'latlong') + pixels = original_pano.data[:, :, 0:3] + + pano_file = os.path.basename(pano_addr) + pano_name = pano_file[:-4] + + if e1 is not None: + e1_ldr, e1_hdr, weights_e1 = apply_the_exposure_get_weights(pixels, e1) + weights_e1_unit = make_weights_binary(e1_ldr, weights_e1, 0.0, 0.95) + mask_e1 = weights_e1_unit[:, :, 0] & weights_e1_unit[:, :, 1] & weights_e1_unit[:, :, 2] + # out_addr = os.path.join(args.out_dir, pano_name + '_e1.exr') + out_addr = os.path.join(args.out_dir, 'lhs_' + pano_name + '.exr') + imwrite(e1_hdr, out_addr) + e1_mask_addr = os.path.join(args.out_dir, 'lhs_' + pano_name + '.png') + # e1_mask_addr = os.path.join(args.out_dir, pano_name + '_e1.png') + Image.fromarray(mask_e1).save(e1_mask_addr) + + if e2 is not None: + e2_ldr, e2_hdr, weights_e2 = apply_the_exposure_get_weights(pixels, e2) + weights_e2_unit = make_weights_binary(e2_ldr, weights_e2, 0.2, 1.0) + mask_e2 = weights_e2_unit[:, :, 0] & weights_e2_unit[:, :, 1] & weights_e2_unit[:, :, 2] + # out_addr = os.path.join(args.out_dir, pano_name + '_e2.exr') + out_addr = os.path.join(args.out_dir, 'rhs_' + pano_name + '.exr') + imwrite(e2_hdr, out_addr) + # e2_mask_addr = os.path.join(args.out_dir, pano_name + '_e2.png') + e2_mask_addr = os.path.join(args.out_dir, 'rhs_' + pano_name + '.png') + Image.fromarray(mask_e2).save(e2_mask_addr) + + # e1_e2_combined = (weights_e2 * e2_data + weights_e1 * e1_data) / (weights_e1 + weights_e2) + # out_addr = os.path.join(args.out_dir, pano_name + '_combine.exr') + # imwrite(e1_e2_combined, out_addr) + +print("Done!") \ No newline at end of file diff --git a/lantern_scripts/hdr_blender.py b/lantern_scripts/hdr_blender.py new file mode 100644 index 000000000..fbd9e839d --- /dev/null +++ b/lantern_scripts/hdr_blender.py @@ -0,0 +1,79 @@ +""" +Script to render IBL renders +Usage: blender -P hdr_blender.py -b +""" +import os +import random +from math import radians +from os import listdir +from os.path import isfile, join + +import bpy +import cycles +from mathutils import Euler, Matrix, Vector + +# path = 'pano-videos/results/panohdr_log/chess_room/' +path = './' + +datasets = ['chess_room','old_class_spot','old_class_neon','small_class'\ + ,'stairway','cafeteria'] +# # out_dirs = ['prof_kitchen/','prof_living_room/','grandma_attic/'\ +# # ,'grandma_guest_room/','chess_room/'\ +# # ,'old_class_spot/','old_class_neon/','small_class/'\ +# # ,'music_room/'] +# # [f'gt/{data}/' for data in datasets] + [f'ldr/{data}/' for data in datasets] +# # [f'blender_multi/gt/{data}/' for data in datasets] + [f'blender_multi/ldr/{data}/' for data in datasets] \ +inp_dirs = ['/mnt/workspace/lantern/renders/lantern-dual-exposure_no_mask/2023-11-01_022121/'] + # [f'results/ft_momo/{data}/mask_pred/' for data in datasets] + \ + # [f'results/1hdr/{data}/mask_pred/' for data in datasets] +out_dirs = ['/mnt/workspace/lantern/renders/lantern-dual-exposure_no_mask/2023-11-01_022121_w_objects/'] + # [f'blender_multi/ft_momo/{data}/mask_pred/' for data in datasets] + \ + # [f'blender_multi/1hdr/{data}/mask_pred/' for data in datasets] + +# inp_dirs = ['final_0/exr/','final_1/exr/','final_2/exr/'] +# out_dirs = ['final_0_render/','final_1_render/','final_2_render/'] +# out_dirs = ['old_class_spot/render_ldr_500000/render/','chess_room/render_ldr_500000/render/', +# 'cafeteria/render_ldr_500000/render/','stairway/render_ldr_500000/render/', +# 'small_class/render_ldr_500000/render/','cafeteria/render_ldr_500000/render/',\ +# 'prof_living_room/render_ldr_500000/render/'] +pairs = list(zip(inp_dirs,out_dirs)) +scn = bpy.context.scene +scn.render.engine = 'CYCLES' +scn.render.image_settings.file_format = 'OPEN_EXR' +scn.world.use_nodes = True +# Enabf_le GPU rendering. +bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA' +bpy.context.preferences.addons['cycles'].preferences.get_devices() +for device in bpy.context.preferences.addons['cycles'].preferences.devices: + if "TITAN V" == device.name: + device.use = True + else: + device.use = False +# select world node tree +wd = scn.world +nt = bpy.data.worlds[wd.name].node_tree + +for pair in pairs: + inp,out = pair + inp = inp + out = os.path.join(path,out) + os.makedirs(out,exist_ok=True) + S = bpy.context.scene + envFiles = [f for f in listdir(inp) if isfile(join(inp, f)) and f.endswith('.exr')] + print(envFiles) + i = 0 + for env_map in envFiles: + env_file = join(inp, env_map) + bpy.ops.image.open(filepath=env_file, directory=inp, files=[ + {"name": env_map, "name": env_map}], show_multiview=False) + output_name = "render_" + env_map.split('.')[0] +'.exr' + # print(output_name) + # bpy.ops.image.save_as(save_as_render=True, copy=True, filepath=join(out_addr, output_name), show_multiview=False, use_multiview=False) + + + backNode = nt.nodes['Environment Texture'] + backNode.image = bpy.data.images.load(env_file) + + bpy.context.scene.render.filepath = os.path.join(out, output_name) + print("Output",os.path.join(out, output_name)) + bpy.ops.render.render(write_still=True) diff --git a/lantern_scripts/tonemap.py b/lantern_scripts/tonemap.py new file mode 100644 index 000000000..a43d560ed --- /dev/null +++ b/lantern_scripts/tonemap.py @@ -0,0 +1,42 @@ +import argparse +from envmap import EnvironmentMap, rotation_matrix +import numpy as np +import os +from tqdm import tqdm +from PIL import Image +from pathlib import Path +import imageio + +def tonemap(img, gamma=2.4): + """Apply gamma, then clip between 0 and 1, finally convert to uint8 [0,255]""" + return (np.clip(np.power(img,1/gamma), 0.0, 1.0)*255).astype('uint8') + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument('--data_dir', type=str, default='/Volumes/Momo/Adobe_2023/Evaluation/test_sets/weber_test_set_original_hdr_panos/') + parser.add_argument('--out_dir', type=str, default='/Users/momo/Desktop/camera_ready/FID_test_set_256x512/') + parser.add_argument("--resolution", type=int, default=256) + args = parser.parse_args() + + test_images = [] + for path in Path(args.data_dir).rglob('*.exr'): + test_images.append(path) + + imageio.plugins.freeimage.download() + + for pano_addr in tqdm(test_images): + # original_pano = imageio.imread(pano_addr, format='HDR-FI') # EnvironmentMap(pano_addr, 'latlong') + # # gt_ldr = tonemap(reexpose_hdr(original_pano.data)[0]) + # gt_ldr = tonemap(reexpose_hdr(original_pano)[0]) + + original_pano = EnvironmentMap(pano_addr, 'latlong') + # gt_ldr = tonemap(reexpose_hdr(original_pano.data)[0], gamma=2.2) # for computing the FID in table 1, sue tune mapping of 2.2 + # gt_ldr = tonemap(reexpose_hdr(original_pano)[0]) + gt_ldr = tonemap(original_pano.data, gamma=2.2) # No Need to re-expose the StyleLight and EverLight + + pano_file = os.path.basename(pano_addr) + pano_file = pano_file.replace('.exr', '.png') + gt_ldr_addr = os.path.join(args.out_dir, pano_file) + # Image.fromarray(gt_ldr).resize((2 * args.resolution, args.resolution), resample=Image.LANCZOS).save(gt_ldr_addr) + Image.fromarray(gt_ldr).save(gt_ldr_addr) +print("Done!") \ No newline at end of file diff --git a/nerfstudio/data/dataparsers/base_dataparser.py b/nerfstudio/data/dataparsers/base_dataparser.py index 80fab739f..56888e8ef 100644 --- a/nerfstudio/data/dataparsers/base_dataparser.py +++ b/nerfstudio/data/dataparsers/base_dataparser.py @@ -69,6 +69,8 @@ class DataparserOutputs: """Transform applied by the dataparser.""" dataparser_scale: float = 1.0 """Scale applied by the dataparser.""" + exposures: Optional[List[float]] = None + """List of the Exposures for each frames.""" def as_dict(self) -> dict: """Returns the dataclass as a dictionary.""" diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index 8fe6ec4e4..582e0d8c3 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -90,6 +90,7 @@ class Nerfstudio(DataParser): mask_filenames = [] depth_filenames = [] poses = [] + exposures = [] fx_fixed = "fl_x" in meta fy_fixed = "fl_y" in meta @@ -160,6 +161,10 @@ class Nerfstudio(DataParser): depth_fname = self._get_fname(depth_filepath, data_dir, downsample_folder_prefix="depths_") depth_filenames.append(depth_fname) + if "exposure" in frame: + exposure_value = float(frame["exposure"]) + exposures.append(exposure_value) + assert len(mask_filenames) == 0 or ( len(mask_filenames) == len(image_filenames) ), """ @@ -172,6 +177,12 @@ class Nerfstudio(DataParser): Different number of image and depth filenames. You should check that depth_file_path is specified for every frame (or zero frames) in transforms.json. """ + assert len(exposures) == 0 or ( + len(exposures) == len(image_filenames) + ), """ + Different number of image and exposure values. + You should check that exposure_value is specified for every frame (or zero frames) in transforms.json. + """ has_split_files_spec = any(f"{split}_filenames" in meta for split in ("train", "val", "test")) if f"{split}_filenames" in meta: @@ -229,6 +240,7 @@ class Nerfstudio(DataParser): image_filenames = [image_filenames[i] for i in indices] mask_filenames = [mask_filenames[i] for i in indices] if len(mask_filenames) > 0 else [] depth_filenames = [depth_filenames[i] for i in indices] if len(depth_filenames) > 0 else [] + exposures = [exposures[i] for i in indices] if len(exposures) > 0 else [] idx_tensor = torch.tensor(indices, dtype=torch.long) poses = poses[idx_tensor] @@ -294,6 +306,7 @@ class Nerfstudio(DataParser): cameras=cameras, scene_box=scene_box, mask_filenames=mask_filenames if len(mask_filenames) > 0 else None, + exposures=exposures if len(exposures) > 0 else None, dataparser_scale=scale_factor, dataparser_transform=transform_matrix, metadata={ diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 319750901..8407abfa4 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -213,9 +213,14 @@ class PixelSampler: c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) c, y, x = c.cpu(), y.cpu(), x.cpu() collated_batch = { - key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None + key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and key != "exposure" and value is not None } + for key, value in batch.items(): + if key == "exposure": + collated_batch["exposure"] = value[c] + + assert collated_batch["image"].shape[0] == num_rays_per_batch # Needed to correct the random indices to their actual camera idx locations. diff --git a/nerfstudio/lantern/dataset.py b/nerfstudio/lantern/dataset.py index d0c842408..27ab92476 100644 --- a/nerfstudio/lantern/dataset.py +++ b/nerfstudio/lantern/dataset.py @@ -110,6 +110,10 @@ class HDRInputDataset(InputDataset): 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]}" + + if self._dataparser_outputs.exposures is not None: + exposure_value = self._dataparser_outputs.exposures[image_idx] + data["exposure"] = exposure_value metadata = self.get_metadata(data) data.update(metadata) return data diff --git a/nerfstudio/lantern/model.py b/nerfstudio/lantern/model.py index fd724d1c6..2ef9be48a 100644 --- a/nerfstudio/lantern/model.py +++ b/nerfstudio/lantern/model.py @@ -189,4 +189,66 @@ class LanternModel(NerfactoModel): ) images_dict[key] = prop_depth_i - return metrics_dict, images_dict \ No newline at end of file + return metrics_dict, images_dict + + + def weighting_function(self, pixels, zmin, zmax): + threshold = 0.5 * (zmin + zmax) + weights = torch.zeros(pixels.shape, dtype = pixels.dtype).to(self.device) + weights = weights + torch.finfo(torch.float32).eps + + # mask_lowerbound = zmin <= pixels <= threshold + mask_lowerbound = torch.ge(pixels, zmin) & torch.le(pixels, threshold) + weights[mask_lowerbound] = 2.0 * (pixels[mask_lowerbound] - zmin) + + # mask_upperbound = threshold <= pixels <= zmax + mask_upperbound = torch.ge(pixels, threshold) & torch.le(pixels, zmax) + weights[mask_upperbound] = 2.0 * (zmax - pixels[mask_upperbound]) + + return weights + + + def get_loss_dict(self, outputs, batch, metrics_dict=None): + loss_dict = {} + weights_for_loss = None + + image = batch["image"].to(self.device) + + if "exposure" in batch: + u = 5000. + img_uncompress = torch.exp(image * torch.log(torch.tensor(u+1.))) - 1. + img_uncompress /= u + + exposures_resized = batch["exposure"].view(batch["exposure"].shape[0], 1) + + img_uncompress_re_exposed = exposures_resized * img_uncompress + weights_for_loss = self.weighting_function(torch.clip(img_uncompress_re_exposed, 0, 1), 0.0, 1.0) + + pred_rgb, gt_rgb = self.renderer_rgb.blend_background_for_loss_computation( + pred_image=outputs["rgb"], + pred_accumulation=outputs["accumulation"], + gt_image=image, + ) + + if weights_for_loss is not None: + loss_dict["rgb_loss"] = (weights_for_loss * ((gt_rgb - pred_rgb) ** 2)).mean() + else: + loss_dict["rgb_loss"] = self.rgb_loss(gt_rgb, pred_rgb) + + if self.training: + loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( + outputs["weights_list"], outputs["ray_samples_list"] + ) + assert metrics_dict is not None and "distortion" in metrics_dict + loss_dict["distortion_loss"] = self.config.distortion_loss_mult * metrics_dict["distortion"] + if self.config.predict_normals: + # orientation loss for computed normals + loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean( + outputs["rendered_orientation_loss"] + ) + + # ground truth supervision for normals + loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( + outputs["rendered_pred_normal_loss"] + ) + return loss_dict \ No newline at end of file diff --git a/nerfstudio/process_data/aligned_pano2plane.py b/nerfstudio/process_data/aligned_pano2plane.py index e7fc09a22..0795115c0 100644 --- a/nerfstudio/process_data/aligned_pano2plane.py +++ b/nerfstudio/process_data/aligned_pano2plane.py @@ -59,7 +59,9 @@ class ProcessAlignedPano(BaseConverterToNerfstudioDataset): """If True, process the .exr files as HDR images.""" use_mask: bool = False """If True, process the .exr files with mask.""" - + use_exposure: bool = False + """If True, using different exposures.""" + def main(self) -> None: """Process images into a nerfstudio dataset.""" @@ -77,7 +79,8 @@ class ProcessAlignedPano(BaseConverterToNerfstudioDataset): 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_mask = self.use_mask, + use_exposure = self.use_exposure ) self.camera_type = "perspective" metadata_dict = io.load_from_json(self.data / "transforms.json") @@ -86,7 +89,8 @@ class ProcessAlignedPano(BaseConverterToNerfstudioDataset): cropped_masks_filename = [] for frame in metadata_dict["frames"]: cropped_images_filename.append(Path(frame["file_path"])) - cropped_masks_filename.append(Path(frame["mask_path"])) + if self.use_mask: + cropped_masks_filename.append(Path(frame["mask_path"])) # Copy images to output directory if self.is_HDR: copied_image_paths = process_data_utils.copy_images_list_EXR( diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py index a4d5f1871..d3ac3f605 100644 --- a/nerfstudio/process_data/equirect_utils.py +++ b/nerfstudio/process_data/equirect_utils.py @@ -23,6 +23,7 @@ import math import sys from pathlib import Path from typing import List, Tuple +import shutil import cv2 import numpy as np @@ -118,7 +119,7 @@ def generate_planar_projections_from_equirectangular( hdr_dir: Path, 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), + HDR_planar_image_size: Tuple[int, int] = (None, None) ) -> Path: """Generate planar projections from an equirectangular image. @@ -471,7 +472,8 @@ def generate_planar_projections_from_equirectangular_GT( mask_dir: Path = None, crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), clip_output: bool = False, - use_mask = False + use_mask = False, + use_exposure = False, ) -> Path: """Given camera pose, generate planar projections from an equirectangular image. And output corresponding camera pose. @@ -495,7 +497,10 @@ def generate_planar_projections_from_equirectangular_GT( device = torch.device("cuda") metadata_dict = io.load_from_json(metadata_path) frames_previous = metadata_dict["frames"] - camera_to_worlds_panos = np.array([frame["transform_matrix"] for frame in frames_previous]).astype(np.float32) + camera_to_worlds_panos = {} + for frame in frames_previous: + camera_to_worlds_panos[frame["file_path"]] = np.array(frame["transform_matrix"]).astype(np.float32) + fov = 120 yaw_pitch_pairs = [] @@ -535,11 +540,19 @@ def generate_planar_projections_from_equirectangular_GT( equi2pers = Equi2Pers(height=planar_image_size[1], width=planar_image_size[0], fov_x=fov, mode="bilinear") frame_dir = image_dir output_dir = image_dir / "planar_projections" - output_dir.mkdir(exist_ok=True) - output_mask_dir = output_dir + if os.path.isdir(output_dir): + CONSOLE.log(f"The directory {output_dir} exists. Deleting the folder...") + shutil.rmtree(output_dir) + output_dir.mkdir(exist_ok=False) + + output_mask_dir = '' if use_mask: output_mask_dir = image_dir/ "mask" / "planar_projections" - output_mask_dir.mkdir(exist_ok=True) + if os.path.isdir(output_mask_dir): + CONSOLE.log(f"The directory {output_mask_dir} exists. Deleting the folder...") + shutil.rmtree(output_mask_dir) + output_mask_dir.mkdir(exist_ok=False) + progress = Progress( TextColumn("[bold blue]Generating Planar Images", justify="right"), BarColumn(), @@ -552,6 +565,8 @@ def generate_planar_projections_from_equirectangular_GT( idx = 0 files_list = [file for file in os.listdir(frame_dir) if os.path.isfile(os.path.join(frame_dir, file))] + files_list = sorted(files_list) + num_ims = len(files_list) with progress: for i in progress.track(files_list, description="", total=num_ims): @@ -565,12 +580,7 @@ def generate_planar_projections_from_equirectangular_GT( im = torch.tensor(im, dtype=torch.float32, device=device) im = torch.permute(im, (2, 0, 1)) / 255.0 count = 0 - current_pano_camera_pose = np.array(camera_to_worlds_panos[idx]) - # Convert from COLMAP's camera coordinate system (OpenCV) to ours (OpenGL) - current_pano_camera_pose[0:3, 1:3] *= -1 - current_pano_camera_pose = current_pano_camera_pose[np.array([1, 0, 2, 3]), :] - current_pano_camera_pose[2, :] *= -1 - + current_pano_camera_pose = camera_to_worlds_panos[f"{i[:-4]}"] current_pano_camera_rotation = current_pano_camera_pose[:3, :3] for u_deg, v_deg in yaw_pitch_pairs: v_rad = torch.pi * v_deg / 180.0 @@ -594,6 +604,11 @@ def generate_planar_projections_from_equirectangular_GT( "file_path": f"{output_dir}/{i[:-4]}_{count}.exr", "transform_matrix": perspective_camera_pose.tolist(), } + if use_exposure: + exposure = 1.0 + if i[0:3] == "rhs": + exposure = 0.009 + frame["exposure"] = exposure else: pers_image *= 255.0 pers_image = (pers_image.permute(1, 2, 0)).type(torch.uint8).to("cpu").numpy() @@ -606,7 +621,7 @@ def generate_planar_projections_from_equirectangular_GT( if use_mask: mask_file_name = file_name + ".png" mask = np.array(cv2.imread(os.path.join(frame_dir, "mask", mask_file_name))).astype("uint8") - mask = ~((mask == 255).all(axis=-1)) + mask = ~((mask == 0).all(axis=-1)) mask = mask.astype("uint8") * 255 # black corresponde to pixel to be ignored mask = torch.tensor(mask[:,:,None], dtype=torch.uint8, device=device)