• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2
3'''
4Planar augmented reality
5==================
6
7This sample shows an example of augmented reality overlay over a planar object
8tracked by PlaneTracker from plane_tracker.py. solvePnP funciton is used to
9estimate the tracked object location in 3d space.
10
11video: http://www.youtube.com/watch?v=pzVbhxx6aog
12
13Usage
14-----
15plane_ar.py [<video source>]
16
17Keys:
18   SPACE  -  pause video
19   c      -  clear targets
20
21Select a textured planar object to track by drawing a box with a mouse.
22Use 'focal' slider to adjust to camera focal length for proper video augmentation.
23'''
24
25import numpy as np
26import cv2
27import video
28import common
29from plane_tracker import PlaneTracker
30
31
32ar_verts = np.float32([[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0],
33                       [0, 0, 1], [0, 1, 1], [1, 1, 1], [1, 0, 1],
34                       [0, 0.5, 2], [1, 0.5, 2]])
35ar_edges = [(0, 1), (1, 2), (2, 3), (3, 0),
36            (4, 5), (5, 6), (6, 7), (7, 4),
37            (0, 4), (1, 5), (2, 6), (3, 7),
38            (4, 8), (5, 8), (6, 9), (7, 9), (8, 9)]
39
40class App:
41    def __init__(self, src):
42        self.cap = video.create_capture(src)
43        self.frame = None
44        self.paused = False
45        self.tracker = PlaneTracker()
46
47        cv2.namedWindow('plane')
48        cv2.createTrackbar('focal', 'plane', 25, 50, common.nothing)
49        self.rect_sel = common.RectSelector('plane', self.on_rect)
50
51    def on_rect(self, rect):
52        self.tracker.add_target(self.frame, rect)
53
54    def run(self):
55        while True:
56            playing = not self.paused and not self.rect_sel.dragging
57            if playing or self.frame is None:
58                ret, frame = self.cap.read()
59                if not ret:
60                    break
61                self.frame = frame.copy()
62
63            vis = self.frame.copy()
64            if playing:
65                tracked = self.tracker.track(self.frame)
66                for tr in tracked:
67                    cv2.polylines(vis, [np.int32(tr.quad)], True, (255, 255, 255), 2)
68                    for (x, y) in np.int32(tr.p1):
69                        cv2.circle(vis, (x, y), 2, (255, 255, 255))
70                    self.draw_overlay(vis, tr)
71
72            self.rect_sel.draw(vis)
73            cv2.imshow('plane', vis)
74            ch = cv2.waitKey(1) & 0xFF
75            if ch == ord(' '):
76                self.paused = not self.paused
77            if ch == ord('c'):
78                self.tracker.clear()
79            if ch == 27:
80                break
81
82    def draw_overlay(self, vis, tracked):
83        x0, y0, x1, y1 = tracked.target.rect
84        quad_3d = np.float32([[x0, y0, 0], [x1, y0, 0], [x1, y1, 0], [x0, y1, 0]])
85        fx = 0.5 + cv2.getTrackbarPos('focal', 'plane') / 50.0
86        h, w = vis.shape[:2]
87        K = np.float64([[fx*w, 0, 0.5*(w-1)],
88                        [0, fx*w, 0.5*(h-1)],
89                        [0.0,0.0,      1.0]])
90        dist_coef = np.zeros(4)
91        ret, rvec, tvec = cv2.solvePnP(quad_3d, tracked.quad, K, dist_coef)
92        verts = ar_verts * [(x1-x0), (y1-y0), -(x1-x0)*0.3] + (x0, y0, 0)
93        verts = cv2.projectPoints(verts, rvec, tvec, K, dist_coef)[0].reshape(-1, 2)
94        for i, j in ar_edges:
95            (x0, y0), (x1, y1) = verts[i], verts[j]
96            cv2.line(vis, (int(x0), int(y0)), (int(x1), int(y1)), (255, 255, 0), 2)
97
98
99if __name__ == '__main__':
100    print __doc__
101
102    import sys
103    try:
104        video_src = sys.argv[1]
105    except:
106        video_src = 0
107    App(video_src).run()
108