Two exposures - with weighted loss

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!")

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!")

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!")

lantern_scripts/hdr_blender.py

@@ -0,0 +1,79 @@
+"""
+Script to render IBL renders
+Usage: blender <blend file> -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)

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!")

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."""

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={

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.

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

nerfstudio/lantern/model.py

@@ -189,4 +189,66 @@ class LanternModel(NerfactoModel):
             )
             images_dict[key] = prop_depth_i
 
-        return metrics_dict, images_dict
+        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

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(

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)