mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 19:04:03 +08:00
metrics for step one
This commit is contained in:
MrKarimiD committed
1 parent
34d5660acf
commit
e3db99e017
1 file changed
+158
@@ -0,0 +1,158 @@
|
||||
import subprocess
|
||||
import argparse
|
||||
import os
|
||||
import shutil
|
||||
|
||||
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)
|
||||
renders_fast = os.path.join(renders_dir, args.metric_name + "_fast")
|
||||
renders_lin = os.path.join(renders_dir, args.metric_name + "_lin")
|
||||
renders_lin_HDR = os.path.join(renders_dir, args.metric_name + "_lin_HDR")
|
||||
lin_renders = os.path.join(renders_dir, args.metric_name + "_lin_renders")
|
||||
|
||||
gt_dir = os.path.join(args.input_dir, "GT")
|
||||
gt_hdr = os.path.join(gt_dir, "GT_hdr")
|
||||
gt_exr = os.path.join(gt_dir, "GT_exr")
|
||||
gt_exr_renders = os.path.join(gt_dir, "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("Linearizing the frames...")
|
||||
subprocess.run([
|
||||
"python", "lantern_scripts/linearize_jpgs.py",
|
||||
"--data_dir", renders_well + "/",
|
||||
"--out_dir", renders_lin + "/",
|
||||
], 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)
|
||||
# import pdb; pdb.set_trace()
|
||||
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}")
|
||||
# import pdb; pdb.set_trace()
|
||||
|
||||
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)
|
||||
Reference in new issue
Block a user