deconvolution.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. #!/usr/bin/env python
  2. '''
  3. Wiener deconvolution.
  4. Sample shows how DFT can be used to perform Weiner deconvolution [1]
  5. of an image with user-defined point spread function (PSF)
  6. Usage:
  7. deconvolution.py [--circle]
  8. [--angle <degrees>]
  9. [--d <diameter>]
  10. [--snr <signal/noise ratio in db>]
  11. [<input image>]
  12. Use sliders to adjust PSF paramitiers.
  13. Keys:
  14. SPACE - switch btw linear/cirular PSF
  15. ESC - exit
  16. Examples:
  17. deconvolution.py --angle 135 --d 22 ../data/licenseplate_motion.jpg
  18. (image source: http://www.topazlabs.com/infocus/_images/licenseplate_compare.jpg)
  19. deconvolution.py --angle 86 --d 31 ../data/text_motion.jpg
  20. deconvolution.py --circle --d 19 ../data/text_defocus.jpg
  21. (image source: compact digital photo camera, no artificial distortion)
  22. [1] http://en.wikipedia.org/wiki/Wiener_deconvolution
  23. '''
  24. # Python 2/3 compatibility
  25. from __future__ import print_function
  26. import numpy as np
  27. import cv2
  28. # local module
  29. from common import nothing
  30. def blur_edge(img, d=31):
  31. h, w = img.shape[:2]
  32. img_pad = cv2.copyMakeBorder(img, d, d, d, d, cv2.BORDER_WRAP)
  33. img_blur = cv2.GaussianBlur(img_pad, (2*d+1, 2*d+1), -1)[d:-d,d:-d]
  34. y, x = np.indices((h, w))
  35. dist = np.dstack([x, w-x-1, y, h-y-1]).min(-1)
  36. w = np.minimum(np.float32(dist)/d, 1.0)
  37. return img*w + img_blur*(1-w)
  38. def motion_kernel(angle, d, sz=65):
  39. kern = np.ones((1, d), np.float32)
  40. c, s = np.cos(angle), np.sin(angle)
  41. A = np.float32([[c, -s, 0], [s, c, 0]])
  42. sz2 = sz // 2
  43. A[:,2] = (sz2, sz2) - np.dot(A[:,:2], ((d-1)*0.5, 0))
  44. kern = cv2.warpAffine(kern, A, (sz, sz), flags=cv2.INTER_CUBIC)
  45. return kern
  46. def defocus_kernel(d, sz=65):
  47. kern = np.zeros((sz, sz), np.uint8)
  48. cv2.circle(kern, (sz, sz), d, 255, -1, cv2.LINE_AA, shift=1)
  49. kern = np.float32(kern) / 255.0
  50. return kern
  51. if __name__ == '__main__':
  52. print(__doc__)
  53. import sys, getopt
  54. opts, args = getopt.getopt(sys.argv[1:], '', ['circle', 'angle=', 'd=', 'snr='])
  55. opts = dict(opts)
  56. try:
  57. fn = args[0]
  58. except:
  59. fn = '../data/licenseplate_motion.jpg'
  60. win = 'deconvolution'
  61. img = cv2.imread(fn, 0)
  62. if img is None:
  63. print('Failed to load fn1:', fn1)
  64. sys.exit(1)
  65. img = np.float32(img)/255.0
  66. cv2.imshow('input', img)
  67. img = blur_edge(img)
  68. IMG = cv2.dft(img, flags=cv2.DFT_COMPLEX_OUTPUT)
  69. defocus = '--circle' in opts
  70. def update(_):
  71. ang = np.deg2rad( cv2.getTrackbarPos('angle', win) )
  72. d = cv2.getTrackbarPos('d', win)
  73. noise = 10**(-0.1*cv2.getTrackbarPos('SNR (db)', win))
  74. if defocus:
  75. psf = defocus_kernel(d)
  76. else:
  77. psf = motion_kernel(ang, d)
  78. cv2.imshow('psf', psf)
  79. psf /= psf.sum()
  80. psf_pad = np.zeros_like(img)
  81. kh, kw = psf.shape
  82. psf_pad[:kh, :kw] = psf
  83. PSF = cv2.dft(psf_pad, flags=cv2.DFT_COMPLEX_OUTPUT, nonzeroRows = kh)
  84. PSF2 = (PSF**2).sum(-1)
  85. iPSF = PSF / (PSF2 + noise)[...,np.newaxis]
  86. RES = cv2.mulSpectrums(IMG, iPSF, 0)
  87. res = cv2.idft(RES, flags=cv2.DFT_SCALE | cv2.DFT_REAL_OUTPUT )
  88. res = np.roll(res, -kh//2, 0)
  89. res = np.roll(res, -kw//2, 1)
  90. cv2.imshow(win, res)
  91. cv2.namedWindow(win)
  92. cv2.namedWindow('psf', 0)
  93. cv2.createTrackbar('angle', win, int(opts.get('--angle', 135)), 180, update)
  94. cv2.createTrackbar('d', win, int(opts.get('--d', 22)), 50, update)
  95. cv2.createTrackbar('SNR (db)', win, int(opts.get('--snr', 25)), 50, update)
  96. update(None)
  97. while True:
  98. ch = cv2.waitKey() & 0xFF
  99. if ch == 27:
  100. break
  101. if ch == ord(' '):
  102. defocus = not defocus
  103. update(None)