1#!/usr/bin/env python 2 3''' 4This module contains some common routines used by other samples. 5''' 6 7import numpy as np 8import cv2 9 10# built-in modules 11import os 12import itertools as it 13from contextlib import contextmanager 14 15image_extensions = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.tiff', '.pbm', '.pgm', '.ppm'] 16 17class Bunch(object): 18 def __init__(self, **kw): 19 self.__dict__.update(kw) 20 def __str__(self): 21 return str(self.__dict__) 22 23def splitfn(fn): 24 path, fn = os.path.split(fn) 25 name, ext = os.path.splitext(fn) 26 return path, name, ext 27 28def anorm2(a): 29 return (a*a).sum(-1) 30def anorm(a): 31 return np.sqrt( anorm2(a) ) 32 33def 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 39def 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 45def 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 55def 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 65def 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 73def 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 77class 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 113cmap_data = { 'jet' : _jet_data } 114 115def 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 130def nothing(*arg, **kw): 131 pass 132 133def clock(): 134 return cv2.getTickCount() / cv2.getTickFrequency() 135 136@contextmanager 137def Timer(msg): 138 print msg, '...', 139 start = clock() 140 try: 141 yield 142 finally: 143 print "%.2f ms" % ((clock()-start)*1000) 144 145class 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 156class 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 192def 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 197def 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 210def getsize(img): 211 h, w = img.shape[:2] 212 return w, h 213 214def mdot(*args): 215 return reduce(np.dot, args) 216 217def 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