Home | History | Annotate | Download | only in python2 
      1 #!/usr/bin/env python
      2 
      3 '''
      4 Planar augmented reality
      5 ==================
      6 
      7 This sample shows an example of augmented reality overlay over a planar object
      8 tracked by PlaneTracker from plane_tracker.py. solvePnP funciton is used to
      9 estimate the tracked object location in 3d space.
     10 
     11 video: http://www.youtube.com/watch?v=pzVbhxx6aog
     12 
     13 Usage
     14 -----
     15 plane_ar.py [<video source>]
     16 
     17 Keys:
     18    SPACE  -  pause video
     19    c      -  clear targets
     20 
     21 Select a textured planar object to track by drawing a box with a mouse.
     22 Use 'focal' slider to adjust to camera focal length for proper video augmentation.
     23 '''
     24 
     25 import numpy as np
     26 import cv2
     27 import video
     28 import common
     29 from plane_tracker import PlaneTracker
     30 
     31 
     32 ar_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]])
     35 ar_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 
     40 class 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 
     99 if __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