如何使用OpenCV将PyBullet仿真坐标投影为渲染帧像素坐标

PyBullet的DIRECT无GUI模式下原生addUserDebugLine()接口无法使用，我们可以通过将3D世界坐标转换为2D像素坐标，再用OpenCV在渲染帧上绘制图形的方式实现同等效果。

原代码偏移的核心原因 #

1. 缺少齐次除法步骤：3D点经过视图、投影矩阵变换后得到齐次坐标，必须除以第四个分量w才能得到合法的标准化设备坐标
2. 像素坐标映射规则错误：标准化设备坐标范围为[-1,1]，需要按规则映射到图像像素范围，同时要适配OpenCV图像y轴向下的坐标约定
3. 矩阵乘法顺序、reshape顺序不符合PyBullet的矩阵存储规则

修复后可运行代码 #

import pybullet as p
import numpy as np
import pybullet_data
import cv2

VIDEO_RESOLUTION = (1280, 720)
MY_COLORS = [(255,0,0), (0,255,0), (0,0,255)]
def capture_frame(base_pos=[0,0,0], _cam_dist=3, _cam_yaw=45, _cam_pitch=-45):
    _render_width, _render_height = VIDEO_RESOLUTION
    view_matrix = p.computeViewMatrixFromYawPitchRoll(
        cameraTargetPosition=base_pos,
        distance=_cam_dist,
        yaw=_cam_yaw,
        pitch=_cam_pitch,
        roll=0,
        upAxisIndex=2)
    proj_matrix = p.computeProjectionMatrixFOV(
        fov=90, aspect=float(_render_width) / _render_height,
        nearVal=0.01, farVal=100.0)
    (_, _, px, _, _) = p.getCameraImage(
        width=_render_width, height=_render_height, viewMatrix=view_matrix,
        projectionMatrix=proj_matrix, renderer=p.ER_TINY_RENDERER)
    rgb_array = np.array(px, dtype=np.uint8)
    rgb_array = np.reshape(rgb_array, (_render_height, _render_width, 4))
    rgb_array = rgb_array[:, :, :3]
    return rgb_array, view_matrix, proj_matrix

def world_to_pixel(world_pos, view_mat, proj_mat, img_w, img_h):
    # 转换为齐次坐标
    world_pos = np.array(world_pos + [1.0])
    # PyBullet的矩阵为列优先存储，reshape使用F顺序
    view_mat = np.array(view_mat).reshape((4,4), order='F')
    proj_mat = np.array(proj_mat).reshape((4,4), order='F')
    # 变换到裁剪空间
    clip_pos = proj_mat @ view_mat @ world_pos
    # 齐次除法得到标准化设备坐标
    ndc_pos = clip_pos[:3] / clip_pos[3]
    # 映射到像素坐标，翻转y轴适配OpenCV坐标约定
    x = (ndc_pos[0] + 1) * 0.5 * img_w
    y = (1 - ndc_pos[1]) * 0.5 * img_h
    return int(round(x)), int(round(y))

def render():
    frame, vmat, pmat = capture_frame()
    w, h = VIDEO_RESOLUTION
    # 计算原点和三轴端点的像素坐标
    origin_px = world_to_pixel([0,0,0], vmat, pmat, w, h)
    x_end_px = world_to_pixel([1,0,0], vmat, pmat, w, h)
    y_end_px = world_to_pixel([0,1,0], vmat, pmat, w, h)
    z_end_px = world_to_pixel([0,0,1], vmat, pmat, w, h)
    # 绘制三轴
    cv2.line(frame, origin_px, x_end_px, color=MY_COLORS[0], thickness=2)
    cv2.line(frame, origin_px, y_end_px, color=MY_COLORS[1], thickness=2)
    cv2.line(frame, origin_px, z_end_px, color=MY_COLORS[2], thickness=2)
    cv2.imwrite("my_rendering.jpg", frame)

if __name__ == '__main__':
    physicsClient = p.connect(p.DIRECT)
    p.setAdditionalSearchPath(pybullet_data.getDataPath())
    p.setGravity(0,0,-10)
    planeId = p.loadURDF("plane.urdf")
    # 加载两个测试机器人
    startPos = [1,0,0.2]
    startOrientation = p.getQuaternionFromEuler([0,0,0])
    boxId1 = p.loadURDF("r2d2.urdf",startPos, startOrientation)
    startPos = [0,2,0.2]
    boxId2 = p.loadURDF("r2d2.urdf",startPos, startOrientation)
    # 步进仿真
    for i in range(2400):
        if i == 2399:
            render()
        p.stepSimulation()
    p.disconnect()

效果说明 #

运行代码后生成的my_rendering.jpg中，红色对应X轴、绿色对应Y轴、蓝色对应Z轴，坐标系会和场景中两个R2D2机器人的实际位置对应，无偏移问题。

内容的提问来源于stack exchange，提问作者avgJoe