1#!/usr/bin/env python
2
3import numpy as np
4import cv2
5
6# built-in modules
7import os
8import sys
9
10# local modules
11import video
12from common import mosaic
13
14from digits import *
15
16def 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
93if __name__ == '__main__':
94    main()
95