该程序需要用到camera_mtx.txt为相机内参数矩阵文本文档，camera_dist.txt为畸变参数矩阵文本文档，TEST.png为标记图像。生成方法在上一章

AR 算法

QQ图片20200426102332.png

ARDemo类实现绘制摄像机帧图像，三维模型

import cv2
import numpy as np
from OpenGL.GL import *
from OpenGL.GLUT import *
from PIL import Image
import logging
from pattern_detector import PatternDetector

class ARDemo:
    def __init__(self, mark_image_name):
        self.camera_mtx = np.loadtxt('camera_mtx.txt')
        self.camera_dist = np.loadtxt('camera_dist.txt')
        self.camera = cv2.VideoCapture(0+cv2.CAP_DSHOW)
        self.camera.open(0+cv2.CAP_DSHOW)
        self.image_width = int(self.camera.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.image_height = int(self.camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.scene_image = None
        self.scene_tex_id = 0       # 帧图像纹理句柄
        self.mark_image = cv2.imread(mark_image_name)
        self.decetor = PatternDetector(self.mark_image, self.camera_mtx, self.camera_dist)
        self.proj_mat = self.decetor.get_gl_proj_mat(self.image_width, self.image_height)
        self.modelview_mat = np.eye(4).flatten()
        self.wait_count = 0
        glutInit()
        glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH)
        glutInitWindowPosition(0, 0)
        glutInitWindowSize(self.image_width, self.image_height)
        glutCreateWindow('AR Demo')
        glutDisplayFunc(self.display_event)
        glutIdleFunc(self.display_event)
        glutWMCloseFunc(self.close_event)
        self.init_gl()

    # 异常处理装饰器, 捕捉func执行时产生的异常, 在异常发生后关闭摄像头并退出
    def process_exception(func):
        def wrap(self):
            try:
                func(self)
            except Exception as e:
                if self.camera:
                    self.camera.release()
                logging.exception(e)
                exit(-1)
        return wrap

    # 初始化OpenGL
    def init_gl(self):
        glEnable(GL_DEPTH_TEST)
        self.scene_tex_id = glGenTextures(1)
        glBindTexture(GL_TEXTURE_2D, self.scene_tex_id)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, self.image_width, self.image_height,
                     0, GL_RGB, GL_UNSIGNED_BYTE, [])
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
        glTexEnvi(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_REPLACE)

    # 绘制摄像机帧图像
    def draw_scene_image(self, image):
        bg_image = Image.fromarray(image)
        width, height = bg_image.size
        bg_data = bg_image.tobytes('raw', 'BGR', 0, -1)
        glBindTexture(GL_TEXTURE_2D, self.scene_tex_id)
        glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE, bg_data)

        glEnable(GL_TEXTURE_2D)
        glDisable(GL_DEPTH_TEST)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()
        # 绘制平面
        glBegin(GL_QUADS)
        glTexCoord2f(0.0, 0.0), glVertex3f(-1.0, -1.0, -1.0)
        glTexCoord2f(1.0, 0.0), glVertex3f(1.0, -1.0, -1.0)
        glTexCoord2f(1.0, 1.0), glVertex3f(1.0, 1.0, -1.0)
        glTexCoord2f(0.0, 1.0), glVertex3f(-1.0, 1.0, -1.0)
        glEnd()
        glDisable(GL_TEXTURE_2D)
        glEnable(GL_DEPTH_TEST)

    # 显示三维模型
    def draw_model(self):
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        glLoadMatrixf(self.proj_mat)
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()
        glLoadMatrixf(self.modelview_mat)
        glTranslate(0.5, 0.5, -0.5)

        glEnable(GL_BLEND)
        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
        glColor4f(0, 0.8, 0, 0.5)
        glutSolidCube(1)
        glDisable(GL_BLEND)

        glDisable(GL_DEPTH_TEST)
        glLineWidth(3)
        glColor3f(0.6, 0, 0)
        glutWireCube(1)
        glEnable(GL_DEPTH_TEST)

    # 获取三维模型的视图模型矩阵
    def get_gl_modelview_mat(self):
        if not self.decetor.find_pattern(self.scene_image):
            return None
        if not self.decetor.compute_pose():
            return None
        return self.decetor.get_gl_modelview_mat()

    # 显示回调函数
    @process_exception
    def display_event(self):
        success, self.scene_image = self.camera.read()
        if not success:
            return
        glClearColor(0, 0, 0, 0)
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        self.draw_scene_image(self.scene_image)
        view_mat = self.get_gl_modelview_mat()
        if view_mat is None:
            if self.wait_count <= 0:
                glutSwapBuffers()
                return
            self.wait_count -= 1
        else:
            self.wait_count = 5
            self.modelview_mat = view_mat
        # 绘制模型
        self.draw_model()
        glutSwapBuffers()

    # 窗口关闭回调函数
    def close_event(self):
        if self.camera:
            self.camera.release()

    # 运行ARDemo
    def run(self):
        glutMainLoop()


if __name__ == '__main__':
    ar_demo = ARDemo('TEST.png')
    ar_demo.run()

PatternDetectoe类实现寻找AR标记物，并计算投影和位姿矩阵

import cv2
import numpy as np

# 寻找AR标记物, 并且计算投影和位姿矩阵
class PatternDetector:

    def __init__(self, mark_image, camera_mtx, camera_dist, MIN_MATCH_COUNT=16):
        self.camera_mtx = camera_mtx
        self.camera_dist = camera_dist
        gray_mark_image = cv2.cvtColor(mark_image, cv2.COLOR_BGR2GRAY)
        self.mark_image = gray_mark_image
        self.mark_kp, self.mark_des = self._extract_features(gray_mark_image)
        self.scene_image = None
        self.scene_image_pts = None
        self.MIN_MATCH_COUNT = MIN_MATCH_COUNT
        self.homo_reproj_threshold = 5
        self.rvec = None
        self.tvec = None

    # 从图像中提取角点和特征描述子
    def _extract_features(self, image):
        detector = cv2.xfeatures2d.SURF_create(200)
        kp, des = detector.detectAndCompute(image, None)
        return kp, des

    # 匹配特征描述子
    def _get_matches(self, des1, des2):
        if des1 is None or des2 is None:
            return []
        # 定义flann匹配器
        FLANN_INDEX_KDTREE = 0
        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        search_params = dict(checks=50)
        matcher = cv2.FlannBasedMatcher(index_params, search_params)
        # 匹配特征描述子
        matches = matcher.knnMatch(des1, des2, k=2)
        # 计算最近与次近距离的比值, 挑选小于阈值的匹配对
        good_matches = []
        for m1, m2 in matches:
            ratio = m1.distance / m2.distance
            if ratio < 0.7:
                good_matches.append(m1)

        return good_matches

    def _find_homography(self, kp1, des1, kp2, des2, matches):
        src_pts = np.float32([kp1[m.queryIdx].pt for m in matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in matches]).reshape(-1, 1, 2)
        homography_mat, _ = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, self.homo_reproj_threshold)
        return homography_mat

    def find_pattern(self, scene_image):
        self.scene_image_pts = None
        gray_scene_image = cv2.cvtColor(scene_image, cv2.COLOR_BGR2GRAY)
        kp1, des1 = self.mark_kp, self.mark_des
        kp2, des2 = self._extract_features(gray_scene_image)
        matches = self._get_matches(des1, des2)
        if len(matches) < self.MIN_MATCH_COUNT:
            return False

        print('matches = %d' % len(matches))
        homography_mat = self._find_homography(kp1, des1, kp2, des2, matches)
        if homography_mat is None:
            return False

        h, w = self.mark_image.shape[:2]
        pts = np.float32([[0, 0], [0, h-1], [w-1, h-1], [w-1, 0]]).reshape(-1, 1, 2)
        self.scene_image_pts = cv2.perspectiveTransform(pts, homography_mat)

        return True

    def compute_pose(self):
        image_pts = self.scene_image_pts
        if image_pts is None:
            return False
        # 准备三维点坐标(0,0,0), (1,0,0), (1,0,0), (1,1,0)
        obj_pts = np.float32([[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0]])
        # 用PnP求解位姿矩阵
        ret, self.rvec, self.tvec = \
            cv2.solvePnP(obj_pts, image_pts, self.camera_mtx, self.camera_dist, False, cv2.SOLVEPNP_P3P)
        return ret

    def get_gl_proj_mat(self, width, height):
        proj_mat = np.zeros(shape=(4, 4), dtype=np.float32)
        fx = self.camera_mtx[0][0]
        fy = self.camera_mtx[1][1]
        cx = self.camera_mtx[0][-1]
        cy = self.camera_mtx[1][-1]
        near = 0.1
        far = 100.0
        proj_mat[0][0] = 2*fx / width
        proj_mat[1][1] = 2*fy / height
        proj_mat[0][2] = 1 - (2*cx / width)
        proj_mat[1][2] = (2*cy / height) - 1
        proj_mat[2][2] = -(far + near) / (far - near)
        proj_mat[3][2] = -1.
        proj_mat[2][3] = -(2*far*near) / (far - near)
        p = proj_mat.T
        return p.flatten()

    def get_gl_modelview_mat(self):
        r_mat, _ = cv2.Rodrigues(self.rvec)
        t_mat = np.hstack((r_mat, self.tvec))
        # 翻转Y轴和Z轴
        flip_mat = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
        t_mat = np.dot(flip_mat, t_mat)
        gl_mat = np.eye(4)
        # 填充齐次矩阵前3行4列
        gl_mat[:3, :] = t_mat
        return gl_mat.T.flatten()

实现效果

AR效果图.png