scripts/calibrate_camera.py

"""Calibrate camera pose from RGB marker observations and arm joint encoders."""
import json
import os
from pathlib import Path
import urllib.request

import cv2
import mujoco
import numpy as np
from PIL import Image

from block_vision import detect_blocks, detect_marker, detect_fiducials, annotate

ROOT = Path(__file__).resolve().parent.parent
RESULTS = ROOT / "results"


def server_url():
    """Resolve the local simulator port at call time for scripts and wrappers."""
    return f"http://127.0.0.1:{int(os.environ.get('SO101_PORT', '8877'))}"


def command(endpoint, **payload):
    request = urllib.request.Request(server_url() + endpoint,
        data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(request, timeout=90) as response:
        result = json.load(response)
    if "error" in result:
        raise RuntimeError(result["error"])
    return result


def main():
    RESULTS.mkdir(parents=True, exist_ok=True)
    model = mujoco.MjModel.from_xml_path(str(ROOT / "scene/scene.xml"))
    scratch = mujoco.MjData(model)
    baseline = np.array([0., -.2, .2, 1.4, 0., .6])
    poses = [baseline.copy()]
    for joint, delta in [(0, .2), (0, -.2), (1, -.18), (2, .18), (3, -.2), (4, .25)]:
        pose = baseline.copy(); pose[joint] += delta; poses.append(pose)
    for pan, lift, elbow, wrist in [(.35, -.45, .1, 1.2), (-.3, -.45, .2, 1.1),
                                    (.1, -.65, .25, 1.05), (.25, -.35, .45, .9),
                                    (-.2, -.55, .45, 1.0), (.3, -.25, .1, 1.2)]:
        poses.append(np.array([pan, lift, elbow, wrist, 0, .6]))
    observations = []
    for index, pose in enumerate(poses):
        result = command("/move", target=pose.tolist(), seconds=1.1, label=f"calibrate-{index}")
        rgb = np.asarray(Image.open(result["image"]).convert("RGB"))
        marker = detect_marker(rgb)
        if marker is None:
            print(f"Sample {index}: marker occluded; skipped", flush=True); continue
        scratch.qpos[:6] = result["qpos"]
        mujoco.mj_forward(model, scratch)
        xyz = scratch.geom("vision_marker").xpos.copy()
        observations.append({"index": index, "qpos": result["qpos"], "xyz": xyz.tolist(),
                             "pixel": marker.centroid.tolist(), "image": result["image"]})
        print(f"Sample {index}: marker pixel {marker.centroid.round(1).tolist()}", flush=True)
    xyz = np.array([sample["xyz"] for sample in observations], dtype=np.float64)
    pixels = np.array([sample["pixel"] for sample in observations], dtype=np.float64)
    focal = 720 / (2 * np.tan(np.deg2rad(42) / 2))
    intrinsic = np.array([[focal, 0, 479.5], [0, focal, 359.5], [0, 0, 1.]])
    if len(xyz) < 6:
        raise RuntimeError("Too few visible calibration samples")
    ok, rotation, translation, inliers = cv2.solvePnPRansac(xyz, pixels, intrinsic, None,
        flags=cv2.SOLVEPNP_EPNP, reprojectionError=3., iterationsCount=500)
    if not ok or inliers is None or len(inliers) < 6:
        raise RuntimeError("Camera pose calibration did not converge")
    idx = inliers.ravel()
    rotation, translation = cv2.solvePnPRefineLM(xyz[idx], pixels[idx], intrinsic, None, rotation, translation)
    final = command("/move", target=baseline.tolist(), seconds=1.2, label="calibrated-home")
    rgb = np.asarray(Image.open(final["image"]).convert("RGB"))
    # Anchor the calibration at the work surface as well as at the moving tip.
    dots = detect_fiducials(rgb)
    board = np.array([[.11, -.16, .0015], [.37, -.16, .0015],
                      [.37, .23, .0015], [.11, .23, .0015]])
    board_pixels, _ = cv2.projectPoints(board, rotation, translation, intrinsic, None)
    for point, predicted in zip(board, board_pixels.reshape(-1, 2)):
        if not dots: break
        closest = min(dots, key=lambda dot: np.linalg.norm(dot.centroid - predicted))
        if np.linalg.norm(closest.centroid - predicted) < 8:
            observations.append({"index": "board", "xyz": point.tolist(), "pixel": closest.centroid.tolist(), "image": final["image"]})
            dots = [dot for dot in dots if dot is not closest]
    xyz = np.array([sample["xyz"] for sample in observations], dtype=np.float64)
    pixels = np.array([sample["pixel"] for sample in observations], dtype=np.float64)
    idx = np.arange(len(xyz))
    rotation, translation = cv2.solvePnPRefineLM(xyz, pixels, intrinsic, None, rotation, translation)
    projected, _ = cv2.projectPoints(xyz, rotation, translation, intrinsic, None)
    errors = np.linalg.norm(projected.reshape(-1, 2) - pixels, axis=1)
    result = {"intrinsic": intrinsic.tolist(), "rotation": rotation.ravel().tolist(),
              "translation": translation.ravel().tolist(), "rms_px": float(np.sqrt(np.mean(errors[idx]**2))),
              "inliers": len(idx), "observations": observations,
              "method": "RGB magenta-marker detections + encoder-derived robot marker positions; no object-state queries"}
    (RESULTS / "camera-calibration.json").write_text(json.dumps(result, indent=2))
    print(f"Calibration: {len(idx)}/{len(xyz)} inliers, {result['rms_px']:.3f} px RMS", flush=True)
    blocks = detect_blocks(rgb)
    Image.fromarray(annotate(rgb, blocks)).save(RESULTS / "calibrated-camera.png")


if __name__ == "__main__":
    main()