我的这个代码里面已经正确定义了如何读取STL文件并合并,以及如何读取一个个文件夹的图像数据(这部分你直接用我这个代码的。)我需要你改进的是calibration算法,也就是如何正确匹配3D和2D点,并计算intrinsic matrix。然后需要将每一张图像匹配的结果像这个代码一样可视化,左图是stl文件渲染图,右图是读取图像,然后特征点连线一下
#!/usr/bin/env python3
"""
Intrinsic calibration & paired-image visualisation for the da Vinci LND tool
Author: Wenzheng Cheng | last update 2025-06-12
左图完全复用 lnd.py 渲染逻辑(1000×800, elev 0, azim 0, roll 120)。
右图 = mask 等比 resize → 1000×800。
输出: *_pair.jpg (左渲右 mask + 彩线匹配)
"""
import os, argparse, math, xml.etree.ElementTree as ET
import cv2, trimesh, numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from natsort import natsorted
from tqdm import tqdm
# ------------------------------------------------- #
# ------------ 常量与 LND 白名单 ------------------- #
# ------------------------------------------------- #
LND_XML_PATH = "/home/iulian/chole_ws/src/drrobot/mujoco_menagerie/lnd/lnd.xml"
LND_ASSET_DIR = "/home/iulian/chole_ws/src/drrobot/mujoco_menagerie/lnd/assets"
WL = {"jaw_1","jaw_2","jaw_pad_1","jaw_pad_2",
"pitch_mech","pitch_screw",
"pitch_wheel_1","pitch_wheel_2","pitch_wheel_3","pitch_wheel_4",
"wheel_1","wheel_2","wheel_3","wheel_4","yaw_screw"}
# 左图渲染固定参数 —— 与 lnd.py 完全一致
RENDER_W, RENDER_H = 1000, 800
CAM_ELEV, CAM_AZIM, CAM_ROLL = 0.0, 0.0, 120.0
# ------------------------------------------------- #
# ---------------- LND 载入 + 渲染 ------------------ #
# ------------------------------------------------- #
def parse(xml, asset_dir=LND_ASSET_DIR):
root = ET.parse(xml).getroot()
return [os.path.join(asset_dir, m.get("file"))
for m in root.findall(".//asset/mesh")
if m.get("name") in WL]
def load_merge(paths):
meshes = []
for p in paths:
m = trimesh.load_mesh(p, process=False)
if isinstance(m, trimesh.Scene):
m = trimesh.util.concatenate(tuple(m.geometry.values()))
m.apply_scale(1000.0) # m → mm
meshes.append(m)
return trimesh.util.concatenate(meshes)
# === lnd.py 原汁渲染 ===
def _plot_trimesh(ax, mesh):
try:
tgt = max(10_000, int(len(mesh.faces)*0.3))
mesh_sub = mesh.simplify_quadratic_decimation(tgt)
except Exception:
mesh_sub = mesh
v, f = mesh_sub.vertices, mesh_sub.faces
ax.add_collection3d(Poly3DCollection(
v[f], facecolor=[.8,.8,.8], edgecolor=[.4,.4,.4], linewidth=0.15))
span = v.max(0) - v.min(0); cen = v.mean(0); R = span.max()*0.6
for setter,c in zip([ax.set_xlim,ax.set_ylim,ax.set_zlim], cen):
setter(c-R, c+R)
def render_lnd(mesh):
fig = plt.figure(figsize=(RENDER_W/100, RENDER_H/100), dpi=100, facecolor="black")
ax = fig.add_subplot(111, projection='3d', facecolor="black")
ax.view_init(elev=CAM_ELEV, azim=CAM_AZIM, roll=CAM_ROLL)
ax.axis('off')
_plot_trimesh(ax, mesh)
plt.tight_layout(pad=0)
canvas = FigureCanvasAgg(fig); canvas.draw()
# Matplotlib ≥3.8:改用 buffer_rgba(),再丢掉 alpha 通道
buf = np.asarray(canvas.buffer_rgba()) # shape (H,W,4)
img = buf[...,:3].copy() # → uint8 RGB
plt.close(fig)
return img
# ------------------------------------------------- #
# ----------- 数学 / 采样 / PnP 工具函数 ----------- #
# ------------------------------------------------- #
def view_to_rvec(elev, azim, roll):
def Rz(t): return np.array([[ math.cos(t),-math.sin(t),0],
[ math.sin(t), math.cos(t),0],
[0,0,1]])
def Rx(t): return np.array([[1,0,0],
[0, math.cos(t),-math.sin(t)],
[0, math.sin(t), math.cos(t)]])
R = Rz(np.radians(azim)) @ Rx(np.radians(elev)) @ Rz(np.radians(roll))
return cv2.Rodrigues(R)[0].astype(np.float32)
def sample_surface(mesh, n):
pts,_ = trimesh.sample.sample_surface(mesh, n)
return pts.astype(np.float32)
def uniform_mask_points(mask, max_n):
ys,xs = np.where(mask>0)
if len(xs)==0: return np.empty((0,2),np.float32)
if len(xs)>max_n:
sel = np.random.choice(len(xs), max_n, False)
xs,ys = xs[sel], ys[sel]
pts = np.stack([xs,ys],1).astype(np.float32)
pts += np.random.rand(*pts.shape)-0.5
return pts
def pnp(obj,img,K):
ok,r,t,_ = cv2.solvePnPRansac(obj,img,K,None,
flags=cv2.SOLVEPNP_EPNP,iterationsCount=800,reprojectionError=3)
if not ok: raise RuntimeError("PnP fail")
return r,t
def mask_consistent(mask,r,t,K,pts3,max_out=1200):
proj,_ = cv2.projectPoints(pts3,r,t,K,None)
proj = proj.reshape(-1,2).astype(int)
h,w = mask.shape
good = (proj[:,0]>=0)&(proj[:,0]<w)&(proj[:,1]>=0)&(proj[:,1]<h)
proj,obj = proj[good], pts3[good]
keep = mask[proj[:,1],proj[:,0]]>0
obj,proj = obj[keep], proj[keep].astype(np.float32)
if len(obj)>max_out:
sel = np.random.choice(len(obj), max_out, False)
obj,proj = obj[sel], proj[sel]
return obj,proj
# ------------------------------------------------- #
# ---------------- 可视化 (单图) ------------------ #
# ------------------------------------------------- #
def scale_pts(pts, sx, sy):
return (pts * np.array([[sx,sy]])).astype(int)
def draw_pair(mask, proj, img_pts, dense_proj, save_path, lnd_img):
# 1) 把 mask resize → 左图同分辨率
mask_resized = cv2.resize(mask, (RENDER_W, RENDER_H),
interpolation=cv2.INTER_NEAREST)
right_img = cv2.cvtColor(mask_resized, cv2.COLOR_GRAY2BGR)
canvas = np.concatenate([lnd_img, right_img], axis=1)
# 2) 坐标缩放系数
h0,w0 = mask.shape
sx, sy = RENDER_W / w0, RENDER_H / h0
img_pts_s = scale_pts(img_pts, sx, sy) + np.array([RENDER_W,0])
proj_s = scale_pts(proj, sx, sy)
dense_scaled= scale_pts(dense_proj, sx, sy)
for p in dense_scaled: cv2.circle(canvas, tuple(p), 1, (80,80,80), -1)
rng = np.random.RandomState(0)
for (x1,y1),(x2,y2),c in zip(proj_s, img_pts_s,
rng.randint(0,255,(len(img_pts_s),3)).tolist()):
cv2.circle(canvas,(x1,y1),3,c,-1)
cv2.circle(canvas,(x2,y2),3,c,-1)
cv2.line(canvas,(x1,y1),(x2,y2),c,1)
cv2.imwrite(save_path, canvas)
# ------------------------------------------------- #
# ----------------- 数据集遍历工具 ----------------- #
# ------------------------------------------------- #
# ---------- 路径同时兼容 seg_masks 与 left_img_dir ----------
def collect_imgs(video):
seg = os.path.join(video, "seg_masks") if os.path.isdir(os.path.join(video, "seg_masks")) else os.path.join(video, "left_img_dir")
return [os.path.join(seg, f) for f in natsorted(os.listdir(seg))
if f.lower().endswith((".png", ".jpg", ".jpeg", ".bmp"))]
def iterate(root):
return [os.path.join(root,d) for d in natsorted(os.listdir(root))
if os.path.isdir(os.path.join(root,d))]
# ---------- 灰度角点替代 uniform 像素 ----------
def detect_corners(gray, n):
c = cv2.goodFeaturesToTrack(gray, maxCorners=n, qualityLevel=0.01, minDistance=3)
return np.empty((0, 2), np.float32) if c is None else c.reshape(-1, 2).astype(np.float32)
def scale_to_canvas(pts, w, h):
if len(pts) == 0: return pts
lo, hi = pts.min(0), pts.max(0)
c = (lo + hi) / 2; span = np.clip(hi - lo, 1e-3, None)
s = min(0.85 * w / span[0], 0.85 * h / span[1])
return (pts - c) * s + np.array([w / 2, h / 2])
# ------------------------------------------------- #
# --------------------------- main ---------------- #
# ------------------------------------------------- #
def main():
ag = argparse.ArgumentParser()
ag.add_argument("--path", required=True)
ag.add_argument("--vis_dir", default="")
ag.add_argument("--samples", type=int, default=10_000)
ag.add_argument("--max_pts", type=int, default=800)
args = ag.parse_args()
if args.vis_dir: os.makedirs(args.vis_dir, exist_ok=True)
mesh = load_merge(parse(LND_XML_PATH))
lnd_img = render_lnd(mesh)
dense_pts = sample_surface(mesh, args.samples)
rvec_fixed = view_to_rvec(CAM_ELEV,CAM_AZIM,CAM_ROLL)
tvec_zero = np.zeros((3,1),np.float32)
for vid in iterate(args.path):
paths = collect_imgs(vid)
if not paths: continue
first_rgb = cv2.imread(paths[0]); h0, w0 = first_rgb.shape[:2]
K0 = np.array([[0.8 * w0, 0, w0 / 2],
[0, 0.8 * w0, h0 / 2],
[0, 0, 1]], float)
dense_proj0, _ = cv2.projectPoints(dense_pts, rvec_fixed, tvec_zero, K0, None)
dense_proj0 = dense_proj0.reshape(-1, 2)
obj_list, img_list = [], []
for p in tqdm(paths, desc=os.path.basename(vid)):
rgb = cv2.imread(p); gray = cv2.cvtColor(rgb, cv2.COLOR_BGR2GRAY)
img_pts = detect_corners(gray, args.max_pts)
if len(img_pts) < 6: continue
obj_guess = dense_pts[np.random.choice(len(dense_pts), len(img_pts), False)]
try:
r, t = pnp(obj_guess, img_pts, K0)
except RuntimeError: continue
full_mask = np.ones_like(gray, np.uint8) * 255 # 让 mask_consistent 不过滤
obj, img_pts_f = mask_consistent(full_mask, r, t, K0, dense_pts, 1200)
if len(obj) < 6: continue
obj_list.append(obj); img_list.append(img_pts_f)
if args.vis_dir:
proj, _ = cv2.projectPoints(obj, r, t, K0, None)
proj = proj.reshape(-1, 2)
fname = os.path.splitext(os.path.basename(p))[0]
right_vis = cv2.resize(rgb, (RENDER_W, RENDER_H))
canvas = np.concatenate([lnd_img, right_vis], 1)
proj_s = scale_to_canvas(proj, RENDER_W, RENDER_H).astype(int)
img_s = scale_to_canvas(img_pts_f, RENDER_W, RENDER_H).astype(int) + np.array([RENDER_W, 0])
for (x1, y1), (x2, y2) in zip(proj_s, img_s):
cv2.circle(canvas, (x1, y1), 3, (0, 255, 0), -1)
cv2.circle(canvas, (x2, y2), 3, (0, 255, 0), -1)
cv2.line(canvas, (x1, y1), (x2, y2), (0, 255, 0), 1)
cv2.imwrite(os.path.join(args.vis_dir, f"{fname}_pair.jpg"), canvas)
if len(obj_list)<3: continue
flag = getattr(cv2,"CALIB_FIX_SKEW",0)
flags = cv2.CALIB_USE_INTRINSIC_GUESS|flag|cv2.CALIB_ZERO_TANGENT_DIST|\
cv2.CALIB_FIX_K3|cv2.CALIB_FIX_K4|cv2.CALIB_FIX_K5|cv2.CALIB_FIX_K6
rms,K,dist,*_ = cv2.calibrateCamera(
obj_list,img_list,(w0,h0),K0,None,flags=flags,
criteria=(cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_COUNT,100,1e-6))
print(f"\n[VIDEO] {os.path.basename(vid)} RMS={rms:.3f}px")
print("K=\n",K,"\n(k1,k2)=",dist.ravel()[:2])
if __name__ == "__main__":
#python intrinsic_matrix.py --path /home/iulian/chole_ws/data/tissue_lift/tissue_1/tissue_lift/
#python intrinsic_matrix.py --path /home/iulian/chole_ws/data/lift/tissue_1/lift/ --vis_dir /home/iulian/chole_ws/src/drrobot/K_vis --max_pts 800
#python intrinsic_matrix.py --path /home/iulian/chole_ws/data/needle_pickup/tissue_1/needle_pickup/
#python intrinsic_matrix.py --path /home/iulian/chole_ws/data/check --vis_dir /home/iulian/chole_ws/src/drrobot/K_vis --max_pts 800
main()