Python+OpenCV：姿态估计（Pose Estimation）

Python+OpenCV：姿态估计（Pose Estimation）

####################################################################################################
# 姿态估计(Pose Estimation)
def lmc_cv_pose_estimation(method):
    """
        函数功能: 姿态估计(Pose Estimation)。
    """

    # Load previously saved data
    with np.load('Calibration parameters.npz') as X:
        camera_matrix, dist_coeffs, rvecs, tvecs = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]

    # 终止条件(termination criteria)
    criteria = (lmc_cv.TERM_CRITERIA_EPS + lmc_cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    # prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
    objp = np.zeros((6 * 7, 3), np.float32)
    objp[:, :2] = np.mgrid[0:7, 0:6].T.reshape(-1, 2)

    # draw a cube or a three-dimensional coordinate axis
    axis = []
    if 0 == method:
        axis = np.float32([[3, 0, 0], [0, 3, 0], [0, 0, -3]]).reshape(-1, 3)

    if 1 == method:
        axis = np.float32([[0, 0, 0], [0, 3, 0], [3, 3, 0], [3, 0, 0],
                           [0, 0, -3], [0, 3, -3], [3, 3, -3], [3, 0, -3]])

    # 读取所有匹配的图像
    stacking_images = []
    corners_image = []
    gray_image = []
    images = glob.glob('D:/99-Research/TestData/cv/stereo/case1/left*.png')
    for image_name in images:
        image = lmc_cv.imread(image_name, lmc_cv.IMREAD_UNCHANGED)
        image = lmc_cv.cvtColor(image, lmc_cv.COLOR_BGR2RGB)
        corners_image = image.copy()
        gray_image = lmc_cv.cvtColor(image, lmc_cv.COLOR_BGR2GRAY)

        # Find the chess board corners
        ret, corners = lmc_cv.findChessboardCorners(gray_image, (7, 6), None)

        # If found, add object points, image points (after refining them)
        if ret:
            # 亚像素角点
            corners2 = lmc_cv.cornerSubPix(gray_image, corners, (11, 11), (-1, -1), criteria)

            # Find the rotation and translation vectors.
            ret, rvecs, tvecs = lmc_cv.solvePnP(objp, corners2, camera_matrix, dist_coeffs)

            # project 3D points to image plane
            imgpts, jac = lmc_cv.projectPoints(axis, rvecs, tvecs, camera_matrix, dist_coeffs)

            # 绘制三维坐标轴
            corners_image = lmc_cv_draw_3d_coordinate_axis(method, corners_image, corners2, imgpts)

            # stacking images side-by-side
            stacking_image = np.hstack((image, corners_image))
            stacking_images.append(stacking_image)

    # 显示图像
    for i in range(len(stacking_images)):
        pyplot.figure('Pose Estimation %d' % (i + 1), figsize=(16, 9))
        # pyplot.subplot(1, 1, 1)
        pyplot.imshow(stacking_images[i], 'gray')
        pyplot.title('Pose Estimation')
        pyplot.xticks([])
        pyplot.yticks([])
        pyplot.savefig('%02d.png' % (i + 1))
    pyplot.show()

####################################################################################################
# 绘制3D轴
def lmc_cv_draw_3d_coordinate_axis(method, image, corners, imgpts):
    if 0 == method:
        corners = np.int32(corners)
        imgpts = np.int32(imgpts)
        corner = tuple(corners[0].ravel())
        image = lmc_cv.line(image, corner, tuple(imgpts[0].ravel()), (255, 0, 0), 5)
        image = lmc_cv.line(image, corner, tuple(imgpts[1].ravel()), (0, 255, 0), 5)
        image = lmc_cv.line(image, corner, tuple(imgpts[2].ravel()), (0, 0, 255), 5)

    if 1 == method:
        imgpts = np.int32(imgpts).reshape(-1, 2)
        # draw ground floor in green
        image = lmc_cv.drawContours(image, [imgpts[:4]], -1, (0, 255, 0), -3)
        # draw pillars in blue color
        for i, j in zip(range(4), range(4, 8)):
            image = lmc_cv.line(image, tuple(imgpts[i]), tuple(imgpts[j]), 255, 3)
        # draw top layer in red color
        image = lmc_cv.drawContours(image, [imgpts[4:]], -1, (0, 0, 255), 3)

    return image