175 lines
6.8 KiB
Python
175 lines
6.8 KiB
Python
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())
|