1793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler#!/usr/bin/env python 2793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 3793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler''' 4793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerThis module contains some common routines used by other samples. 5793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler''' 6793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 7793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerimport numpy as np 8793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerimport cv2 9793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 10793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler# built-in modules 11793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerimport os 12793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerimport itertools as it 13793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerfrom contextlib import contextmanager 14793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 15793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerimage_extensions = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.tiff', '.pbm', '.pgm', '.ppm'] 16793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 17793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerclass Bunch(object): 18793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler def __init__(self, **kw): 19793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.__dict__.update(kw) 20793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler def __str__(self): 21793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return str(self.__dict__) 22793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 23793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef splitfn(fn): 24793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler path, fn = os.path.split(fn) 25793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler name, ext = os.path.splitext(fn) 26793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return path, name, ext 27793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 28793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef anorm2(a): 29793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return (a*a).sum(-1) 30793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef anorm(a): 31793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return np.sqrt( anorm2(a) ) 32793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 33793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef homotrans(H, x, y): 34793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler xs = H[0, 0]*x + H[0, 1]*y + H[0, 2] 35793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler ys = H[1, 0]*x + H[1, 1]*y + H[1, 2] 36793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler s = H[2, 0]*x + H[2, 1]*y + H[2, 2] 37793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return xs/s, ys/s 38793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 39793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef to_rect(a): 40793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler a = np.ravel(a) 41793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if len(a) == 2: 42793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler a = (0, 0, a[0], a[1]) 43793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return np.array(a, np.float64).reshape(2, 2) 44793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 45793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef rect2rect_mtx(src, dst): 46793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler src, dst = to_rect(src), to_rect(dst) 47793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cx, cy = (dst[1] - dst[0]) / (src[1] - src[0]) 48793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler tx, ty = dst[0] - src[0] * (cx, cy) 49793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler M = np.float64([[ cx, 0, tx], 50793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler [ 0, cy, ty], 51793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler [ 0, 0, 1]]) 52793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return M 53793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 54793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 55793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef lookat(eye, target, up = (0, 0, 1)): 56793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler fwd = np.asarray(target, np.float64) - eye 57793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler fwd /= anorm(fwd) 58793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler right = np.cross(fwd, up) 59793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler right /= anorm(right) 60793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler down = np.cross(fwd, right) 61793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler R = np.float64([right, down, fwd]) 62793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler tvec = -np.dot(R, eye) 63793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return R, tvec 64793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 65793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef mtx2rvec(R): 66793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler w, u, vt = cv2.SVDecomp(R - np.eye(3)) 67793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler p = vt[0] + u[:,0]*w[0] # same as np.dot(R, vt[0]) 68793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler c = np.dot(vt[0], p) 69793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler s = np.dot(vt[1], p) 70793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler axis = np.cross(vt[0], vt[1]) 71793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return axis * np.arctan2(s, c) 72793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 73793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef draw_str(dst, (x, y), s): 74793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.putText(dst, s, (x+1, y+1), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 0), thickness = 2, lineType=cv2.LINE_AA) 75793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.putText(dst, s, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), lineType=cv2.LINE_AA) 76793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 77793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerclass Sketcher: 78793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler def __init__(self, windowname, dests, colors_func): 79793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.prev_pt = None 80793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.windowname = windowname 81793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.dests = dests 82793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.colors_func = colors_func 83793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.dirty = False 84793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.show() 85793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.setMouseCallback(self.windowname, self.on_mouse) 86793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 87793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler def show(self): 88793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.imshow(self.windowname, self.dests[0]) 89793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 90793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler def on_mouse(self, event, x, y, flags, param): 91793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler pt = (x, y) 92793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if event == cv2.EVENT_LBUTTONDOWN: 93793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.prev_pt = pt 94793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif event == cv2.EVENT_LBUTTONUP: 95793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.prev_pt = None 96793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 97793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if self.prev_pt and flags & cv2.EVENT_FLAG_LBUTTON: 98793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler for dst, color in zip(self.dests, self.colors_func()): 99793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.line(dst, self.prev_pt, pt, color, 5) 100793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.dirty = True 101793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.prev_pt = pt 102793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.show() 103793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 104793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 105793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler# palette data from matplotlib/_cm.py 106793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler_jet_data = {'red': ((0., 0, 0), (0.35, 0, 0), (0.66, 1, 1), (0.89,1, 1), 107793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler (1, 0.5, 0.5)), 108793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 'green': ((0., 0, 0), (0.125,0, 0), (0.375,1, 1), (0.64,1, 1), 109793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler (0.91,0,0), (1, 0, 0)), 110793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 'blue': ((0., 0.5, 0.5), (0.11, 1, 1), (0.34, 1, 1), (0.65,0, 0), 111793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler (1, 0, 0))} 112793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 113793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslercmap_data = { 'jet' : _jet_data } 114793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 115793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef make_cmap(name, n=256): 116793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler data = cmap_data[name] 117793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler xs = np.linspace(0.0, 1.0, n) 118793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler channels = [] 119793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler eps = 1e-6 120793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler for ch_name in ['blue', 'green', 'red']: 121793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler ch_data = data[ch_name] 122793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler xp, yp = [], [] 123793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler for x, y1, y2 in ch_data: 124793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler xp += [x, x+eps] 125793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler yp += [y1, y2] 126793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler ch = np.interp(xs, xp, yp) 127793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler channels.append(ch) 128793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return np.uint8(np.array(channels).T*255) 129793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 130793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef nothing(*arg, **kw): 131793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler pass 132793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 133793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef clock(): 134793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return cv2.getTickCount() / cv2.getTickFrequency() 135793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 136793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler@contextmanager 137793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef Timer(msg): 138793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler print msg, '...', 139793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler start = clock() 140793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler try: 141793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler yield 142793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler finally: 143793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler print "%.2f ms" % ((clock()-start)*1000) 144793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 145793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerclass StatValue: 146793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler def __init__(self, smooth_coef = 0.5): 147793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.value = None 148793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.smooth_coef = smooth_coef 149793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler def update(self, v): 150793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if self.value is None: 151793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.value = v 152793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler else: 153793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler c = self.smooth_coef 154793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.value = c * self.value + (1.0-c) * v 155793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 156793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerclass RectSelector: 157793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler def __init__(self, win, callback): 158793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.win = win 159793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.callback = callback 160793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.setMouseCallback(win, self.onmouse) 161793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.drag_start = None 162793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.drag_rect = None 163793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler def onmouse(self, event, x, y, flags, param): 164793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler x, y = np.int16([x, y]) # BUG 165793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if event == cv2.EVENT_LBUTTONDOWN: 166793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.drag_start = (x, y) 167793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if self.drag_start: 168793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if flags & cv2.EVENT_FLAG_LBUTTON: 169793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler xo, yo = self.drag_start 170793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler x0, y0 = np.minimum([xo, yo], [x, y]) 171793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler x1, y1 = np.maximum([xo, yo], [x, y]) 172793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.drag_rect = None 173793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if x1-x0 > 0 and y1-y0 > 0: 174793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.drag_rect = (x0, y0, x1, y1) 175793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler else: 176793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rect = self.drag_rect 177793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.drag_start = None 178793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.drag_rect = None 179793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if rect: 180793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler self.callback(rect) 181793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler def draw(self, vis): 182793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if not self.drag_rect: 183793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return False 184793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler x0, y0, x1, y1 = self.drag_rect 185793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.rectangle(vis, (x0, y0), (x1, y1), (0, 255, 0), 2) 186793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return True 187793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler @property 188793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler def dragging(self): 189793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return self.drag_rect is not None 190793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 191793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 192793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef grouper(n, iterable, fillvalue=None): 193793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler '''grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx''' 194793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler args = [iter(iterable)] * n 195793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return it.izip_longest(fillvalue=fillvalue, *args) 196793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 197793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef mosaic(w, imgs): 198793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler '''Make a grid from images. 199793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 200793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler w -- number of grid columns 201793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler imgs -- images (must have same size and format) 202793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler ''' 203793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler imgs = iter(imgs) 204793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler img0 = imgs.next() 205793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler pad = np.zeros_like(img0) 206793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler imgs = it.chain([img0], imgs) 207793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rows = grouper(w, imgs, pad) 208793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return np.vstack(map(np.hstack, rows)) 209793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 210793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef getsize(img): 211793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler h, w = img.shape[:2] 212793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return w, h 213793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 214793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef mdot(*args): 215793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler return reduce(np.dot, args) 216793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 217793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef draw_keypoints(vis, keypoints, color = (0, 255, 255)): 218793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler for kp in keypoints: 219793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler x, y = kp.pt 220793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.circle(vis, (int(x), int(y)), 2, color) 221