176 lines
4.6 KiB
Python
176 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import json
|
|
import numpy as np
|
|
import cv2
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# Load JSON
|
|
# ------------------------------------------------------------
|
|
|
|
def load_json(path):
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# Robot model: marker centers in world coordinates
|
|
# ------------------------------------------------------------
|
|
|
|
def load_robot_markers(robot_json):
|
|
markers = {}
|
|
for m in robot_json["Marker"]:
|
|
if m.get("on") == "Base":
|
|
mid = int(m["id"])
|
|
markers[mid] = np.array(m["relPos"], dtype=np.float32)
|
|
return markers
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# Marker geometry (world frame)
|
|
# ------------------------------------------------------------
|
|
|
|
def marker_corners_world(center, size_m):
|
|
"""
|
|
Returns 4 corners in consistent OpenCV order:
|
|
TL, TR, BR, BL
|
|
Marker lies in XY plane (z=0)
|
|
"""
|
|
h = size_m / 2.0
|
|
x, y, z = center
|
|
|
|
return np.array([
|
|
[x - h, y + h, z],
|
|
[x + h, y + h, z],
|
|
[x + h, y - h, z],
|
|
[x - h, y - h, z],
|
|
], dtype=np.float32)
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# Build correspondences for one camera
|
|
# ------------------------------------------------------------
|
|
|
|
def build_correspondences(camera_id, scene_markers, robot_markers, marker_size_m):
|
|
|
|
obj_pts = []
|
|
img_pts = []
|
|
|
|
for marker_id, marker_data in scene_markers.items():
|
|
mid = int(marker_id)
|
|
|
|
if mid not in robot_markers:
|
|
continue
|
|
|
|
# Find observations for this camera
|
|
for obs in marker_data.get("observations", []):
|
|
if obs.get("camera_id") == camera_id:
|
|
center = robot_markers[mid]
|
|
obj_corners = marker_corners_world(center, marker_size_m)
|
|
img_corners = np.array(obs["corners_px"], dtype=np.float32)
|
|
|
|
obj_pts.append(obj_corners)
|
|
img_pts.append(img_corners)
|
|
|
|
if len(obj_pts) == 0:
|
|
return None, None
|
|
|
|
obj_pts = np.vstack(obj_pts).astype(np.float32)
|
|
img_pts = np.vstack(img_pts).astype(np.float32)
|
|
|
|
return obj_pts, img_pts
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# Solve PnP
|
|
# ------------------------------------------------------------
|
|
|
|
def solve_camera(obj_pts, img_pts, K, dist):
|
|
|
|
if obj_pts is None or len(obj_pts) < 6:
|
|
raise RuntimeError("Not enough correspondences for PnP")
|
|
|
|
ok, rvec, tvec = cv2.solvePnP(
|
|
obj_pts,
|
|
img_pts,
|
|
K,
|
|
dist,
|
|
flags=cv2.SOLVEPNP_ITERATIVE
|
|
)
|
|
|
|
if not ok:
|
|
raise RuntimeError("solvePnP failed")
|
|
|
|
R, _ = cv2.Rodrigues(rvec)
|
|
return R, tvec
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# Main
|
|
# ------------------------------------------------------------
|
|
|
|
def main():
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument("-scene", required=True)
|
|
parser.add_argument("-robot", required=True)
|
|
parser.add_argument("--marker_size", type=float, default=0.025)
|
|
parser.add_argument("-out", default="camera_poses.json")
|
|
|
|
args = parser.parse_args()
|
|
|
|
scene = load_json(args.scene)
|
|
robot = load_json(args.robot)
|
|
|
|
robot_markers = load_robot_markers(robot)
|
|
|
|
result = {
|
|
"camera_poses": {}
|
|
}
|
|
|
|
# --------------------------------------------------------
|
|
# Each camera independently
|
|
# --------------------------------------------------------
|
|
|
|
for cam_id, cam in scene["cameras"].items():
|
|
|
|
print(f"[INFO] Solving camera {cam_id}")
|
|
|
|
K = np.array(cam["camera_matrix"], dtype=np.float32)
|
|
dist = np.array(cam["distortion_coefficients"], dtype=np.float32)
|
|
|
|
obj_pts, img_pts = build_correspondences(
|
|
cam_id,
|
|
scene["markers"],
|
|
robot_markers,
|
|
args.marker_size
|
|
)
|
|
|
|
if obj_pts is None:
|
|
print(f"[WARN] Camera {cam_id}: no valid markers")
|
|
continue
|
|
|
|
R, t = solve_camera(obj_pts, img_pts, K, dist)
|
|
|
|
result["camera_poses"][cam_id] = {
|
|
"R_world_from_cam": R.tolist(),
|
|
"t_world_from_cam": t.flatten().tolist()
|
|
}
|
|
|
|
print(f"[OK] Camera {cam_id} solved")
|
|
|
|
# --------------------------------------------------------
|
|
# Save
|
|
# --------------------------------------------------------
|
|
|
|
with open(args.out, "w", encoding="utf-8") as f:
|
|
json.dump(result, f, indent=2)
|
|
|
|
print(f"[DONE] Saved -> {args.out}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |