1#!/usr/bin/env python
2
3'''
4Video capture sample.
5
6Sample shows how VideoCapture class can be used to acquire video
7frames from a camera of a movie file. Also the sample provides
8an example of procedural video generation by an object, mimicking
9the VideoCapture interface (see Chess class).
10
11'create_capture' is a convinience function for capture creation,
12falling back to procedural video in case of error.
13
14Usage:
15    video.py [--shotdir <shot path>] [source0] [source1] ...'
16
17    sourceN is an
18     - integer number for camera capture
19     - name of video file
20     - synth:<params> for procedural video
21
22Synth examples:
23    synth:bg=../data/lena.jpg:noise=0.1
24    synth:class=chess:bg=../data/lena.jpg:noise=0.1:size=640x480
25
26Keys:
27    ESC    - exit
28    SPACE  - save current frame to <shot path> directory
29
30'''
31
32import numpy as np
33from numpy import pi, sin, cos
34
35import cv2
36
37# built-in modules
38from time import clock
39
40# local modules
41import common
42
43class VideoSynthBase(object):
44    def __init__(self, size=None, noise=0.0, bg = None, **params):
45        self.bg = None
46        self.frame_size = (640, 480)
47        if bg is not None:
48            self.bg = cv2.imread(bg, 1)
49            h, w = self.bg.shape[:2]
50            self.frame_size = (w, h)
51
52        if size is not None:
53            w, h = map(int, size.split('x'))
54            self.frame_size = (w, h)
55            self.bg = cv2.resize(self.bg, self.frame_size)
56
57        self.noise = float(noise)
58
59    def render(self, dst):
60        pass
61
62    def read(self, dst=None):
63        w, h = self.frame_size
64
65        if self.bg is None:
66            buf = np.zeros((h, w, 3), np.uint8)
67        else:
68            buf = self.bg.copy()
69
70        self.render(buf)
71
72        if self.noise > 0.0:
73            noise = np.zeros((h, w, 3), np.int8)
74            cv2.randn(noise, np.zeros(3), np.ones(3)*255*self.noise)
75            buf = cv2.add(buf, noise, dtype=cv2.CV_8UC3)
76        return True, buf
77
78    def isOpened(self):
79        return True
80
81class Chess(VideoSynthBase):
82    def __init__(self, **kw):
83        super(Chess, self).__init__(**kw)
84
85        w, h = self.frame_size
86
87        self.grid_size = sx, sy = 10, 7
88        white_quads = []
89        black_quads = []
90        for i, j in np.ndindex(sy, sx):
91            q = [[j, i, 0], [j+1, i, 0], [j+1, i+1, 0], [j, i+1, 0]]
92            [white_quads, black_quads][(i + j) % 2].append(q)
93        self.white_quads = np.float32(white_quads)
94        self.black_quads = np.float32(black_quads)
95
96        fx = 0.9
97        self.K = np.float64([[fx*w, 0, 0.5*(w-1)],
98                        [0, fx*w, 0.5*(h-1)],
99                        [0.0,0.0,      1.0]])
100
101        self.dist_coef = np.float64([-0.2, 0.1, 0, 0])
102        self.t = 0
103
104    def draw_quads(self, img, quads, color = (0, 255, 0)):
105        img_quads = cv2.projectPoints(quads.reshape(-1, 3), self.rvec, self.tvec, self.K, self.dist_coef) [0]
106        img_quads.shape = quads.shape[:2] + (2,)
107        for q in img_quads:
108            cv2.fillConvexPoly(img, np.int32(q*4), color, cv2.LINE_AA, shift=2)
109
110    def render(self, dst):
111        t = self.t
112        self.t += 1.0/30.0
113
114        sx, sy = self.grid_size
115        center = np.array([0.5*sx, 0.5*sy, 0.0])
116        phi = pi/3 + sin(t*3)*pi/8
117        c, s = cos(phi), sin(phi)
118        ofs = np.array([sin(1.2*t), cos(1.8*t), 0]) * sx * 0.2
119        eye_pos = center + np.array([cos(t)*c, sin(t)*c, s]) * 15.0 + ofs
120        target_pos = center + ofs
121
122        R, self.tvec = common.lookat(eye_pos, target_pos)
123        self.rvec = common.mtx2rvec(R)
124
125        self.draw_quads(dst, self.white_quads, (245, 245, 245))
126        self.draw_quads(dst, self.black_quads, (10, 10, 10))
127
128
129classes = dict(chess=Chess)
130
131presets = dict(
132    empty = 'synth:',
133    lena = 'synth:bg=../data/lena.jpg:noise=0.1',
134    chess = 'synth:class=chess:bg=../data/lena.jpg:noise=0.1:size=640x480'
135)
136
137
138def create_capture(source = 0, fallback = presets['chess']):
139    '''source: <int> or '<int>|<filename>|synth [:<param_name>=<value> [:...]]'
140    '''
141    source = str(source).strip()
142    chunks = source.split(':')
143    # handle drive letter ('c:', ...)
144    if len(chunks) > 1 and len(chunks[0]) == 1 and chunks[0].isalpha():
145        chunks[1] = chunks[0] + ':' + chunks[1]
146        del chunks[0]
147
148    source = chunks[0]
149    try: source = int(source)
150    except ValueError: pass
151    params = dict( s.split('=') for s in chunks[1:] )
152
153    cap = None
154    if source == 'synth':
155        Class = classes.get(params.get('class', None), VideoSynthBase)
156        try: cap = Class(**params)
157        except: pass
158    else:
159        cap = cv2.VideoCapture(source)
160        if 'size' in params:
161            w, h = map(int, params['size'].split('x'))
162            cap.set(cv2.CAP_PROP_FRAME_WIDTH, w)
163            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, h)
164    if cap is None or not cap.isOpened():
165        print 'Warning: unable to open video source: ', source
166        if fallback is not None:
167            return create_capture(fallback, None)
168    return cap
169
170if __name__ == '__main__':
171    import sys
172    import getopt
173
174    print __doc__
175
176    args, sources = getopt.getopt(sys.argv[1:], '', 'shotdir=')
177    args = dict(args)
178    shotdir = args.get('--shotdir', '.')
179    if len(sources) == 0:
180        sources = [ 0 ]
181
182    caps = map(create_capture, sources)
183    shot_idx = 0
184    while True:
185        imgs = []
186        for i, cap in enumerate(caps):
187            ret, img = cap.read()
188            imgs.append(img)
189            cv2.imshow('capture %d' % i, img)
190        ch = 0xFF & cv2.waitKey(1)
191        if ch == 27:
192            break
193        if ch == ord(' '):
194            for i, img in enumerate(imgs):
195                fn = '%s/shot_%d_%03d.bmp' % (shotdir, i, shot_idx)
196                cv2.imwrite(fn, img)
197                print fn, 'saved'
198            shot_idx += 1
199    cv2.destroyAllWindows()
200