hermes-bot/ros/hermes_vision/hermes_vision/calibrate.py

"""Camera calibration from a checkerboard, writing the camera_info YAML that camera_ros loads.

  python3 -m hermes_vision.calibrate --cols 9 --rows 6 --square-mm 25

COLS x ROWS are the board's *inner corners* (a board of 10 x 7 squares has 9 x 6). With the robot
running, hold the board in front of the camera at different distances, tilts and positions, keeping the
whole board in view; it collects views that differ enough from those already taken, solves with OpenCV
and writes ~/.ros/camera_info/<camera name>_<w>x<h>.yaml (restart the robot to load it). A good result
has a reprojection error under about 0.5 px. Calibration is only valid for the focus it was taken at,
so set the lens focus first, and for the image size in use.
"""
from __future__ import annotations

import argparse
import os
import time

import cv2
import numpy as np

DEFAULT_NAME = 'ov5647__base_soc_i2c0mux_i2c_1_ov5647_36'      # what camera_ros calls this camera


def board_points(cols: int, rows: int, square_m: float) -> np.ndarray:
    """The board's inner corners in its own frame (z = 0), row by row, in metres."""
    p = np.zeros((rows * cols, 3), np.float32)
    p[:, :2] = np.mgrid[0:cols, 0:rows].T.reshape(-1, 2) * square_m
    return p


def view_signature(corners: np.ndarray, size: tuple[int, int]) -> np.ndarray:
    """A few numbers that describe where and how a board appears: centre, apparent size, and the
    skew of its first row against its last (a tilt cue), all relative to the image."""
    c = corners.reshape(-1, 2)
    w, h = size
    centre = c.mean(axis=0) / (w, h)
    span = (c.max(axis=0) - c.min(axis=0)) / (w, h)
    return np.array([centre[0], centre[1], span[0], span[1], span[0] / max(span[1], 1e-6)])


def novel_enough(sig: np.ndarray, taken: list[np.ndarray], min_distance: float = 0.06) -> bool:
    """True if this view differs enough from every view already taken (a spread of poses matters far
    more than the number of views: twenty copies of one pose pin down nothing)."""
    return all(np.linalg.norm(sig - t) >= min_distance for t in taken)


def solve(obj_pts: list[np.ndarray], img_pts: list[np.ndarray], size: tuple[int, int]):
    """OpenCV calibration. Returns (rms reprojection error in px, camera matrix, distortion)."""
    rms, k, d, _, _ = cv2.calibrateCamera(obj_pts, img_pts, size, None, None)
    return float(rms), k, d.reshape(-1)


def camera_info_yaml(name: str, size: tuple[int, int], k: np.ndarray, d: np.ndarray) -> str:
    """The ROS camera_info YAML (the format camera_calibration_parsers reads)."""
    w, h = size
    d = list(d[:5]) + [0.0] * (5 - len(d[:5]))
    f = lambda xs: ', '.join(f'{x:.9g}' for x in xs)           # noqa: E731
    proj = [k[0, 0], 0, k[0, 2], 0, 0, k[1, 1], k[1, 2], 0, 0, 0, 1, 0]
    return (f'image_width: {w}\nimage_height: {h}\ncamera_name: {name}\n'
            f'camera_matrix:\n  rows: 3\n  cols: 3\n  data: [{f(k.reshape(-1))}]\n'
            f'distortion_model: plumb_bob\n'
            f'distortion_coefficients:\n  rows: 1\n  cols: 5\n  data: [{f(d)}]\n'
            f'rectification_matrix:\n  rows: 3\n  cols: 3\n  data: [1, 0, 0, 0, 1, 0, 0, 0, 1]\n'
            f'projection_matrix:\n  rows: 3\n  cols: 4\n  data: [{f(proj)}]\n')


def find_board(bgr: np.ndarray, cols: int, rows: int):
    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    ok, corners = cv2.findChessboardCornersSB(gray, (cols, rows), cv2.CALIB_CB_EXHAUSTIVE | cv2.CALIB_CB_ACCURACY)
    return (corners.astype(np.float32) if ok else None)


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('--cols', type=int, required=True, help='inner corners across')
    p.add_argument('--rows', type=int, required=True, help='inner corners down')
    p.add_argument('--square-mm', type=float, required=True, help='side of one square, measured')
    p.add_argument('--views', type=int, default=20)
    p.add_argument('--timeout', type=float, default=240)
    p.add_argument('--topic', default='/camera/image_raw/compressed')
    p.add_argument('--name', default=DEFAULT_NAME)
    p.add_argument('--out-dir', default=os.path.expanduser('~/.ros/camera_info'))
    a = p.parse_args()

    import rclpy
    from rclpy.node import Node
    from sensor_msgs.msg import CompressedImage
    rclpy.init()
    node = Node('hermes_calibrate')
    obj = board_points(a.cols, a.rows, a.square_mm / 1000.0)
    taken_sigs, obj_pts, img_pts, state = [], [], [], {'size': None, 'last': 0.0, 'seen': 0}

    def on_image(msg):
        if len(img_pts) >= a.views or time.monotonic() - state['last'] < 0.6:
            return
        bgr = cv2.imdecode(np.frombuffer(msg.data, np.uint8), cv2.IMREAD_COLOR)
        if bgr is None:
            return
        state['size'] = (bgr.shape[1], bgr.shape[0])
        state['last'] = time.monotonic()
        corners = find_board(bgr, a.cols, a.rows)
        if corners is None:
            state['seen'] += 1
            if state['seen'] % 8 == 1:
                print('  no board found: show all of it, flat and well lit, not too close', flush=True)
            return
        sig = view_signature(corners, state['size'])
        if not novel_enough(sig, taken_sigs):
            print('  board found, but too like one already taken: change distance, tilt or position', flush=True)
            return
        taken_sigs.append(sig)
        obj_pts.append(obj)
        img_pts.append(corners)
        print(f'  view {len(img_pts)}/{a.views} taken', flush=True)

    node.create_subscription(CompressedImage, a.topic, on_image, 1)
    print(f'Hold the {a.cols}x{a.rows}-corner board ({a.square_mm:g} mm squares) in front of the camera; '
          f'vary distance, tilt and position. Waiting for {a.views} views...')
    t0 = time.monotonic()
    while len(img_pts) < a.views and time.monotonic() - t0 < a.timeout:
        rclpy.spin_once(node, timeout_sec=0.2)
    rclpy.shutdown()
    if len(img_pts) < 10:
        print(f'only {len(img_pts)} usable views; nothing written (need at least 10, ideally {a.views})')
        return 1
    size = state['size']
    rms, k, d = solve(obj_pts, img_pts, size)
    name = f'{a.name}_{size[0]}x{size[1]}'
    os.makedirs(a.out_dir, exist_ok=True)
    path = os.path.join(a.out_dir, name + '.yaml')
    with open(path, 'w') as f:
        f.write(camera_info_yaml(name, size, k, d))
    print(f'\n{len(img_pts)} views, reprojection error {rms:.3f} px ({"good" if rms < 0.5 else "mediocre: retake with more varied poses"})')
    print(f'focal length {k[0, 0]:.1f} x {k[1, 1]:.1f} px, centre ({k[0, 2]:.0f}, {k[1, 2]:.0f}) of {size[0]}x{size[1]}')
    print(f'wrote {path}\nrestart the robot to load it:  sudo systemctl restart robot')
    return 0


if __name__ == '__main__':
    raise SystemExit(main())