用python做简易AR

实现用现实中手持红色的笔接屏幕中掉落的小方块；

import pyglet
from pyglet.window import key
import cv2 as cv
import pymunk
from pymunk.pyglet_util import DrawOptions
import numpy as np

opt = DrawOptions()
cap = cv.VideoCapture(0)

def get_pencil_pos(image):
    HSV = cv.cvtColor(image,cv.COLOR_BGR2HSV)
    lower_red = np.array([0,195,35])
    upper_red = np.array([7,230,255])
    gray = cv.inRange(HSV, lower_red, upper_red)
    edges = cv.Canny(gray,50,150,apertureSize=3)
    try:
        lines = cv.HoughLinesP(edges,1,np.pi/180,60,minLineLength=60,maxLineGap=10)[:,0,:]
        if len(lines)>0:
            x1,y1,x2,y2 = lines[0]
            return 620-x1,480-y1,620-x2,480-y2
    except Exception as e:
        pass

class App(pyglet.window.Window):
    def __init__(self):
        super().__init__()
        _, self.image = cap.read()
        self._image =cv.flip(self.image,-1).tostring()
        self.img = pyglet.image.ImageData(640,480,"BGR",self._image)
        self.pencil_pos = [0,0,0,0]
        self.space = pymunk.Space()
        self.space.gravity = 0,-100
        self.body = pymunk.Body(1,1666)
        self.body.position = 300,400
        self.poly = pymunk.Poly.create_box(self.body,(50,50))
        self.poly.elasticity = 1
        self.poly.friction = 0.6
        self.brick_body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)
        self.brick_body.position = 573, 77
        self.brick_shape = pymunk.Poly.create_box(self.brick_body, (100,10))
        self.brick_shape.elasticity = 0.7
        self.brick_shape.friction = 0.6
        self.space.add(self.body, self.poly,self.brick_body,self.brick_shape)
        
        

    def update(self,dt):
        _, self.image = cap.read()
        self.pencil_pos = get_pencil_pos(self.image)
        try:
            self.brick_body.position = self.pencil_pos[:2]
        except Exception as e:
            pass

        self._image =cv.flip(self.image,-1).tostring()
        self.img = pyglet.image.ImageData(640,480,"BGR",self._image)
        self.space.step(dt)

    def on_draw(self):
        self.clear()
        self.img.blit(0,0)
        self.space.debug_draw(opt)
        
    def on_key_press(self,symbol, modifiers):
        if symbol == key.A:
            print('The "A" key was pressed.')
        elif symbol == key.LEFT:
            print('The left arrow key was pressed.')
        elif symbol == key.ENTER:
            print('The enter key was pressed.')

app = App()
pyglet.clock.schedule_interval(app.update, 1/30)
pyglet.app.run()
cap.release()