Home | History | Annotate | Download | only in python2
      1 #!/usr/bin/env python
      2 
      3 '''
      4 This module contains some common routines used by other samples.
      5 '''
      6 
      7 import numpy as np
      8 import cv2
      9 
     10 # built-in modules
     11 import os
     12 import itertools as it
     13 from contextlib import contextmanager
     14 
     15 image_extensions = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.tiff', '.pbm', '.pgm', '.ppm']
     16 
     17 class Bunch(object):
     18     def __init__(self, **kw):
     19         self.__dict__.update(kw)
     20     def __str__(self):
     21         return str(self.__dict__)
     22 
     23 def splitfn(fn):
     24     path, fn = os.path.split(fn)
     25     name, ext = os.path.splitext(fn)
     26     return path, name, ext
     27 
     28 def anorm2(a):
     29     return (a*a).sum(-1)
     30 def anorm(a):
     31     return np.sqrt( anorm2(a) )
     32 
     33 def homotrans(H, x, y):
     34     xs = H[0, 0]*x + H[0, 1]*y + H[0, 2]
     35     ys = H[1, 0]*x + H[1, 1]*y + H[1, 2]
     36     s  = H[2, 0]*x + H[2, 1]*y + H[2, 2]
     37     return xs/s, ys/s
     38 
     39 def to_rect(a):
     40     a = np.ravel(a)
     41     if len(a) == 2:
     42         a = (0, 0, a[0], a[1])
     43     return np.array(a, np.float64).reshape(2, 2)
     44 
     45 def rect2rect_mtx(src, dst):
     46     src, dst = to_rect(src), to_rect(dst)
     47     cx, cy = (dst[1] - dst[0]) / (src[1] - src[0])
     48     tx, ty = dst[0] - src[0] * (cx, cy)
     49     M = np.float64([[ cx,  0, tx],
     50                     [  0, cy, ty],
     51                     [  0,  0,  1]])
     52     return M
     53 
     54 
     55 def lookat(eye, target, up = (0, 0, 1)):
     56     fwd = np.asarray(target, np.float64) - eye
     57     fwd /= anorm(fwd)
     58     right = np.cross(fwd, up)
     59     right /= anorm(right)
     60     down = np.cross(fwd, right)
     61     R = np.float64([right, down, fwd])
     62     tvec = -np.dot(R, eye)
     63     return R, tvec
     64 
     65 def mtx2rvec(R):
     66     w, u, vt = cv2.SVDecomp(R - np.eye(3))
     67     p = vt[0] + u[:,0]*w[0]    # same as np.dot(R, vt[0])
     68     c = np.dot(vt[0], p)
     69     s = np.dot(vt[1], p)
     70     axis = np.cross(vt[0], vt[1])
     71     return axis * np.arctan2(s, c)
     72 
     73 def draw_str(dst, (x, y), s):
     74     cv2.putText(dst, s, (x+1, y+1), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 0), thickness = 2, lineType=cv2.LINE_AA)
     75     cv2.putText(dst, s, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), lineType=cv2.LINE_AA)
     76 
     77 class Sketcher:
     78     def __init__(self, windowname, dests, colors_func):
     79         self.prev_pt = None
     80         self.windowname = windowname
     81         self.dests = dests
     82         self.colors_func = colors_func
     83         self.dirty = False
     84         self.show()
     85         cv2.setMouseCallback(self.windowname, self.on_mouse)
     86 
     87     def show(self):
     88         cv2.imshow(self.windowname, self.dests[0])
     89 
     90     def on_mouse(self, event, x, y, flags, param):
     91         pt = (x, y)
     92         if event == cv2.EVENT_LBUTTONDOWN:
     93             self.prev_pt = pt
     94         elif event == cv2.EVENT_LBUTTONUP:
     95             self.prev_pt = None
     96 
     97         if self.prev_pt and flags & cv2.EVENT_FLAG_LBUTTON:
     98             for dst, color in zip(self.dests, self.colors_func()):
     99                 cv2.line(dst, self.prev_pt, pt, color, 5)
    100             self.dirty = True
    101             self.prev_pt = pt
    102             self.show()
    103 
    104 
    105 # palette data from matplotlib/_cm.py
    106 _jet_data =   {'red':   ((0., 0, 0), (0.35, 0, 0), (0.66, 1, 1), (0.89,1, 1),
    107                          (1, 0.5, 0.5)),
    108                'green': ((0., 0, 0), (0.125,0, 0), (0.375,1, 1), (0.64,1, 1),
    109                          (0.91,0,0), (1, 0, 0)),
    110                'blue':  ((0., 0.5, 0.5), (0.11, 1, 1), (0.34, 1, 1), (0.65,0, 0),
    111                          (1, 0, 0))}
    112 
    113 cmap_data = { 'jet' : _jet_data }
    114 
    115 def make_cmap(name, n=256):
    116     data = cmap_data[name]
    117     xs = np.linspace(0.0, 1.0, n)
    118     channels = []
    119     eps = 1e-6
    120     for ch_name in ['blue', 'green', 'red']:
    121         ch_data = data[ch_name]
    122         xp, yp = [], []
    123         for x, y1, y2 in ch_data:
    124             xp += [x, x+eps]
    125             yp += [y1, y2]
    126         ch = np.interp(xs, xp, yp)
    127         channels.append(ch)
    128     return np.uint8(np.array(channels).T*255)
    129 
    130 def nothing(*arg, **kw):
    131     pass
    132 
    133 def clock():
    134     return cv2.getTickCount() / cv2.getTickFrequency()
    135 
    136 @contextmanager
    137 def Timer(msg):
    138     print msg, '...',
    139     start = clock()
    140     try:
    141         yield
    142     finally:
    143         print "%.2f ms" % ((clock()-start)*1000)
    144 
    145 class StatValue:
    146     def __init__(self, smooth_coef = 0.5):
    147         self.value = None
    148         self.smooth_coef = smooth_coef
    149     def update(self, v):
    150         if self.value is None:
    151             self.value = v
    152         else:
    153             c = self.smooth_coef
    154             self.value = c * self.value + (1.0-c) * v
    155 
    156 class RectSelector:
    157     def __init__(self, win, callback):
    158         self.win = win
    159         self.callback = callback
    160         cv2.setMouseCallback(win, self.onmouse)
    161         self.drag_start = None
    162         self.drag_rect = None
    163     def onmouse(self, event, x, y, flags, param):
    164         x, y = np.int16([x, y]) # BUG
    165         if event == cv2.EVENT_LBUTTONDOWN:
    166             self.drag_start = (x, y)
    167         if self.drag_start:
    168             if flags & cv2.EVENT_FLAG_LBUTTON:
    169                 xo, yo = self.drag_start
    170                 x0, y0 = np.minimum([xo, yo], [x, y])
    171                 x1, y1 = np.maximum([xo, yo], [x, y])
    172                 self.drag_rect = None
    173                 if x1-x0 > 0 and y1-y0 > 0:
    174                     self.drag_rect = (x0, y0, x1, y1)
    175             else:
    176                 rect = self.drag_rect
    177                 self.drag_start = None
    178                 self.drag_rect = None
    179                 if rect:
    180                     self.callback(rect)
    181     def draw(self, vis):
    182         if not self.drag_rect:
    183             return False
    184         x0, y0, x1, y1 = self.drag_rect
    185         cv2.rectangle(vis, (x0, y0), (x1, y1), (0, 255, 0), 2)
    186         return True
    187     @property
    188     def dragging(self):
    189         return self.drag_rect is not None
    190 
    191 
    192 def grouper(n, iterable, fillvalue=None):
    193     '''grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx'''
    194     args = [iter(iterable)] * n
    195     return it.izip_longest(fillvalue=fillvalue, *args)
    196 
    197 def mosaic(w, imgs):
    198     '''Make a grid from images.
    199 
    200     w    -- number of grid columns
    201     imgs -- images (must have same size and format)
    202     '''
    203     imgs = iter(imgs)
    204     img0 = imgs.next()
    205     pad = np.zeros_like(img0)
    206     imgs = it.chain([img0], imgs)
    207     rows = grouper(w, imgs, pad)
    208     return np.vstack(map(np.hstack, rows))
    209 
    210 def getsize(img):
    211     h, w = img.shape[:2]
    212     return w, h
    213 
    214 def mdot(*args):
    215     return reduce(np.dot, args)
    216 
    217 def draw_keypoints(vis, keypoints, color = (0, 255, 255)):
    218     for kp in keypoints:
    219             x, y = kp.pt
    220             cv2.circle(vis, (int(x), int(y)), 2, color)
    221