Fetch Bracket Bot Stereo Calibration

fetch-stereo-calibration.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.8"
# dependencies = [
#     "opencv-python>=4.8.0",
#     "numpy>=1.24.0",
#     "gdown>=4.7.0",
# ]
# ///
"""
Download a calibration JSON from Google Drive, convert to OpenCV YAML (fisheye), and install locally.

Assumes you are already SSH'ed into the robot (no ssh/scp used).
Writes to: ~/BracketBotOS/bbos/daemons/depth/cache/{stereo_calibration_fisheye.yaml,.json}

Usage:
    python3 upload_calibration_to_robot.py https://drive.google.com/file/d/FILE_ID/view
"""
from __future__ import annotations

import argparse
import json
import math
import sys
import tempfile
from pathlib import Path
import shutil
from typing import Any, Dict, Tuple

import numpy as np
import gdown

# ------------------------ Converter (inlined) ------------------------

def try_import_cv2():
    try:
        import cv2  # type: ignore
        return cv2
    except Exception:
        return None

cv2 = try_import_cv2()

def _get(d: Dict[str,Any], path: str, default=None):
    cur = d
    for key in path.split("/"):
        if not isinstance(cur, dict) or key not in cur:
            return default
        cur = cur[key]
    return cur

def cam_params(cam: Dict[str,Any]) -> Tuple[float,float,float,float,Tuple[float,float,float,float]]:
    prm = _get(cam, "model/ptr_wrapper/data/parameters")
    if prm is None:
        raise KeyError("Missing camera parameters in JSON at model/ptr_wrapper/data/parameters")
    f  = float(prm["f"]["val"])
    ar = float(prm.get("ar", {}).get("val", 1.0) or 1.0)
    cx = float(prm["cx"]["val"])
    cy = float(prm["cy"]["val"])
    k1 = float(prm["k1"]["val"])
    k2 = float(prm["k2"]["val"])
    k3 = float(prm.get("k3", {}).get("val", 0.0) or 0.0)
    k4 = float(prm.get("k4", {}).get("val", 0.0) or 0.0)
    return f, ar, cx, cy, (k1,k2,k3,k4)

def rodrigues_to_R(rvec: np.ndarray) -> np.ndarray:
    rvec = np.asarray(rvec, dtype=np.float64).reshape(3,1)
    if cv2 is not None:
        try:
            R, _ = cv2.Rodrigues(rvec)
            return np.asarray(R, dtype=np.float64)
        except Exception:
            pass
    theta = float(np.linalg.norm(rvec))
    if theta < 1e-12:
        return np.eye(3, dtype=np.float64)
    k = (rvec/theta).reshape(3)
    K = np.array([[0, -k[2], k[1]],[k[2], 0, -k[0]],[-k[1], k[0], 0]], dtype=np.float64)
    R = np.eye(3) + math.sin(theta)*K + (1-math.cos(theta))*(K@K)
    return R

def write_cv_yaml(path: str, mats: Dict[str,np.ndarray]) -> None:
    def mat_tag(name: str, arr: np.ndarray) -> str:
        arr = np.asarray(arr, dtype=np.float64)
        rows, cols = arr.shape
        flat = ", ".join(f"{x:.17g}" for x in arr.reshape(-1))
        return f"""{name}: !!opencv-matrix
   rows: {rows}
   cols: {cols}
   dt: d
   data: [ {flat} ]
"""
    with open(path, "w") as f:
        f.write("%YAML:1.0\n---\n")
        for k in ["mtx_l","dist_l","mtx_r","dist_r","R","T","R1","R2","P1","P2","Q"]:
            f.write(mat_tag(k, mats[k]))

def convert(in_json: str, out_yaml: str, balance: float=0.0, fov_scale: float=1.0, zero_disparity: bool=True) -> None:
    with open(in_json, "r") as f:
        data = json.load(f)
    cal = data.get("Calibration") or data.get("calibration") or data
    cams = cal["cameras"]
    if not isinstance(cams, list) or len(cams) < 2:
        raise ValueError("Expected at least two cameras in Calibration.cameras")

    # Left, Right
    f_l, ar_l, cx_l, cy_l, kd_l = cam_params(cams[0])
    f_r, ar_r, cx_r, cy_r, kd_r = cam_params(cams[1])

    # Image size
    img_size = _get(cams[0], "model/ptr_wrapper/data/CameraModelCRT/CameraModelBase/imageSize") \
               or _get(cams[0], "model/ptr_wrapper/data/CameraModelBase/imageSize")
    if img_size is None:
        raise KeyError("Missing image size in JSON (CameraModelBase/imageSize)")
    w = int(img_size["width"]); h = int(img_size["height"])

    # Intrinsics
    fx_l, fy_l = f_l, f_l * ar_l
    fx_r, fy_r = f_r, f_r * ar_r
    K1 = np.array([[fx_l, 0., cx_l],
                   [0.,   fy_l, cy_l],
                   [0.,   0.,    1. ]], dtype=np.float64)
    K2 = np.array([[fx_r, 0., cx_r],
                   [0.,   fy_r, cy_r],
                   [0.,   0.,    1. ]], dtype=np.float64)
    D1 = np.array(kd_l, dtype=np.float64).reshape(4,1)
    D2 = np.array(kd_r, dtype=np.float64).reshape(4,1)

    # Extrinsics (right wrt left)
    rx = float(_get(cams[1], "transform/rotation/rx", 0.0))
    ry = float(_get(cams[1], "transform/rotation/ry", 0.0))
    rz = float(_get(cams[1], "transform/rotation/rz", 0.0))
    rvec = np.array([rx, ry, rz], dtype=np.float64).reshape(3,1)
    R = rodrigues_to_R(rvec)

    tx = float(_get(cams[1], "transform/translation/x", 0.0))
    ty = float(_get(cams[1], "transform/translation/y", 0.0))
    tz = float(_get(cams[1], "transform/translation/z", 0.0))
    # Use millimeters for T to match many OpenCV YAMLs
    T_mm = np.array([tx, ty, tz], dtype=np.float64).reshape(3,1) * 1000.0

    # Defaults; refined if cv2 available
    R1 = np.eye(3, dtype=np.float64)
    R2 = R.copy()
    P1 = np.array([[fx_l,0,cx_l,0],[0,fy_l,cy_l,0],[0,0,1,0]], dtype=np.float64)
    P2 = P1.copy()
    Q  = np.zeros((4,4), dtype=np.float64)

    if cv2 is not None and hasattr(cv2, "fisheye") and hasattr(cv2.fisheye, "stereoRectify"):
        flags = cv2.CALIB_ZERO_DISPARITY if zero_disparity else 0
        R1, R2, P1, P2, Q = cv2.fisheye.stereoRectify(
            K1, D1, K2, D2, (w,h), R, T_mm,
            flags=flags, balance=float(balance), fov_scale=float(fov_scale)
        )
    elif cv2 is not None:
        # Pinhole fallback
        flags = cv2.CALIB_ZERO_DISPARITY if zero_disparity else 0
        Z = np.zeros((1,5), dtype=np.float64)
        R1, R2, P1, P2, Q, _, _ = cv2.stereoRectify(
            K1, Z, K2, Z, (w,h), R, T_mm, flags=flags, alpha=0.0
        )
    # else: leave defaults

    mats = {
        "mtx_l": K1, "dist_l": D1,
        "mtx_r": K2, "dist_r": D2,
        "R": R, "T": T_mm,
        "R1": R1, "R2": R2, "P1": P1, "P2": P2, "Q": Q,
    }
    write_cv_yaml(out_yaml, mats)

# ------------------------ Orchestration (download + install) ------------------------

def download_calibration(url: str, output_path: Path) -> None:
    print(f"[i] Downloading calibration from: {url}")
    try:
        gdown.download(str(url), str(output_path), quiet=False, fuzzy=True)
    except Exception as e:
        print(f"[!] Error downloading file: {e}", file=sys.stderr)
        sys.exit(1)
    if not output_path.exists() or output_path.stat().st_size == 0:
        print(f"[!] Download failed or produced empty file: {output_path}", file=sys.stderr)
        sys.exit(1)
    print(f"[✓] Downloaded -> {output_path}")

def install_locally(yaml_path: Path, json_path: Path, dest_dir: Path) -> None:
    dest = dest_dir.expanduser().resolve()
    dest.mkdir(parents=True, exist_ok=True)
    final_yaml = dest / "stereo_calibration_fisheye.yaml"
    final_json = dest / "stereo_calibration_fisheye.json"
    shutil.copy2(yaml_path, final_yaml)
    shutil.copy2(json_path, final_json)
    print(f"[✓] Installed YAML -> {final_yaml}")
    print(f"[✓] Stored JSON  -> {final_json}")

def main():
    p = argparse.ArgumentParser(
        description="Download calibration JSON, convert to YAML (fisheye), and install locally on the robot."
        )
    p.add_argument(
    "url",
    help="Google Drive URL of the calibration JSON (gdown-compatible)"
    )
    args = p.parse_args()

    url = args.url
    dest_dir = Path("~/BracketBotOS/bbos/daemons/depth/cache")
    balance = 0.0
    fov_scale = 1.0
    zero_disparity = True

    with tempfile.TemporaryDirectory() as td:
        tmp = Path(td)
        json_path = tmp / "calibration.json"
        yaml_path = tmp / "stereo_calibration_fisheye.yaml"

        download_calibration(args.url, json_path)

        print("[i] Converting JSON -> YAML ...")
        try:
            convert(str(json_path), str(yaml_path),
                    balance=balance, fov_scale=fov_scale,
                    zero_disparity=zero_disparity)
        except Exception as e:
            print(f"[!] Conversion error: {e}", file=sys.stderr)
            sys.exit(1)
        if not yaml_path.exists() or yaml_path.stat().st_size == 0:
            print(f"[!] Conversion produced no output: {yaml_path}", file=sys.stderr)
            sys.exit(1)
        print(f"[✓] Converted -> {yaml_path}")

        install_locally(yaml_path, json_path, dest_dir)

    print("\n✓ Done. Restart the depth daemon to pick up the new calibration.")

if __name__ == "__main__":
    main()