162 lines
5.1 KiB
Python
162 lines
5.1 KiB
Python
import numpy as np
|
|
import cv2
|
|
import os
|
|
from contextlib import contextmanager
|
|
|
|
image_extensions = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.tiff', '.pbm', '.pgm', '.ppm']
|
|
|
|
def splitfn(fn):
|
|
path, fn = os.path.split(fn)
|
|
name, ext = os.path.splitext(fn)
|
|
return path, name, ext
|
|
|
|
def anorm2(a):
|
|
return (a*a).sum(-1)
|
|
def anorm(a):
|
|
return np.sqrt( anorm2(a) )
|
|
|
|
def homotrans(H, x, y):
|
|
xs = H[0, 0]*x + H[0, 1]*y + H[0, 2]
|
|
ys = H[1, 0]*x + H[1, 1]*y + H[1, 2]
|
|
s = H[2, 0]*x + H[2, 1]*y + H[2, 2]
|
|
return xs/s, ys/s
|
|
|
|
def to_rect(a):
|
|
a = np.ravel(a)
|
|
if len(a) == 2:
|
|
a = (0, 0, a[0], a[1])
|
|
return np.array(a, np.float64).reshape(2, 2)
|
|
|
|
def rect2rect_mtx(src, dst):
|
|
src, dst = to_rect(src), to_rect(dst)
|
|
cx, cy = (dst[1] - dst[0]) / (src[1] - src[0])
|
|
tx, ty = dst[0] - src[0] * (cx, cy)
|
|
M = np.float64([[ cx, 0, tx],
|
|
[ 0, cy, ty],
|
|
[ 0, 0, 1]])
|
|
return M
|
|
|
|
|
|
def lookat(eye, target, up = (0, 0, 1)):
|
|
fwd = np.asarray(target, np.float64) - eye
|
|
fwd /= anorm(fwd)
|
|
right = np.cross(fwd, up)
|
|
right /= anorm(right)
|
|
down = np.cross(fwd, right)
|
|
R = np.float64([right, down, fwd])
|
|
tvec = -np.dot(R, eye)
|
|
return R, tvec
|
|
|
|
def mtx2rvec(R):
|
|
w, u, vt = cv2.SVDecomp(R - np.eye(3))
|
|
p = vt[0] + u[:,0]*w[0] # same as np.dot(R, vt[0])
|
|
c = np.dot(vt[0], p)
|
|
s = np.dot(vt[1], p)
|
|
axis = np.cross(vt[0], vt[1])
|
|
return axis * np.arctan2(s, c)
|
|
|
|
def draw_str(dst, (x, y), s):
|
|
cv2.putText(dst, s, (x+1, y+1), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 0), thickness = 2, linetype=cv2.CV_AA)
|
|
cv2.putText(dst, s, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), linetype=cv2.CV_AA)
|
|
|
|
class Sketcher:
|
|
def __init__(self, windowname, dests, colors_func):
|
|
self.prev_pt = None
|
|
self.windowname = windowname
|
|
self.dests = dests
|
|
self.colors_func = colors_func
|
|
self.dirty = False
|
|
self.show()
|
|
cv2.setMouseCallback(self.windowname, self.on_mouse)
|
|
|
|
def show(self):
|
|
cv2.imshow(self.windowname, self.dests[0])
|
|
|
|
def on_mouse(self, event, x, y, flags, param):
|
|
pt = (x, y)
|
|
if event == cv2.EVENT_LBUTTONDOWN:
|
|
self.prev_pt = pt
|
|
if self.prev_pt and flags & cv2.EVENT_FLAG_LBUTTON:
|
|
for dst, color in zip(self.dests, self.colors_func()):
|
|
cv2.line(dst, self.prev_pt, pt, color, 5)
|
|
self.dirty = True
|
|
self.prev_pt = pt
|
|
self.show()
|
|
else:
|
|
self.prev_pt = None
|
|
|
|
|
|
# palette data from matplotlib/_cm.py
|
|
_jet_data = {'red': ((0., 0, 0), (0.35, 0, 0), (0.66, 1, 1), (0.89,1, 1),
|
|
(1, 0.5, 0.5)),
|
|
'green': ((0., 0, 0), (0.125,0, 0), (0.375,1, 1), (0.64,1, 1),
|
|
(0.91,0,0), (1, 0, 0)),
|
|
'blue': ((0., 0.5, 0.5), (0.11, 1, 1), (0.34, 1, 1), (0.65,0, 0),
|
|
(1, 0, 0))}
|
|
|
|
cmap_data = { 'jet' : _jet_data }
|
|
|
|
def make_cmap(name, n=256):
|
|
data = cmap_data[name]
|
|
xs = np.linspace(0.0, 1.0, n)
|
|
channels = []
|
|
eps = 1e-6
|
|
for ch_name in ['blue', 'green', 'red']:
|
|
ch_data = data[ch_name]
|
|
xp, yp = [], []
|
|
for x, y1, y2 in ch_data:
|
|
xp += [x, x+eps]
|
|
yp += [y1, y2]
|
|
ch = np.interp(xs, xp, yp)
|
|
channels.append(ch)
|
|
return np.uint8(np.array(channels).T*255)
|
|
|
|
def nothing(*arg, **kw):
|
|
pass
|
|
|
|
def clock():
|
|
return cv2.getTickCount() / cv2.getTickFrequency()
|
|
|
|
@contextmanager
|
|
def Timer(msg):
|
|
print msg, '...',
|
|
start = clock()
|
|
try:
|
|
yield
|
|
finally:
|
|
print "%.2f ms" % ((clock()-start)*1000)
|
|
|
|
class RectSelector:
|
|
def __init__(self, win, callback):
|
|
self.win = win
|
|
self.callback = callback
|
|
cv2.setMouseCallback(win, self.onmouse)
|
|
self.drag_start = None
|
|
self.drag_rect = None
|
|
def onmouse(self, event, x, y, flags, param):
|
|
x, y = np.int16([x, y]) # BUG
|
|
if event == cv2.EVENT_LBUTTONDOWN:
|
|
self.drag_start = (x, y)
|
|
if self.drag_start:
|
|
if flags & cv2.EVENT_FLAG_LBUTTON:
|
|
#h, w = self.frame.shape[:2]
|
|
xo, yo = self.drag_start
|
|
x0, y0 = np.minimum([xo, yo], [x, y])
|
|
x1, y1 = np.maximum([xo, yo], [x, y])
|
|
#x0, y0 = np.maximum(0, np.minimum([xo, yo], [x, y]))
|
|
#x1, y1 = np.minimum([w, h], np.maximum([xo, yo], [x, y]))
|
|
self.drag_rect = None
|
|
if x1-x0 > 0 and y1-y0 > 0:
|
|
self.drag_rect = (x0, y0, x1, y1)
|
|
else:
|
|
rect = self.drag_rect
|
|
self.drag_start = None
|
|
self.drag_rect = None
|
|
if rect:
|
|
self.callback(rect)
|
|
def draw(self, vis):
|
|
if not self.drag_rect:
|
|
return
|
|
x0, y0, x1, y1 = self.drag_rect
|
|
cv2.rectangle(vis, (x0, y0), (x1, y1), (0, 255, 0), 2)
|