118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""光度标定入口脚本
|
|
|
|
用法:
|
|
python run_calibration.py \
|
|
--model_dir /path/to/colmap_model \
|
|
--rgb_dir /path/to/rgb \
|
|
--mask_dir /path/to/masks \
|
|
--output_dir /path/to/output
|
|
"""
|
|
|
|
import argparse
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from calib_io import run_calibration
|
|
from calibration import calibrate, CalibrationOptions
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description="Photometric calibration using COLMAP reconstruction",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
|
)
|
|
parser.add_argument(
|
|
"--model_dir", "-m",
|
|
type=str, required=True,
|
|
help="COLMAP sparse reconstruction directory (contains cameras.bin/images.bin/points3D.bin)",
|
|
)
|
|
parser.add_argument(
|
|
"--rgb_dir", "-r",
|
|
type=str, required=True,
|
|
help="Directory containing input RGB images",
|
|
)
|
|
parser.add_argument(
|
|
"--mask_dir", "-mk",
|
|
type=str, required=True,
|
|
help="Directory containing foreground masks (.png, same names as RGB)",
|
|
)
|
|
parser.add_argument(
|
|
"--output_dir", "-o",
|
|
type=str, required=True,
|
|
help="Output directory for corrected images and report",
|
|
)
|
|
|
|
# ---- 可选参数 ----
|
|
parser.add_argument("--mask_erosion_radius", type=int, default=2)
|
|
parser.add_argument("--min_shared_tracks", type=int, default=3)
|
|
parser.add_argument("--valid_color_min", type=int, default=1)
|
|
parser.add_argument("--valid_color_max", type=int, default=254)
|
|
parser.add_argument("--huber_k", type=float, default=1.345)
|
|
parser.add_argument("--max_iterations", type=int, default=20)
|
|
parser.add_argument("--tolerance", type=float, default=1e-6)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
# 路径校验
|
|
model_dir = Path(args.model_dir)
|
|
rgb_dir = Path(args.rgb_dir)
|
|
mask_dir = Path(args.mask_dir)
|
|
output_dir = Path(args.output_dir)
|
|
|
|
if not model_dir.exists():
|
|
print(f"[ERROR] model_dir does not exist: {model_dir}", file=sys.stderr)
|
|
sys.exit(1)
|
|
if not rgb_dir.exists():
|
|
print(f"[ERROR] rgb_dir does not exist: {rgb_dir}", file=sys.stderr)
|
|
sys.exit(1)
|
|
if not mask_dir.exists():
|
|
print(f"[ERROR] mask_dir does not exist: {mask_dir}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# 参数
|
|
options = CalibrationOptions(
|
|
mask_erosion_radius=args.mask_erosion_radius,
|
|
min_shared_tracks=args.min_shared_tracks,
|
|
valid_color_min=args.valid_color_min,
|
|
valid_color_max=args.valid_color_max,
|
|
huber_k=args.huber_k,
|
|
max_iterations=args.max_iterations,
|
|
tolerance=args.tolerance,
|
|
)
|
|
|
|
print("=" * 50)
|
|
print("Photometric Calibration")
|
|
print("=" * 50)
|
|
print(f" model_dir : {model_dir}")
|
|
print(f" rgb_dir : {rgb_dir}")
|
|
print(f" mask_dir : {mask_dir}")
|
|
print(f" output_dir : {output_dir}")
|
|
print(f" min_shared_tracks: {options.min_shared_tracks}")
|
|
print(f" mask_erosion_radius: {options.mask_erosion_radius}")
|
|
print("-" * 50)
|
|
|
|
# 执行标定
|
|
print("Starting calibration...")
|
|
print(f"model_dir={model_dir}")
|
|
outputs = run_calibration(
|
|
model_dir=str(model_dir),
|
|
rgb_dir=str(rgb_dir),
|
|
mask_dir=str(mask_dir),
|
|
output_dir=str(output_dir),
|
|
calibrator=lambda dataset: calibrate(dataset, options=options),
|
|
)
|
|
|
|
print("=" * 50)
|
|
print(f"[DONE] Corrected images: {len(outputs.image_paths)}")
|
|
print(f" Output dir : {output_dir}")
|
|
print(f" Report : {outputs.report_path}")
|
|
print("=" * 50)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |