Added code to split synthetic dataset in colmap_baseline and exposed frames

lantern_scripts/split_synthetic_colmap.py

@@ -0,0 +1,111 @@
+"""
+This script works on synthetic data, all HDR frames with a fixed exposure.
+It removes 10 frames from the training set to use it for colmap calibration. 
+This subset is identified with the "colmap" suffix, where both left and right have the same exposure.
+The rest of the frames are used for two-exposure training (e1 and e2).
+
+The masks are generated by the script lantern_scripts/combine_left_right.py, where 0 means unknown, and 1 means known.
+
+Quantization and gamma-tonemapping are used on all images to simulate a camera.
+
+input: left, right in HDR
+
+output (all in LDR):
+left_colmap, left_colmap_mask, 
+right_colmap, right_colmap_mask, 
+left_e1, left_e1_mask,
+right_e2, right_e2_mask
+"""
+import argparse
+import os
+
+import cv2
+import numpy as np
+from tqdm import tqdm
+
+from lantern_scripts.combine_left_right import apply_the_exposure_get_weights, make_weights_binary
+from lantern_scripts.tonemap import tonemap
+from PIL import Image
+
+CALIB_FRAMES = [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]
+# all the rest is used for two-exposures
+
+def process_one_image(pixels, out_dir, name, has_exposure, exposure_value, zmin, zmax):
+    if 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, f'{name}_mask.png')
+        # e1_mask_addr = os.path.join(args.out_dir, pano_name + '_e1.png')
+        Image.fromarray(mask_e1).save(e1_mask_addr)
+    else:
+        raise NotImplementedError
+    # out_addr = os.path.join(args.out_dir, pano_name + '_e1.exr')
+
+    e1_ldr_tonemapped = tonemap(e1_ldr)
+    out_addr = os.path.join(out_dir, f'{name}.png')
+    os.makedirs(os.path.dirname(out_addr), exist_ok=True)
+
+    cv2.imwrite(out_addr, e1_ldr_tonemapped) 
+
+def process_images(images, out_dir, prefix, has_exposure, exposure_value, zmin, zmax):
+    os.makedirs(out_dir)
+    print("Processing images in: ", out_dir)
+    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)
+
+def 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.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!'
+    
+
+
+    paths_left_all = args.left_dir + '/HDR_Normal/'
+    images_left_all = [os.path.join(paths_left_all, f) for f in os.listdir(paths_left_all) if os.path.isfile(os.path.join(paths_left_all, f)) and f.endswith('.exr')]
+    # images_left_all.sort()
+    # sort images named like: r_00.exr, r_01.exr, ..., r_99.exr, r_100.exr, r_101.exr, ..., r_199.exr, r_200.exr, ...
+    # in order
+    images_left_all.sort(key=lambda f: int(''.join(filter(str.isdigit, f))))
+    images_left_calib = [images_left_all[i] for i in CALIB_FRAMES]
+    images_left_2exp = [images_left_all[i] for i in range(len(images_left_all)) if i not in CALIB_FRAMES]
+
+    paths_right_all = args.right_dir + '/HDR_Normal/'
+    images_right_all = [os.path.join(paths_right_all, f) for f in os.listdir(paths_right_all) if os.path.isfile(os.path.join(paths_right_all, f)) and f.endswith('.exr')]
+    # images_right_all.sort()
+    # sort images named like: r_00.exr, r_01.exr, ..., r_99.exr, r_100.exr, r_101.exr, ..., r_199.exr, r_200.exr, ...
+    # in order
+    images_right_all.sort(key=lambda f: int(''.join(filter(str.isdigit, f))))
+    images_right_calib = [images_right_all[i] for i in CALIB_FRAMES]
+    images_right_2exp = [images_right_all[i] for i in range(len(images_right_all)) if i not in CALIB_FRAMES]
+
+    process_images(images_left_calib, args.out_dir + '/left_colmap_baseline/', 'lhs_', args.has_exposure, args.exposure1, 0.0, 0.95)
+    process_images(images_right_calib, args.out_dir + '/right_colmap_baseline/', 'rhs_', args.has_exposure, args.exposure1, 0.0, 0.95)
+
+    process_images(images_left_2exp, args.out_dir + '/left_e1/', 'lhs_', args.has_exposure, args.exposure1, 0.0, 0.95)
+    process_images(images_right_2exp, args.out_dir + '/right_e2/', 'rhs_', args.has_exposure, args.exposure2, 0.2, 1.0)
+
+    print("Done!")
+
+if __name__ == "__main__":
+    main()