源图均衡

2026-08-22 16:17:44 +08:00
1 parent 56f3c93216
commit adedcfd4f1
@@ -0,0 +1,118 @@
+#!/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()