该程序需要用到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