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