Home | History | Annotate | Download | only in python2 
      1 #!/usr/bin/env python
      2 
      3 '''
      4 Video capture sample.
      5 
      6 Sample shows how VideoCapture class can be used to acquire video
      7 frames from a camera of a movie file. Also the sample provides
      8 an example of procedural video generation by an object, mimicking
      9 the VideoCapture interface (see Chess class).
     10 
     11 'create_capture' is a convinience function for capture creation,
     12 falling back to procedural video in case of error.
     13 
     14 Usage:
     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 
     22 Synth examples:
     23     synth:bg=../data/lena.jpg:noise=0.1
     24     synth:class=chess:bg=../data/lena.jpg:noise=0.1:size=640x480
     25 
     26 Keys:
     27     ESC    - exit
     28     SPACE  - save current frame to <shot path> directory
     29 
     30 '''
     31 
     32 import numpy as np
     33 from numpy import pi, sin, cos
     34 
     35 import cv2
     36 
     37 # built-in modules
     38 from time import clock
     39 
     40 # local modules
     41 import common
     42 
     43 class 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 
     81 class 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 
    129 classes = dict(chess=Chess)
    130 
    131 presets = 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 
    138 def 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 
    170 if __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