Speckle-Scanner/02_Calibration/main.py

#!/usr/bin/env python3
"""
Calibration entry point.

Default (2-step pipeline):
  1. detect_features.py — corners/ellipses → per-image JSON
  2. calibrate.py       — mono intrinsics + stereo (lc vs rc/rg/ir)

Legacy one-shot mode: --legacy (detect + calibrate in memory, single partner)
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path.home() / "Speckle-Scanner"))
sys.path.insert(0, str(Path(__file__).resolve().parent))

import argparse
import threading
from typing import Optional, Tuple

import config
from calibrationclasses.calibration import StereoCalibration
from calibrationclasses.calibration_engine import (
    run_mono_calibration,
    run_stereo_calibration,
)
from calibrationclasses.cli_common import (
    add_board_args,
    add_session_args,
    add_troubleshooting_arg,
    build_board_config,
    parse_chessboard_size,
    resolve_input_path,
)
from calibrationclasses.feature_detection import DetectionConfig, run_detection
from calibrationclasses.session import STEREO_PARTNERS


def parse_args():
    parser = argparse.ArgumentParser(
        description="Stereo camera calibration (2-step pipeline by default)"
    )
    add_session_args(parser)
    add_board_args(parser)
    parser.add_argument(
        "--step",
        choices=("detect", "calibrate", "all"),
        default="all",
        help="Pipeline step: detect JSONs, calibrate from JSONs, or both (default)",
    )
    parser.add_argument(
        "--legacy",
        action="store_true",
        help="Old one-shot flow: detect in memory, one stereo partner only",
    )
    parser.add_argument(
        "--left_camera",
        type=str,
        default="lc",
        choices=("lc", "lc-ir", "lc_ir"),
        help="Left camera folder for stereo (default: lc)",
    )
    parser.add_argument(
        "--right_camera",
        type=str,
        default="rc",
        choices=("rc", "rgb", "rg", "ir"),
        help="Stereo partner (legacy mode only; 2-step uses lc vs all partners)",
    )
    parser.add_argument(
        "--time_window",
        type=float,
        default=0.1,
        help="Stereo pair time window in seconds (default: 0.1)",
    )
    parser.add_argument(
        "--partners",
        type=str,
        default="rc,rg,ir",
        help="Stereo partners in 2-step mode (default: rc,rg,ir)",
    )
    parser.add_argument(
        "--ir_mode",
        choices=("auto", "chessboard", "ellipse"),
        default="auto",
        help="IR feature detection mode for step 1",
    )
    add_troubleshooting_arg(parser)
    return parser.parse_args()


def run_legacy(
    input_path,
    chessboard_size=(8, 7),
    square_size=0.045,
    chessboard_size_left: Optional[Tuple[int, int]] = None,
    chessboard_size_right: Optional[Tuple[int, int]] = None,
    square_size_left: Optional[float] = None,
    square_size_right: Optional[float] = None,
    preprocessing="None",
    left_camera="lc",
    right_camera="rc",
    troubleshooting=False,
):
    chessboard_size_left = chessboard_size_left or chessboard_size
    chessboard_size_right = chessboard_size_right or chessboard_size
    square_size_left = square_size if square_size_left is None else square_size_left
    square_size_right = (
        square_size if square_size_right is None else square_size_right
    )

    stereo_calibrator = StereoCalibration(
        input_path,
        chessboard_size,
        square_size,
        preprocessing,
        chessboard_size_left=chessboard_size_left,
        chessboard_size_right=chessboard_size_right,
        square_size_left=square_size_left,
        square_size_right=square_size_right,
        left_camera=left_camera,
        right_camera=right_camera,
        troubleshooting=troubleshooting,
    )
    if stereo_calibrator._preprocessing_enabled():
        print(f"[INFO] Preprocessing for corner detection enabled: {preprocessing!r}")

    t1 = threading.Thread(target=stereo_calibrator.create_chessboard_points_left)
    t2 = threading.Thread(target=stereo_calibrator.create_chessboard_points_right)
    t1.start()
    t2.start()
    t1.join()
    t2.join()

    stereo_calibrator.build_pairs_cal()
    stereo_calibrator.calibrate()
    stereo_calibrator.save_stereo_calibration()

    if troubleshooting:
        stereo_calibrator.rectify_calibration_images()


def run_two_step(args):
    input_path = resolve_input_path(args)
    board_sizes, square_sizes = build_board_config(args)
    per_camera_board = {
        name: {"board_size": board_sizes[name], "square_size": square_sizes[name]}
        for name in board_sizes
    }
    left_camera = args.left_camera.lower().replace("_", "-")
    partners = tuple(p.strip() for p in args.partners.split(",") if p.strip())

    if args.step in ("detect", "all"):
        print("\n=== Step 1: Feature detection → JSON ===")
        config_det = DetectionConfig(
            chessboard_size=args.chessboard_size,
            square_size=args.square_size,
            preprocessing=args.preprocessing,
            ir_mode=args.ir_mode,
            troubleshooting=args.troubleshooting,
        )
        run_detection(input_path, config_det, per_camera_board=per_camera_board)

    if args.step in ("calibrate", "all"):
        print("\n=== Step 2a: Mono intrinsics ===")
        mono_results = run_mono_calibration(
            input_path,
            board_sizes,
            square_sizes,
            troubleshooting=args.troubleshooting,
        )
        print("\n=== Step 2b: Stereo (lc vs partners) ===")
        run_stereo_calibration(
            input_path,
            left_camera=left_camera,
            mono_results=mono_results,
            board_sizes=board_sizes,
            square_sizes=square_sizes,
            time_window_sec=args.time_window,
            partners=partners or STEREO_PARTNERS,
            troubleshooting=args.troubleshooting,
        )


if __name__ == "__main__":
    args = parse_args()
    input_path = str(resolve_input_path(args))

    if args.legacy:
        run_legacy(
            input_path=input_path,
            chessboard_size=args.chessboard_size,
            square_size=args.square_size,
            chessboard_size_left=args.left_chessboard_size,
            chessboard_size_right=args.right_chessboard_size,
            square_size_left=args.left_square_size,
            square_size_right=args.right_square_size,
            preprocessing=args.preprocessing,
            left_camera=args.left_camera,
            right_camera=args.right_camera,
            troubleshooting=args.troubleshooting,
        )
    else:
        run_two_step(args)