samples/python2/digits_video.py

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