Add AprilTag calibration, assets, and tracking demo
This commit is contained in:
174
calibrate_charuco.py
Normal file
174
calibrate_charuco.py
Normal file
@@ -0,0 +1,174 @@
|
||||
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())
|
||||
Reference in New Issue
Block a user