← Back to branch summary

~paparazzi-uav/paparazzi/v5.0-manual

« back to all changes in this revision

Viewing changes to sw/ext/opencv_bebop/opencv/samples/python/plane_ar.py

  • Committer: Paparazzi buildbot
  • Date: 2016-05-18 15:00:29 UTC
  • Revision ID: felix.ruess+docbot@gmail.com-20160518150029-e8lgzi5kvb4p7un9
Manual import commit 4b8bbb730080dac23cf816b98908dacfabe2a8ec from v5.0 branch.

Show diffs side-by-side

added added

removed removed

Lines of Context:
sw/ext/opencv_bebop/opencv/samples/python/plane_ar.py
 
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
# Python 2/3 compatibility
 
26
from __future__ import print_function
 
27
 
 
28
import numpy as np
 
29
import cv2
 
30
import video
 
31
import common
 
32
from plane_tracker import PlaneTracker
 
33
 
 
34
 
 
35
ar_verts = np.float32([[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0],
 
36
                       [0, 0, 1], [0, 1, 1], [1, 1, 1], [1, 0, 1],
 
37
                       [0, 0.5, 2], [1, 0.5, 2]])
 
38
ar_edges = [(0, 1), (1, 2), (2, 3), (3, 0),
 
39
            (4, 5), (5, 6), (6, 7), (7, 4),
 
40
            (0, 4), (1, 5), (2, 6), (3, 7),
 
41
            (4, 8), (5, 8), (6, 9), (7, 9), (8, 9)]
 
42
 
 
43
class App:
 
44
    def __init__(self, src):
 
45
        self.cap = video.create_capture(src)
 
46
        self.frame = None
 
47
        self.paused = False
 
48
        self.tracker = PlaneTracker()
 
49
 
 
50
        cv2.namedWindow('plane')
 
51
        cv2.createTrackbar('focal', 'plane', 25, 50, common.nothing)
 
52
        self.rect_sel = common.RectSelector('plane', self.on_rect)
 
53
 
 
54
    def on_rect(self, rect):
 
55
        self.tracker.add_target(self.frame, rect)
 
56
 
 
57
    def run(self):
 
58
        while True:
 
59
            playing = not self.paused and not self.rect_sel.dragging
 
60
            if playing or self.frame is None:
 
61
                ret, frame = self.cap.read()
 
62
                if not ret:
 
63
                    break
 
64
                self.frame = frame.copy()
 
65
 
 
66
            vis = self.frame.copy()
 
67
            if playing:
 
68
                tracked = self.tracker.track(self.frame)
 
69
                for tr in tracked:
 
70
                    cv2.polylines(vis, [np.int32(tr.quad)], True, (255, 255, 255), 2)
 
71
                    for (x, y) in np.int32(tr.p1):
 
72
                        cv2.circle(vis, (x, y), 2, (255, 255, 255))
 
73
                    self.draw_overlay(vis, tr)
 
74
 
 
75
            self.rect_sel.draw(vis)
 
76
            cv2.imshow('plane', vis)
 
77
            ch = cv2.waitKey(1) & 0xFF
 
78
            if ch == ord(' '):
 
79
                self.paused = not self.paused
 
80
            if ch == ord('c'):
 
81
                self.tracker.clear()
 
82
            if ch == 27:
 
83
                break
 
84
 
 
85
    def draw_overlay(self, vis, tracked):
 
86
        x0, y0, x1, y1 = tracked.target.rect
 
87
        quad_3d = np.float32([[x0, y0, 0], [x1, y0, 0], [x1, y1, 0], [x0, y1, 0]])
 
88
        fx = 0.5 + cv2.getTrackbarPos('focal', 'plane') / 50.0
 
89
        h, w = vis.shape[:2]
 
90
        K = np.float64([[fx*w, 0, 0.5*(w-1)],
 
91
                        [0, fx*w, 0.5*(h-1)],
 
92
                        [0.0,0.0,      1.0]])
 
93
        dist_coef = np.zeros(4)
 
94
        ret, rvec, tvec = cv2.solvePnP(quad_3d, tracked.quad, K, dist_coef)
 
95
        verts = ar_verts * [(x1-x0), (y1-y0), -(x1-x0)*0.3] + (x0, y0, 0)
 
96
        verts = cv2.projectPoints(verts, rvec, tvec, K, dist_coef)[0].reshape(-1, 2)
 
97
        for i, j in ar_edges:
 
98
            (x0, y0), (x1, y1) = verts[i], verts[j]
 
99
            cv2.line(vis, (int(x0), int(y0)), (int(x1), int(y1)), (255, 255, 0), 2)
 
100
 
 
101
 
 
102
if __name__ == '__main__':
 
103
    print(__doc__)
 
104
 
 
105
    import sys
 
106
    try:
 
107
        video_src = sys.argv[1]
 
108
    except:
 
109
        video_src = 0
 
110
    App(video_src).run()