calibrate.py

from __future__ import annotations

import argparse
from dataclasses import dataclass
from enum import Enum
import json
from pathlib import Path
import shutil
import subprocess
import sys
import tempfile
import time
from typing import Literal

import cv2 as cv
import numpy as np

from cam_util import OpenCvFrameReader, WpilibFrameReader


class InputKind(Enum):
    PORT = "port"
    IMAGES = "images"


@dataclass(frozen=True)
class CalibrationInput:
    kind: InputKind
    value: int | str
    
    @classmethod
    def PORT(cls, port: int) -> CalibrationInput:
        """Build a live-camera input from a device index."""
        return CalibrationInput(InputKind.PORT, port)

    @classmethod
    def IMAGES(cls, path: str) -> CalibrationInput:
        """Build a folder input from an image directory."""
        return CalibrationInput(InputKind.IMAGES, path)


@dataclass(frozen=True)
class CapturedSample:
    """Stores one captured frame and its ChArUco detections."""
    frame: np.ndarray
    charuco_corners: np.ndarray
    charuco_ids: np.ndarray


def draw_marker_dots(
    image, marker_corners_list, color=(255, 0, 255), min_radius=3, max_radius=15
):
    """Draw dots for captured ArUco markers."""
    if marker_corners_list is None or len(marker_corners_list) == 0:
        return

    marker_sizes = []
    marker_centers = []

    for corners in marker_corners_list:
        corners_2d = (
            corners.copy().reshape(-1, 2) if len(corners.shape) > 2 else corners.copy()
        )

        center = np.mean(corners_2d, axis=0).astype(int)
        marker_centers.append(center)
        edge_lengths = []
        for i in range(4):
            p1 = corners_2d[i]
            p2 = corners_2d[(i + 1) % 4]
            edge_lengths.append(np.linalg.norm(p2 - p1))
        marker_sizes.append(np.mean(edge_lengths))

    min_size = min(marker_sizes)
    max_size = max(marker_sizes)
    size_range = max_size - min_size if max_size > min_size else 1.0

    for center, size in zip(marker_centers, marker_sizes):
        if size_range > 0:
            normalized_size = (size - min_size) / size_range
            radius = int(min_radius + normalized_size * (max_radius - min_radius))
        else:
            radius = (min_radius + max_radius) // 2

        cv.circle(image, tuple(center), radius, color, -1)


class CharucoCapture:
    """Runs a capture-and-detect loop and accumulates detections for calibration."""

    def __init__(
        self,
        detection_skip_frames: int,
        detector: "cv.aruco.CharucoDetector",
        manual_capture: bool = False,
        aggressive_refinement: bool = False,
    ):
        """Store detector settings for the capture loop."""
        self.detection_skip_frames = max(0, detection_skip_frames)
        self.detector = detector
        self.manual_capture = manual_capture
        self.aggressive_refinement = aggressive_refinement
        self.sampled_detection_interval = self.detection_skip_frames + 1

    def _scale_detected_points(self, points, scale: float):
        """Scale detected corner coordinates back to the original image size."""
        if points is None or scale == 1.0:
            return points
        return np.asarray(points, dtype=np.float32).copy() / scale

    def _scale_marker_corners(self, marker_corners, scale: float):
        """Scale detected marker corner lists back to the original image size."""
        if marker_corners is None or scale == 1.0:
            return marker_corners
        return [np.asarray(corners, dtype=np.float32).copy() / scale for corners in marker_corners]

    def _detect_on_image(self, image, scale: float = 1.0):
        """Run board detection and map results back to the source image space."""
        charuco_corners, charuco_ids, marker_corners, marker_ids = self.detector.detectBoard(image)
        return (
            self._scale_detected_points(charuco_corners, scale),
            charuco_ids,
            self._scale_marker_corners(marker_corners, scale),
            marker_ids,
        )

    def _build_refined_candidates(self, frame_gray):
        """Generate refined grayscale variants for offline image detection."""
        candidates = [(frame_gray, 1.0)]

        clahe = cv.createCLAHE(clipLimit=2.5, tileGridSize=(8, 8)).apply(frame_gray)
        candidates.append((clahe, 1.0))

        blurred = cv.GaussianBlur(frame_gray, (5, 5), 0)
        candidates.append((blurred, 1.0))

        median = cv.medianBlur(frame_gray, 5)
        candidates.append((median, 1.0))

        sharpen_kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], dtype=np.float32)
        sharpened = cv.filter2D(clahe, -1, sharpen_kernel)
        candidates.append((sharpened, 1.0))

        otsu = cv.threshold(clahe, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)[1]
        candidates.append((otsu, 1.0))

        adaptive = cv.adaptiveThreshold(
            clahe,
            255,
            cv.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv.THRESH_BINARY,
            31,
            5,
        )
        candidates.append((adaptive, 1.0))

        max_dimension = max(frame_gray.shape)
        upscale_factor = 1.0
        if max_dimension < 1200:
            upscale_factor = 2.0
        elif max_dimension < 2200:
            upscale_factor = 1.5

        if upscale_factor > 1.0:
            upscaled = cv.resize(
                frame_gray,
                None,
                fx=upscale_factor,
                fy=upscale_factor,
                interpolation=cv.INTER_CUBIC,
            )
            candidates.append((upscaled, upscale_factor))

            upscaled_clahe = cv.createCLAHE(clipLimit=2.5, tileGridSize=(8, 8)).apply(upscaled)
            candidates.append((upscaled_clahe, upscale_factor))

            upscaled_sharpened = cv.filter2D(upscaled_clahe, -1, sharpen_kernel)
            candidates.append((upscaled_sharpened, upscale_factor))

        return candidates

    def _pick_best_detection(self, frame_gray):
        """Try multiple variants and return the strongest detection."""
        candidates = (
            self._build_refined_candidates(frame_gray)
            if self.aggressive_refinement
            else [(frame_gray, 1.0)]
        )

        best_result = (None, None, None, None)
        best_score = (-1, -1)

        for candidate_image, scale in candidates:
            detection = self._detect_on_image(candidate_image, scale=scale)
            charuco_corners, charuco_ids, _marker_corners, marker_ids = detection
            score = (
                0 if charuco_ids is None else len(charuco_ids),
                0 if marker_ids is None else len(marker_ids),
            )
            if score > best_score:
                best_result = detection
                best_score = score

        return best_result

    def _record_capture(
        self,
        frame,
        charuco_corners,
        charuco_ids,
        marker_corners,
        all_corners,
        all_ids,
        all_counter,
        all_marker_corners,
        captured_samples,
    ) -> None:
        """Store one valid detection for calibration and preview."""
        all_corners.extend(charuco_corners)
        all_ids.extend(charuco_ids)
        all_counter.append(len(charuco_ids))

        if marker_corners is not None and len(marker_corners) > 0:
            for marker_corner in marker_corners:
                all_marker_corners.append(marker_corner.copy())

        captured_samples.append(
            CapturedSample(
                frame=frame.copy(),
                charuco_corners=np.asarray(charuco_corners, dtype=np.float32).copy(),
                charuco_ids=np.asarray(charuco_ids, dtype=np.int32).copy(),
            )
        )

    def run(self, frame_source) -> tuple[list, list, list, tuple[int, int], list[CapturedSample]]:
        """Collect detections from the provided frame source."""
        frame_count = 0
        frame_shape = (0, 0)
        all_corners = []
        all_ids = []
        all_counter = []
        all_marker_corners = []
        captured_samples = []
        prev_time = None
        fps = 0.0
        status_text = ""
        status_color = (0, 255, 0)
        status_expires_at = 0.0

        try:
            for frame in frame_source.frames():
                frame_count += 1

                now = time.time()
                if prev_time is not None:
                    instant_fps = 1.0 / max(now - prev_time, 1e-6)
                    fps = fps * 0.9 + instant_fps * 0.1
                else:
                    fps = 0.0
                prev_time = now

                frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
                if frame_shape != (0, 0) and frame_shape != frame_gray.shape:
                    raise RuntimeError(
                        "All calibration images must have the same resolution."
                    )
                frame_shape = frame_gray.shape

                should_detect_this_frame = self.manual_capture or (
                    frame_count % self.sampled_detection_interval == 0
                )
                debug_image = frame.copy()
                charuco_corners = None
                charuco_ids = None
                marker_corners = None
                marker_ids = None

                if should_detect_this_frame:
                    charuco_corners, charuco_ids, marker_corners, marker_ids = self._pick_best_detection(frame_gray)

                    if marker_corners is not None and len(marker_corners) > 0:
                        debug_image = cv.aruco.drawDetectedMarkers(
                            debug_image, marker_corners, marker_ids
                        )

                    if charuco_ids is not None and len(charuco_ids) > 0:
                        debug_image = cv.aruco.drawDetectedCornersCharuco(
                            debug_image, charuco_corners, charuco_ids
                        )

                    if (
                        not self.manual_capture
                        and charuco_ids is not None
                        and len(charuco_ids) >= 4
                    ):
                        self._record_capture(
                            frame,
                            charuco_corners,
                            charuco_ids,
                            marker_corners,
                            all_corners,
                            all_ids,
                            all_counter,
                            all_marker_corners,
                            captured_samples,
                        )

                draw_marker_dots(debug_image, all_marker_corners)

                cv.putText(
                    debug_image,
                    f"Images captured: {len(all_counter)}",
                    (10, 30),
                    cv.FONT_HERSHEY_SIMPLEX,
                    0.5,
                    (0, 255, 0),
                    2,
                )

                cv.putText(
                    debug_image,
                    f"Resolution: {frame_shape[1]}x{frame_shape[0]}",
                    (10, 60),
                    cv.FONT_HERSHEY_SIMPLEX,
                    0.5,
                    (0, 255, 0),
                    2,
                )

                cv.putText(
                    debug_image,
                    f"FPS: {fps:.1f}",
                    (10, 90),
                    cv.FONT_HERSHEY_SIMPLEX,
                    0.5,
                    (0, 255, 0),
                    2,
                )

                capture_mode_text = (
                    "Mode: manual (SPACE capture, q calibrate)"
                    if self.manual_capture
                    else "Mode: auto (q calibrate)"
                )
                cv.putText(
                    debug_image,
                    capture_mode_text,
                    (10, 120),
                    cv.FONT_HERSHEY_SIMPLEX,
                    0.5,
                    (0, 255, 0),
                    2,
                )

                if status_text and time.time() < status_expires_at:
                    cv.putText(
                        debug_image,
                        status_text,
                        (10, 150),
                        cv.FONT_HERSHEY_SIMPLEX,
                        0.5,
                        status_color,
                        2,
                    )

                cv.imshow("Frame", debug_image)
                key = cv.waitKey(1) & 0xFF
                if key == ord("q"):
                    break
                if key == ord(" ") and self.manual_capture:
                    if charuco_ids is not None and len(charuco_ids) >= 4:
                        self._record_capture(
                            frame,
                            charuco_corners,
                            charuco_ids,
                            marker_corners,
                            all_corners,
                            all_ids,
                            all_counter,
                            all_marker_corners,
                            captured_samples,
                        )

                        status_text = f"Captured image {len(all_counter)}"
                        status_color = (0, 255, 0)
                    else:
                        status_text = "Capture skipped: need at least 4 ChArUco corners"
                        status_color = (0, 165, 255)
                    status_expires_at = time.time() + 1.5
        finally:
            frame_source.close()
            cv.destroyAllWindows()

        return all_corners, all_ids, all_counter, frame_shape, captured_samples


class FolderFrameReader:
    """Reads pre-captured calibration images from a folder."""

    SUPPORTED_SUFFIXES = {
        ".bmp",
        ".heic",
        ".heif",
        ".jpeg",
        ".jpg",
        ".png",
        ".tif",
        ".tiff",
        ".webp",
    }

    def __init__(self, image_folder_path: str):
        """Store the folder path used for image playback."""
        self.image_folder_path = Path(image_folder_path)

    def _read_with_heif_convert(self, image_path: Path):
        """Decode a HEIC image through heif-convert when available."""
        heif_convert = shutil.which("heif-convert")
        if heif_convert is None:
            return None

        temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        temp_path = Path(temp_file.name)
        temp_file.close()

        try:
            result = subprocess.run(
                [heif_convert, str(image_path), str(temp_path)],
                capture_output=True,
                text=True,
                check=False,
            )
            if result.returncode != 0:
                return None
            return cv.imread(str(temp_path))
        finally:
            temp_path.unlink(missing_ok=True)

    def _read_with_ffmpeg(self, image_path: Path):
        """Decode a HEIC image through ffmpeg as a fallback."""
        ffmpeg = shutil.which("ffmpeg")
        if ffmpeg is None:
            return None

        temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        temp_path = Path(temp_file.name)
        temp_file.close()

        try:
            result = subprocess.run(
                [ffmpeg, "-y", "-i", str(image_path), "-frames:v", "1", str(temp_path)],
                capture_output=True,
                text=True,
                check=False,
            )
            if result.returncode != 0:
                return None
            return cv.imread(str(temp_path))
        finally:
            temp_path.unlink(missing_ok=True)

    def _read_image(self, image_path: Path):
        """Read one image, including HEIC fallbacks on macOS."""
        frame = cv.imread(str(image_path))
        if frame is not None:
            return frame

        if image_path.suffix.lower() in {".heic", ".heif"}:
            frame = self._read_with_heif_convert(image_path)
            if frame is not None:
                return frame
            return self._read_with_ffmpeg(image_path)

        return None

    def frames(self):
        """Yield supported images from the folder in sorted order."""
        if not self.image_folder_path.exists():
            raise RuntimeError(f"Image folder does not exist: {self.image_folder_path}")
        if not self.image_folder_path.is_dir():
            raise RuntimeError(f"Image folder is not a directory: {self.image_folder_path}")

        image_paths = sorted(
            path
            for path in self.image_folder_path.iterdir()
            if path.is_file() and path.suffix.lower() in self.SUPPORTED_SUFFIXES
        )

        if not image_paths:
            raise RuntimeError(
                f"No supported images found in folder: {self.image_folder_path}"
            )

        yielded_any = False
        for image_path in image_paths:
            frame = self._read_image(image_path)
            if frame is None:
                continue
            yielded_any = True
            yield frame

        if not yielded_any:
            raise RuntimeError(
                "No images could be decoded from the folder. For HEIC inputs, install a working HEIC decoder such as heif-convert."
            )

    def close(self) -> None:
        """Keep a no-op close API for reader parity."""
        pass


class Calibration:
    """Configures and runs a ChArUco camera calibration."""

    def __init__(self):
        """Initialize calibration settings with defaults."""
        self.input_source = None
        self.output_file_path = None
        self.square_size = None
        self.marker_size = None
        self.width = None
        self.height = None
        self.limit_fps = 0
        self.fps = 50
        self.use_wpilib = False
        self.set_fps = False
        self.manual_capture = False
        self.max_detection_accuracy = False
        self.preview_rectified_captures = False
        self.resolution_width = 640
        self.resolution_height = 480

    def with_input(self, input_source: Literal[CalibrationInput.PORT, CalibrationInput.IMAGES]) -> "Calibration":
        """Set the calibration input source."""
        self.input_source = input_source
        return self

    def with_output_file_path(self, output_file_path: str) -> "Calibration":
        """Set the JSON output path."""
        self.output_file_path = output_file_path
        return self

    def with_square_size(self, square_size: float) -> "Calibration":
        """Set the physical ChArUco square size in meters."""
        self.square_size = square_size
        return self

    def with_marker_size(self, marker_size: float) -> "Calibration":
        """Set the physical marker size in meters."""
        self.marker_size = marker_size
        return self

    def with_width(self, width: int) -> "Calibration":
        """Set the board width in squares."""
        self.width = width
        return self

    def with_height(self, height: int) -> "Calibration":
        """Set the board height in squares."""
        self.height = height
        return self

    def with_limit_fps(self, limit_fps: int = 0) -> "Calibration":
        """Set how many frames to skip between detections."""
        self.limit_fps = max(0, int(limit_fps))
        return self

    def with_fps(self, fps: int) -> "Calibration":
        """Set the requested device FPS."""
        self.fps = fps
        return self

    def with_use_wpilib(self, use_wpilib: bool = True) -> "Calibration":
        """Choose the WPILib frame reader."""
        self.use_wpilib = use_wpilib
        return self

    def with_set_fps(self, set_fps: bool = True) -> "Calibration":
        """Control whether camera FPS is requested explicitly."""
        self.set_fps = set_fps
        return self

    def with_manual_capture(self, manual_capture: bool = True) -> "Calibration":
        """Enable capture-on-space manual mode."""
        self.manual_capture = manual_capture
        return self

    def with_max_detection_accuracy(self, max_detection_accuracy: bool = True) -> "Calibration":
        """Process every frame with the heaviest detection path."""
        self.max_detection_accuracy = max_detection_accuracy
        return self

    def with_preview_rectified_captures(self, preview_rectified_captures: bool = True) -> "Calibration":
        """Show a post-calibration slideshow of rectified captures."""
        self.preview_rectified_captures = preview_rectified_captures
        return self

    def with_resolution(self, width: int, height: int) -> "Calibration":
        """Set the requested capture resolution."""
        self.resolution_width = width
        self.resolution_height = height
        return self

    def with_resolution_width(self, resolution_width: int) -> "Calibration":
        """Set only the requested capture width."""
        self.resolution_width = resolution_width
        return self

    def with_resolution_height(self, resolution_height: int) -> "Calibration":
        """Set only the requested capture height."""
        self.resolution_height = resolution_height
        return self

    @classmethod
    def with_defaults(cls) -> "Calibration":
        """Create a calibration object with default settings."""
        return cls()

    def _validate(self) -> None:
        """Validate that the current settings are usable."""
        missing_fields = []
        for field_name in (
            "output_file_path",
            "square_size",
            "marker_size",
            "width",
            "height",
        ):
            if getattr(self, field_name) is None:
                missing_fields.append(field_name)

        if missing_fields:
            formatted = ", ".join(missing_fields)
            raise ValueError(f"Missing calibration settings: {formatted}")

        if self.input_source is None:
            raise ValueError("Missing calibration source: use with_input(PORT(...)) or with_input(IMAGES(...)).")

        if self.manual_capture and self.input_source.kind == InputKind.IMAGES:
            raise ValueError("Manual capture mode is only supported with a live camera input.")

    def _build_detector(self) -> tuple["cv.aruco.CharucoBoard", "cv.aruco.CharucoDetector"]:
        """Build the ChArUco board and detector."""
        aruco_dict = cv.aruco.getPredefinedDictionary(cv.aruco.DICT_5X5_1000)
        charuco_board = cv.aruco.CharucoBoard(
            (self.width, self.height),
            self.square_size,
            self.marker_size,
            aruco_dict,
        )
        return charuco_board, cv.aruco.CharucoDetector(charuco_board)

    def _build_frame_reader(self):
        """Create the configured frame reader."""
        if self.input_source.kind == InputKind.IMAGES:
            return FolderFrameReader(str(self.input_source.value))

        reader_cls = WpilibFrameReader if self.use_wpilib else OpenCvFrameReader
        return reader_cls(
            port=int(self.input_source.value),
            resolution_width=self.resolution_width,
            resolution_height=self.resolution_height,
            set_fps=self.set_fps,
            requested_fps=self.fps,
        )

    def _solve_sample_pose(
        self,
        charuco_board,
        sample: CapturedSample,
        camera_matrix,
        dist_coeffs,
    ):
        """Solve the board pose for one captured sample."""
        if len(sample.charuco_ids) < 4:
            return None

        object_points = charuco_board.getChessboardCorners()[
            sample.charuco_ids.flatten()
        ].astype(np.float32)
        image_points = sample.charuco_corners.reshape(-1, 2).astype(np.float32)

        success, rvec, tvec = cv.solvePnP(
            object_points,
            image_points,
            camera_matrix,
            dist_coeffs,
            flags=cv.SOLVEPNP_ITERATIVE,
        )
        if not success:
            return None
        return rvec, tvec

    def _rectify_sample(
        self,
        charuco_board,
        sample: CapturedSample,
        camera_matrix,
        dist_coeffs,
    ):
        """Create a fronto-parallel rectified board view for one sample."""
        pose = self._solve_sample_pose(charuco_board, sample, camera_matrix, dist_coeffs)
        undistorted = cv.undistort(sample.frame, camera_matrix, dist_coeffs)
        if pose is None:
            return undistorted

        rvec, tvec = pose
        board_width = (self.width - 1) * self.square_size
        board_height = (self.height - 1) * self.square_size
        board_corners = np.array(
            [
                [0.0, 0.0, 0.0],
                [board_width, 0.0, 0.0],
                [board_width, board_height, 0.0],
                [0.0, board_height, 0.0],
            ],
            dtype=np.float32,
        )
        projected_corners, _ = cv.projectPoints(
            board_corners,
            rvec,
            tvec,
            camera_matrix,
            np.zeros_like(dist_coeffs),
        )
        source = projected_corners.reshape(-1, 2).astype(np.float32)

        pixels_per_square = 180
        output_width = max(1, int((self.width - 1) * pixels_per_square))
        output_height = max(1, int((self.height - 1) * pixels_per_square))
        destination = np.array(
            [
                [0.0, 0.0],
                [output_width - 1.0, 0.0],
                [output_width - 1.0, output_height - 1.0],
                [0.0, output_height - 1.0],
            ],
            dtype=np.float32,
        )

        transform = cv.getPerspectiveTransform(source, destination)
        return cv.warpPerspective(undistorted, transform, (output_width, output_height))

    def _preview_rectified_captures(
        self,
        charuco_board,
        captured_samples: list[CapturedSample],
        camera_matrix,
        dist_coeffs,
    ) -> None:
        """Cycle through rectified captures after calibration."""
        if not captured_samples:
            return

        print("Showing rectified captures. Press q to close, space to pause/resume.")
        window_name = "Rectified Captures"
        paused = False
        index = 0

        try:
            while True:
                rectified = self._rectify_sample(
                    charuco_board,
                    captured_samples[index],
                    camera_matrix,
                    dist_coeffs,
                )
                cv.imshow(window_name, rectified)

                delay_ms = 0 if paused else 800
                key = cv.waitKey(delay_ms) & 0xFF
                if key in (ord("q"), 27):
                    break
                if key == ord(" "):
                    paused = not paused
                    continue
                if key == ord("p"):
                    paused = True
                    index = (index - 1) % len(captured_samples)
                    continue

                index = (index + 1) % len(captured_samples)
        finally:
            cv.destroyWindow(window_name)

    def run(self) -> dict:
        """Run capture, calibrate the camera, and save the result."""
        self._validate()

        charuco_board, charuco_detector = self._build_detector()
        frame_reader = self._build_frame_reader()
        force_full_processing = (
            self.input_source.kind == InputKind.IMAGES or self.max_detection_accuracy
        )
        charuco_capture = CharucoCapture(
            detection_skip_frames=0 if force_full_processing else self.limit_fps,
            detector=charuco_detector,
            manual_capture=self.manual_capture,
            aggressive_refinement=force_full_processing,
        )
        all_corners, all_ids, all_counter, frame_shape, captured_samples = charuco_capture.run(frame_reader)

        if (
            frame_shape == (0, 0)
            or len(all_counter) == 0
            or len(all_corners) == 0
            or len(all_ids) == 0
        ):
            raise RuntimeError(
                "No valid detections captured. Ensure the Charuco board is visible and try again."
            )

        if len(all_counter) < 3:
            raise RuntimeError(
                f"Not enough images for calibration: got {len(all_counter)}, need at least 3."
            )

        print("Calibrating camera...")
        (
            _,
            camera_matrix,
            dist_coeffs,
            _r_vecs,
            _t_vecs,
            _std_dev_intrinsics,
            _std_dev_extrinsics,
            per_view_errors,
        ) = cv.aruco.calibrateCameraArucoExtended(
            np.array(all_corners),
            np.array(all_ids),
            np.array(all_counter),
            charuco_board,
            frame_shape,
            np.array([]),
            np.array([]),
            np.array([]),
            np.array([]),
            np.array([]),
            np.array([]),
            np.array([]),
        )

        dist_coeffs = dist_coeffs.flatten()[:5].reshape(1, 5)
        camera_model = {
            "camera_matrix": camera_matrix.flatten().tolist(),
            "distortion_coefficients": dist_coeffs.flatten().tolist(),
            "avg_reprojection_error": float(np.average(np.array(per_view_errors))),
            "num_images": len(all_counter),
        }

        print("Saving calibration results...")
        with open(self.output_file_path, "w", encoding="utf-8") as output_file:
            output_file.write(json.dumps(camera_model, indent=4))
        print("Calibration complete!")

        if self.preview_rectified_captures:
            self._preview_rectified_captures(
                charuco_board,
                captured_samples,
                camera_matrix,
                dist_coeffs,
            )

        return camera_model


def build_arg_parser() -> argparse.ArgumentParser:
    """Build the command-line parser for calibration runs."""
    arg_parser = argparse.ArgumentParser(prog="calibrate")
    arg_parser.add_argument("--port", help="port of input camera", type=int)
    arg_parser.add_argument(
        "--image_folder", help="folder containing pre-captured calibration images", type=str
    )
    arg_parser.add_argument("--output_file_path", help="path to output file", type=str)
    arg_parser.add_argument("--square_size", help="width of square in meters", type=float)
    arg_parser.add_argument("--marker_size", help="width of marker in meters", type=float)
    arg_parser.add_argument("--width", help="width of board", type=int)
    arg_parser.add_argument("--height", help="height of board", type=int)
    arg_parser.add_argument(
        "--limit_fps",
        help="number of frames to skip between detections (0 = detect every frame)",
        type=int,
        default=0,
    )
    arg_parser.add_argument("--fps", help="fps", type=int, default=50)
    arg_parser.add_argument(
        "--use_wpilib", help="use WPILib cscore to open camera", action="store_true"
    )
    arg_parser.add_argument(
        "--set_fps",
        help="explicitly request device FPS (off by default; can crash on macOS)",
        action="store_true",
    )
    arg_parser.add_argument(
        "--manual_capture",
        help="render detections continuously and only capture when SPACE is pressed",
        action="store_true",
    )
    arg_parser.add_argument(
        "--max_detection_accuracy",
        help="process every frame with the heaviest detection path",
        action="store_true",
    )
    arg_parser.add_argument(
        "--preview_rectified_captures",
        help="show a post-calibration slideshow of rectified captures",
        action="store_true",
    )
    arg_parser.add_argument(
        "--resolution_width", help="resolution width", type=int, default=640
    )
    arg_parser.add_argument(
        "--resolution_height", help="resolution height", type=int, default=480
    )
    return arg_parser


def calibration_from_args(args: argparse.Namespace) -> Calibration:
    """Create a Calibration object from parsed CLI arguments."""
    if args.port is not None and args.image_folder is not None:
        raise ValueError("Choose only one input source: --port or --image_folder.")

    calibration = (
        Calibration()
        .with_output_file_path(args.output_file_path)
        .with_square_size(args.square_size)
        .with_marker_size(args.marker_size)
        .with_width(args.width)
        .with_height(args.height)
        .with_limit_fps(args.limit_fps)
        .with_fps(args.fps)
        .with_use_wpilib(args.use_wpilib)
        .with_set_fps(args.set_fps)
        .with_manual_capture(args.manual_capture)
        .with_max_detection_accuracy(args.max_detection_accuracy)
        .with_preview_rectified_captures(args.preview_rectified_captures)
        .with_resolution(args.resolution_width, args.resolution_height)
    )

    if args.port is not None:
        calibration.with_input(CalibrationInput.PORT(args.port))
    elif args.image_folder is not None:
        calibration.with_input(CalibrationInput.IMAGES(args.image_folder))

    return calibration


def main() -> int:
    """Run the CLI entry point and return a process code."""
    args = build_arg_parser().parse_args()
    try:
        calibration_from_args(args).run()
    except (RuntimeError, ValueError) as exc:
        print(exc)
        return 2
    return 0


if __name__ == "__main__":
    sys.exit(main())