import subprocess import argparse import os import shutil import sys if __name__ == '__main__': argparser = argparse.ArgumentParser() argparser.add_argument("--input_dir", type=str, required=True, default="/mnt/data/scene/scene_ns/", help="Directory containing input data (e.g. videos, GT, etc.)") argparser.add_argument("--checkpoint", type=str, required=True, default="/mnt/workspace/lantern/nerfstudio-HDR/outputs/coffee_room2_ns/lantern-nerfacto/2025-01-22_194346/", help="Directory containing training data (e.g. config.yml and checkpoint info)") argparser.add_argument("--metric_name", type=str, required=True, default="unaligned", help="Name of the method, used for naming output directories") args = argparser.parse_args() if not args.input_dir.endswith('/'): args.input_dir += '/' metrics_dir = os.path.join(args.input_dir, "metrics") renders_dir = os.path.join(metrics_dir, "images") renders_well = os.path.join(renders_dir, args.metric_name) if os.path.exists(renders_well): shutil.rmtree(renders_well) renders_fast = os.path.join(renders_dir, args.metric_name + "_fast") if os.path.exists(renders_fast): shutil.rmtree(renders_fast) renders_lin = os.path.join(renders_dir, args.metric_name + "_lin") if os.path.exists(renders_lin): shutil.rmtree(renders_lin) renders_lin_HDR = os.path.join(renders_dir, args.metric_name + "_lin_HDR") if os.path.exists(renders_lin_HDR): shutil.rmtree(renders_lin_HDR) lin_renders = os.path.join(renders_dir, args.metric_name + "_lin_renders") if os.path.exists(lin_renders): shutil.rmtree(lin_renders) gt_dir = os.path.join(args.input_dir, "GT") gt_hdr = os.path.join(gt_dir, "GT_hdr") if os.path.exists(gt_hdr): shutil.rmtree(gt_hdr) gt_exr = os.path.join(gt_dir, "GT_exr") gt_exr_renders = os.path.join(gt_dir, "GT_exr_renders") if os.path.exists(gt_exr_renders): shutil.rmtree(gt_exr_renders) # 1. Get ground truth transformations: print("Running ns-process-data to get GT transformations...") subprocess.run([ "ns-process-data", "lantern-GT-HDR", "--data", args.input_dir, "--output-dir", metrics_dir, "--metadata", os.path.join(args.input_dir, "sfm/reconstruction.json"), "--checkpoint", args.checkpoint ], check=True) # 2. Render panos at GT positions for well and fast exposures: print("Rendering well-exposed panos...") subprocess.run([ "ns-render", "camera-path", "--load-config", os.path.join(args.checkpoint, "config.yml"), "--camera-path-filename", os.path.join(metrics_dir, "GT_transforms.json"), "--output-path", renders_well, "--output-format", "images" ], check=True) print("Rendering fast-exposed panos...") subprocess.run([ "ns-render", "camera-path", "--load-config", os.path.join(args.checkpoint, "config.yml"), "--camera-path-filename", os.path.join(metrics_dir, "GT_transforms.json"), "--output-path", renders_fast, "--output-format", "images", "--rendered_output_names", "rgb_fast" ], check=True) # 3. Combine both well and fast exposed panos: print("Combining well and fast exposure renders...") subprocess.run([ "python", "lantern_scripts/combine.py", "--well_dir", renders_well + "/", "--fast_dir", renders_fast + "/", "--out_dir", renders_lin + "/", "--experiment_location", os.path.join(args.input_dir, "data") + "/", "--do_linearization" ], check=True) # 4. Rename the render files: print("Step 4a: Renaming render filenames for consistency between GT and renders...") render_files = sorted([f for f in os.listdir(renders_lin) if f.endswith('.exr')]) print("Render files found:", render_files) total = len(render_files) number_strings = [str(i) for i in range(1, total + 1)] lex_order = sorted(number_strings) for old_path, num in zip(render_files, lex_order): new_filename = f"GT{num}.exr" new_path = os.path.join(renders_lin, new_filename) old_path = os.path.join(renders_lin, old_path) shutil.move(old_path, new_path) print(f"Renamed {os.path.basename(old_path)} to {new_filename}") print("Step 4b: Verifying filename consistency between GT and renders...") try: gt_files = set(f for f in os.listdir(gt_exr) if f.endswith('.exr')) render_files = set(f for f in os.listdir(renders_lin) if f.endswith('.exr')) missing_in_renders = gt_files - render_files extra_in_renders = render_files - gt_files if missing_in_renders or extra_in_renders: if missing_in_renders: print("Error: The following GT files are missing in the renders directory:", missing_in_renders) if extra_in_renders: print("Error: There are extra files in the renders directory not found in GT:", extra_in_renders) sys.exit(1) else: print("All GT filenames are present in the renders directory.") except Exception as e: print("Error during filename verification:", e) sys.exit(1) # 5. Calculate PSNR, SSIM and LPIPS: print("Calculating LDR metrics (PSNR, SSIM, LPIPS)...") subprocess.run([ "python", "lantern_scripts/ldr_res.py", "--gt_dir", gt_exr + "/", "--data_dir", renders_lin + "/" ], check=True) # 6. Generate GT and result renders with Blender: expected_render_dir = gt_exr_renders rendered_files_exist = any( f.endswith('.exr') for f in os.listdir(expected_render_dir) ) if os.path.exists(expected_render_dir) else False if not rendered_files_exist: print("Rendering GT images with Blender...") subprocess.run([ "blender", "--background", "lantern_scripts/scene_new.blend", "-P", "lantern_scripts/hdr_blender.py", "--", os.path.join(gt_exr, ""), os.path.join(gt_exr_renders, "") ], check=True) else: print("Skipping Blender render — output files already exist.") print("Rendering result images with Blender...") subprocess.run([ "blender", "--background", "lantern_scripts/scene_new.blend", "-P", "lantern_scripts/hdr_blender.py", "--", renders_lin + "/", lin_renders + "/" ], check=True) # 7. Calculate si-RMSE, RMSE, RGB angular error and PSNR (LDR rendered): print("Calculating additional metrics (si-RMSE, RMSE, angular error, PSNR)...") subprocess.run([ "python", "lantern_scripts/res_table_1.py", "--results_dataset_roots", lin_renders + "/", "--gt_dataset_root", gt_exr_renders + "/", "--output_folder", metrics_dir + "/results_table/", "--dataset_names", args.metric_name ], check=True) # 8. Convert EXR panos to HDR format: print("Converting EXR panos to HDR...") subprocess.run([ "python", "lantern_scripts/exr2hdr.py", "--hdr_dir", renders_lin, "--output_dir", renders_lin_HDR ], check=True) subprocess.run([ "python", "lantern_scripts/exr2hdr.py", "--hdr_dir", gt_exr, "--output_dir", gt_hdr ], check=True) # 9. Convert EXR panos to LDR format: print("Converting EXR panos to LDR visualization...") subprocess.run([ "python", "lantern_scripts/tonemap.py", "--data_dir", renders_lin, "--out_dir", renders_lin ], check=True) subprocess.run([ "python", "lantern_scripts/tonemap.py", "--data_dir", gt_exr, "--out_dir", gt_exr ], check=True) subprocess.run([ "python", "lantern_scripts/tonemap.py", "--data_dir", lin_renders, "--out_dir", lin_renders ], check=True) subprocess.run([ "python", "lantern_scripts/tonemap.py", "--data_dir", gt_exr_renders, "--out_dir", gt_exr_renders ], check=True) # 9. Convert EXR panos to LDR format: print("Preparing the HTML visualization...") subprocess.run([ "python", "lantern_scripts/results_visualization.py", "--gt_dir", gt_exr, "--gt_render_dir", gt_exr_renders, "--pandora_dir", renders_lin, "--pandora_render_dir", lin_renders, "--out_dir", args.input_dir, "--name", args.metric_name ], check=True)