common.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117
  1. import numpy as np
  2. import cv2
  3. import os
  4. image_extensions = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.tiff', '.pbm', '.pgm', '.ppm']
  5. def splitfn(fn):
  6. path, fn = os.path.split(fn)
  7. name, ext = os.path.splitext(fn)
  8. return path, name, ext
  9. def anorm2(a):
  10. return (a*a).sum(-1)
  11. def anorm(a):
  12. return np.sqrt( anorm2(a) )
  13. def homotrans(H, x, y):
  14. xs = H[0, 0]*x + H[0, 1]*y + H[0, 2]
  15. ys = H[1, 0]*x + H[1, 1]*y + H[1, 2]
  16. s = H[2, 0]*x + H[2, 1]*y + H[2, 2]
  17. return xs/s, ys/s
  18. def to_rect(a):
  19. a = np.ravel(a)
  20. if len(a) == 2:
  21. a = (0, 0, a[0], a[1])
  22. return np.array(a, np.float64).reshape(2, 2)
  23. def rect2rect_mtx(src, dst):
  24. src, dst = to_rect(src), to_rect(dst)
  25. cx, cy = (dst[1] - dst[0]) / (src[1] - src[0])
  26. tx, ty = dst[0] - src[0] * (cx, cy)
  27. M = np.float64([[ cx, 0, tx],
  28. [ 0, cy, ty],
  29. [ 0, 0, 1]])
  30. return M
  31. def lookat(eye, target, up = (0, 0, 1)):
  32. fwd = np.asarray(target, np.float64) - eye
  33. fwd /= anorm(fwd)
  34. right = np.cross(fwd, up)
  35. right /= anorm(right)
  36. down = np.cross(fwd, right)
  37. R = np.float64([right, down, fwd])
  38. tvec = -np.dot(R, eye)
  39. return R, tvec
  40. def mtx2rvec(R):
  41. w, u, vt = cv2.SVDecomp(R - np.eye(3))
  42. p = vt[0] + u[:,0]*w[0] # same as np.dot(R, vt[0])
  43. c = np.dot(vt[0], p)
  44. s = np.dot(vt[1], p)
  45. axis = np.cross(vt[0], vt[1])
  46. return axis * np.arctan2(s, c)
  47. def draw_str(dst, (x, y), s):
  48. cv2.putText(dst, s, (x+1, y+1), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 0), thickness = 2, linetype=cv2.CV_AA)
  49. cv2.putText(dst, s, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), linetype=cv2.CV_AA)
  50. class Sketcher:
  51. def __init__(self, windowname, dests, colors_func):
  52. self.prev_pt = None
  53. self.windowname = windowname
  54. self.dests = dests
  55. self.colors_func = colors_func
  56. self.dirty = False
  57. self.show()
  58. cv2.setMouseCallback(self.windowname, self.on_mouse)
  59. def show(self):
  60. cv2.imshow(self.windowname, self.dests[0])
  61. def on_mouse(self, event, x, y, flags, param):
  62. pt = (x, y)
  63. if event == cv2.EVENT_LBUTTONDOWN:
  64. self.prev_pt = pt
  65. if self.prev_pt and flags & cv2.EVENT_FLAG_LBUTTON:
  66. for dst, color in zip(self.dests, self.colors_func()):
  67. cv2.line(dst, self.prev_pt, pt, color, 5)
  68. self.dirty = True
  69. self.prev_pt = pt
  70. self.show()
  71. else:
  72. self.prev_pt = None
  73. # palette data from matplotlib/_cm.py
  74. _jet_data = {'red': ((0., 0, 0), (0.35, 0, 0), (0.66, 1, 1), (0.89,1, 1),
  75. (1, 0.5, 0.5)),
  76. 'green': ((0., 0, 0), (0.125,0, 0), (0.375,1, 1), (0.64,1, 1),
  77. (0.91,0,0), (1, 0, 0)),
  78. 'blue': ((0., 0.5, 0.5), (0.11, 1, 1), (0.34, 1, 1), (0.65,0, 0),
  79. (1, 0, 0))}
  80. cmap_data = { 'jet' : _jet_data }
  81. def make_cmap(name, n=256):
  82. data = cmap_data[name]
  83. xs = np.linspace(0.0, 1.0, n)
  84. channels = []
  85. eps = 1e-6
  86. for ch_name in ['blue', 'green', 'red']:
  87. ch_data = data[ch_name]
  88. xp, yp = [], []
  89. for x, y1, y2 in ch_data:
  90. xp += [x, x+eps]
  91. yp += [y1, y2]
  92. ch = np.interp(xs, xp, yp)
  93. channels.append(ch)
  94. return np.uint8(np.array(channels).T*255)
  95. def nothing(*arg, **kw):
  96. pass
  97. def clock():
  98. return cv2.getTickCount() / cv2.getTickFrequency()