mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 19:04:03 +08:00
111 lines
5.2 KiB
Python
111 lines
5.2 KiB
Python
"""
|
|
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() |