digits_video.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109
  1. #!/usr/bin/env python
  2. '''
  3. Digit recognition from video.
  4. Run digits.py before, to train and save the SVM.
  5. Usage:
  6. digits_video.py [{camera_id|video_file}]
  7. '''
  8. # Python 2/3 compatibility
  9. from __future__ import print_function
  10. import numpy as np
  11. import cv2 as cv
  12. # built-in modules
  13. import os
  14. import sys
  15. # local modules
  16. import video
  17. from common import mosaic
  18. from digits import *
  19. def main():
  20. try:
  21. src = sys.argv[1]
  22. except:
  23. src = 0
  24. cap = video.create_capture(src, fallback='synth:bg={}:noise=0.05'.format(cv.samples.findFile('sudoku.png')))
  25. classifier_fn = 'digits_svm.dat'
  26. if not os.path.exists(classifier_fn):
  27. print('"%s" not found, run digits.py first' % classifier_fn)
  28. return
  29. model = cv.ml.SVM_load(classifier_fn)
  30. while True:
  31. _ret, frame = cap.read()
  32. gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
  33. bin = cv.adaptiveThreshold(gray, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY_INV, 31, 10)
  34. bin = cv.medianBlur(bin, 3)
  35. contours, heirs = cv.findContours( bin.copy(), cv.RETR_CCOMP, cv.CHAIN_APPROX_SIMPLE)
  36. try:
  37. heirs = heirs[0]
  38. except:
  39. heirs = []
  40. for cnt, heir in zip(contours, heirs):
  41. _, _, _, outer_i = heir
  42. if outer_i >= 0:
  43. continue
  44. x, y, w, h = cv.boundingRect(cnt)
  45. if not (16 <= h <= 64 and w <= 1.2*h):
  46. continue
  47. pad = max(h-w, 0)
  48. x, w = x - (pad // 2), w + pad
  49. cv.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0))
  50. bin_roi = bin[y:,x:][:h,:w]
  51. m = bin_roi != 0
  52. if not 0.1 < m.mean() < 0.4:
  53. continue
  54. '''
  55. gray_roi = gray[y:,x:][:h,:w]
  56. v_in, v_out = gray_roi[m], gray_roi[~m]
  57. if v_out.std() > 10.0:
  58. continue
  59. s = "%f, %f" % (abs(v_in.mean() - v_out.mean()), v_out.std())
  60. cv.putText(frame, s, (x, y), cv.FONT_HERSHEY_PLAIN, 1.0, (200, 0, 0), thickness = 1)
  61. '''
  62. s = 1.5*float(h)/SZ
  63. m = cv.moments(bin_roi)
  64. c1 = np.float32([m['m10'], m['m01']]) / m['m00']
  65. c0 = np.float32([SZ/2, SZ/2])
  66. t = c1 - s*c0
  67. A = np.zeros((2, 3), np.float32)
  68. A[:,:2] = np.eye(2)*s
  69. A[:,2] = t
  70. bin_norm = cv.warpAffine(bin_roi, A, (SZ, SZ), flags=cv.WARP_INVERSE_MAP | cv.INTER_LINEAR)
  71. bin_norm = deskew(bin_norm)
  72. if x+w+SZ < frame.shape[1] and y+SZ < frame.shape[0]:
  73. frame[y:,x+w:][:SZ, :SZ] = bin_norm[...,np.newaxis]
  74. sample = preprocess_hog([bin_norm])
  75. digit = model.predict(sample)[1].ravel()
  76. cv.putText(frame, '%d'%digit, (x, y), cv.FONT_HERSHEY_PLAIN, 1.0, (200, 0, 0), thickness = 1)
  77. cv.imshow('frame', frame)
  78. cv.imshow('bin', bin)
  79. ch = cv.waitKey(1)
  80. if ch == 27:
  81. break
  82. print('Done')
  83. if __name__ == '__main__':
  84. print(__doc__)
  85. main()
  86. cv.destroyAllWindows()