From c549928cd7aaeeae106900d1f650c8ffb76a54e2 Mon Sep 17 00:00:00 2001 From: lefreud Date: Tue, 2 Jan 2024 19:00:28 +0000 Subject: [PATCH] Added code to split synthetic dataset in colmap_baseline and exposed frames --- .gitignore | 4 + lantern_scripts/split_synthetic_colmap.py | 111 ++++++++++++++++++++++ 2 files changed, 115 insertions(+) create mode 100644 lantern_scripts/split_synthetic_colmap.py diff --git a/.gitignore b/.gitignore index 873fe16d9..661bf0136 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,7 @@ +# Ours +Dual_Cameras +gt_test + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/lantern_scripts/split_synthetic_colmap.py b/lantern_scripts/split_synthetic_colmap.py new file mode 100644 index 000000000..23767c975 --- /dev/null +++ b/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() \ No newline at end of file