You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
118 lines
3.6 KiB
118 lines
3.6 KiB
#!/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() |