import argparse import json from pathlib import Path import cv2 import numpy as np IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"} def parse_args() -> argparse.Namespace: ap = argparse.ArgumentParser(description="Select good ChArUco calibration images and calibrate the camera.") ap.add_argument("--assets-dir", default="assets", help="Directory containing calibration images.") ap.add_argument("--output-dir", default="calibration_output", help="Directory where reports and calibration files are written.") ap.add_argument("--squares-x", type=int, default=7, help="Number of chessboard squares along X.") ap.add_argument("--squares-y", type=int, default=5, help="Number of chessboard squares along Y.") ap.add_argument("--square-length-mm", type=float, default=25.0, help="Square side length in millimeters.") ap.add_argument("--marker-length-mm", type=float, default=18.75, help="Inner marker side length in millimeters.") ap.add_argument("--dictionary", default="DICT_4X4_50", help="OpenCV aruco dictionary name.") ap.add_argument("--min-charuco-corners", type=int, default=12, help="Minimum detected ChArUco corners required to accept an image.") ap.add_argument("--min-markers", type=int, default=8, help="Minimum detected ArUco markers required to accept an image.") ap.add_argument("--min-images", type=int, default=8, help="Minimum accepted images required before calibration.") return ap.parse_args() def get_dictionary(name: str): if not hasattr(cv2.aruco, name): raise ValueError(f"Unknown dictionary: {name}") return cv2.aruco.getPredefinedDictionary(getattr(cv2.aruco, name)) def collect_images(assets_dir: Path) -> list[Path]: return sorted([p for p in assets_dir.iterdir() if p.is_file() and p.suffix.lower() in IMAGE_EXTS]) def write_yaml(path: Path, camera_matrix: np.ndarray, dist_coeffs: np.ndarray) -> None: fs = cv2.FileStorage(str(path), cv2.FILE_STORAGE_WRITE) fs.write("camera_matrix", camera_matrix) fs.write("dist_coeffs", dist_coeffs) fs.release() def main() -> int: args = parse_args() assets_dir = Path(args.assets_dir) output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) if not assets_dir.exists(): raise SystemExit(f"Assets directory not found: {assets_dir}") dictionary = get_dictionary(args.dictionary) board = cv2.aruco.CharucoBoard( (args.squares_x, args.squares_y), args.square_length_mm / 1000.0, args.marker_length_mm / 1000.0, dictionary, ) detector = cv2.aruco.CharucoDetector(board) files = collect_images(assets_dir) if not files: raise SystemExit(f"No images found in {assets_dir}") accepted: list[dict] = [] rejected: list[dict] = [] all_charuco_corners = [] all_charuco_ids = [] image_size = None for path in files: img = cv2.imread(str(path)) if img is None: rejected.append({"file": path.name, "reason": "read_fail", "markers": 0, "charuco_corners": 0}) continue gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if image_size is None: image_size = (gray.shape[1], gray.shape[0]) charuco_corners, charuco_ids, marker_corners, marker_ids = detector.detectBoard(gray) markers = 0 if marker_ids is None else len(marker_ids) corners = 0 if charuco_ids is None else len(charuco_ids) record = { "file": path.name, "markers": int(markers), "charuco_corners": int(corners), "width": int(gray.shape[1]), "height": int(gray.shape[0]), } if markers >= args.min_markers and corners >= args.min_charuco_corners: accepted.append(record) all_charuco_corners.append(charuco_corners) all_charuco_ids.append(charuco_ids) else: reasons = [] if markers < args.min_markers: reasons.append(f"markers<{args.min_markers}") if corners < args.min_charuco_corners: reasons.append(f"charuco<{args.min_charuco_corners}") record["reason"] = ",".join(reasons) if reasons else "rejected" rejected.append(record) (output_dir / "accepted_images.json").write_text(json.dumps(accepted, indent=2), encoding="utf-8") (output_dir / "rejected_images.json").write_text(json.dumps(rejected, indent=2), encoding="utf-8") summary = { "assets_dir": str(assets_dir.resolve()), "dictionary": args.dictionary, "squares_x": args.squares_x, "squares_y": args.squares_y, "square_length_mm": args.square_length_mm, "marker_length_mm": args.marker_length_mm, "min_markers": args.min_markers, "min_charuco_corners": args.min_charuco_corners, "total_images": len(files), "accepted_images": len(accepted), "rejected_images": len(rejected), } print(f"Images found: {len(files)}") print(f"Accepted: {len(accepted)}") print(f"Rejected: {len(rejected)}") if len(accepted) < args.min_images: summary["status"] = "not_enough_images" (output_dir / "calibration_report.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") raise SystemExit(f"Not enough accepted images for calibration: {len(accepted)} < {args.min_images}") retval, camera_matrix, dist_coeffs, rvecs, tvecs = cv2.aruco.calibrateCameraCharuco( charucoCorners=all_charuco_corners, charucoIds=all_charuco_ids, board=board, imageSize=image_size, cameraMatrix=None, distCoeffs=None, ) np.savez( output_dir / "camera_calibration.npz", camera_matrix=camera_matrix, dist_coeffs=dist_coeffs, image_width=image_size[0], image_height=image_size[1], dictionary=args.dictionary, squares_x=args.squares_x, squares_y=args.squares_y, square_length_mm=args.square_length_mm, marker_length_mm=args.marker_length_mm, ) write_yaml(output_dir / "camera_calibration.yaml", camera_matrix, dist_coeffs) summary.update( { "status": "ok", "image_width": image_size[0], "image_height": image_size[1], "rms_reprojection_error": float(retval), "camera_matrix": camera_matrix.tolist(), "dist_coeffs": dist_coeffs.tolist(), } ) (output_dir / "calibration_report.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") print(f"Calibration RMS reprojection error: {retval:.6f}") print(f"Report written to: {output_dir / 'calibration_report.json'}") print(f"NPZ written to: {output_dir / 'camera_calibration.npz'}") print(f"YAML written to: {output_dir / 'camera_calibration.yaml'}") return 0 if __name__ == "__main__": raise SystemExit(main())