From 2719a4457bc098454702dca56de8d5bd7cd1c25d Mon Sep 17 00:00:00 2001 From: lefreud Date: Tue, 7 May 2024 23:02:36 +0000 Subject: [PATCH] Generate small subset for HDR Nerf --- README-pandora.md | 3 + .../generate_small_subset_hdr_nerf.py | 135 ++++++++++++++++++ 2 files changed, 138 insertions(+) create mode 100644 lantern_scripts/generate_small_subset_hdr_nerf.py diff --git a/README-pandora.md b/README-pandora.md index c7a16dd89..06a749fed 100644 --- a/README-pandora.md +++ b/README-pandora.md @@ -4,6 +4,9 @@ ## 0. Setup + +Open this folder in VS Code Dev Container extension! + ``` export PYTHONPATH=$pwd:$PYTHONPATH pip install skylibs equilib diff --git a/lantern_scripts/generate_small_subset_hdr_nerf.py b/lantern_scripts/generate_small_subset_hdr_nerf.py new file mode 100644 index 000000000..36f34b8bc --- /dev/null +++ b/lantern_scripts/generate_small_subset_hdr_nerf.py @@ -0,0 +1,135 @@ +import argparse +import json +import math +import os +import re +import shutil +from pathlib import Path + +import cv2 +import numpy as np + + +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 is_mask_all_valid(mask_path): + mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) + return np.all(mask > 200) + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument('--input_json', type=str, required=True) + parser.add_argument('--output_dir', type=str, required=True) + parser.add_argument('--train_split_fraction', type=float, default= 0.8) + parser.add_argument('--eval_split_fraction', type=float, default= 0.1) + parser.add_argument('--min_frame_number', type=int, default=100) + parser.add_argument('--frames_per_direction', type=int, default=100) + args = parser.parse_args() + + assert os.path.isfile(args.input_json) and args.input_json.endswith('.json'), 'The input JSON file is missing!' + + os.makedirs(args.output_dir, exist_ok=True) + + transforms_data = load_from_json( Path(args.input_json) ) + exposures_data = load_from_json( Path(args.input_json).parent / 'exposures.json' ) + + + selected_left_frames = [] + selected_right_frames = [] + new_exposures = dict() + + os.makedirs(os.path.join(args.output_dir, 'images'), exist_ok=True) + os.makedirs(os.path.join(args.output_dir, 'masks'), exist_ok=True) + for i, frame in enumerate(transforms_data['frames']): + file_path = frame['file_path'] # e.g. images/left_e1_output_0010_0_linear.exr (frame num is 0010) + mask_path = str(Path(args.input_json).parent / frame['mask_path']) # e.g. masks/left_e1_output_0010_0.png + # get frame number + frame_num = int(re.search(r'(\d{4})', file_path).group(0)) + direction = re.search(r'(left|right)', file_path, re.IGNORECASE).group(0).lower() + + if int(frame_num) >= args.min_frame_number: + if direction == 'left': + if len(selected_left_frames) < args.frames_per_direction and is_mask_all_valid(mask_path): + selected_left_frames.append(frame) + + shutil.copy(str(Path(args.input_json).parent / file_path), os.path.join(args.output_dir, file_path)) + shutil.copy(str(Path(args.input_json).parent / file_path.replace('_linear.exr', '.png')), os.path.join(args.output_dir, file_path.replace('_linear.exr', '.png'))) + shutil.copy(mask_path, os.path.join(args.output_dir, frame['mask_path'])) + shutil.copy(mask_path.replace('.png', '_saturation_mask.png'), os.path.join(args.output_dir, frame['mask_path'].replace('.png', '_saturation_mask.png'))) + + new_exposures[file_path] = exposures_data[file_path] + else: + if len(selected_right_frames) < args.frames_per_direction and is_mask_all_valid(mask_path): + selected_right_frames.append(frame) + + shutil.copy(str(Path(args.input_json).parent / file_path), os.path.join(args.output_dir, file_path)) + shutil.copy(str(Path(args.input_json).parent / file_path.replace('_linear.exr', '.png')), os.path.join(args.output_dir, file_path.replace('_linear.exr', '.png'))) + shutil.copy(mask_path, os.path.join(args.output_dir, frame['mask_path'])) + shutil.copy(mask_path.replace('.png', '_saturation_mask.png'), os.path.join(args.output_dir, frame['mask_path'].replace('.png', '_saturation_mask.png'))) + + new_exposures[file_path] = exposures_data[file_path] + + + new_transforms_data = {} + all_keys = transforms_data.keys() + for key in all_keys: + new_transforms_data[key] = transforms_data[key] + + new_transforms_data['frames'] = selected_left_frames + selected_right_frames + + num_images_per_direction = len(selected_left_frames) + + num_train_images = math.ceil(num_images_per_direction * args.train_split_fraction) + num_eval_test_images = num_images_per_direction - num_train_images + i_all = np.arange(num_images_per_direction) + i_train = np.linspace( + 0, num_images_per_direction - 1, num_train_images, dtype=int + ) # equally spaced training images starting and ending at 0 and num_images-1 + i_eval_test = np.setdiff1d(i_all, i_train) # eval images are the remaining images + assert len(i_eval_test) == num_eval_test_images + + num_eval_images = math.ceil(num_images_per_direction * args.eval_split_fraction) + num_test_images = num_eval_test_images - num_train_images + i_eval = np.linspace( + 0, num_eval_test_images - 1, num_eval_images, dtype=int + ) # equally spaced training images starting and ending at 0 and num_images-1 + i_eval = i_eval_test[i_eval] + i_test = np.setdiff1d(i_eval_test, i_eval) # eval images are the remaining images + assert len(i_eval) == num_eval_images + + + # has_split_files_spec = any(f"{split}_filenames" in meta for split in ("train", "val", "test")) + + train_list = [] + for i in i_train: + train_list.append(selected_left_frames[i]["file_path"]) + train_list.append(selected_right_frames[i]["file_path"]) + new_transforms_data["train_filenames"] = train_list + + test_list = [] + for i in i_test: + test_list.append(selected_left_frames[i]["file_path"]) + test_list.append(selected_right_frames[i]["file_path"]) + new_transforms_data["test_filenames"] = test_list + + eval_list = [] + for i in i_eval: + eval_list.append(selected_left_frames[i]["file_path"]) + eval_list.append(selected_right_frames[i]["file_path"]) + new_transforms_data["val_filenames"] = eval_list + + with open(os.path.join(args.output_dir, 'transforms.json'), "w") as outfile: + json.dump(new_transforms_data, outfile, indent=4) + + with open(os.path.join(args.output_dir, 'exposures.json'), "w") as outfile: + json.dump(new_exposures, outfile, indent=4) + + print("All done!")