From 65e235060623b0e58ad96b126c7bcf256874a159 Mon Sep 17 00:00:00 2001
From: Alexander Mordvintsev <no@email>
Date: Thu, 12 Jul 2012 11:51:27 +0000
Subject: [PATCH] work on digits_video.py

---
 samples/python2/digits_video.py | 65 ++++++++++++++++++++-------------
 samples/python2/morphology.py   |  1 +
 2 files changed, 40 insertions(+), 26 deletions(-)

diff --git a/samples/python2/digits_video.py b/samples/python2/digits_video.py
index 2a46daad3..999ec9b24 100644
--- a/samples/python2/digits_video.py
+++ b/samples/python2/digits_video.py
@@ -1,70 +1,83 @@
 import numpy as np
 import cv2
 import os
+import sys
 import video
 from common import mosaic
 
 from digits import *
 
-
 def main():
-    cap = video.create_capture()
+    try: src = sys.argv[1]
+    except: src = 0
+    cap = video.create_capture(src)
 
     classifier_fn = 'digits_svm.dat'
     if not os.path.exists(classifier_fn):
         print '"%s" not found, run digits.py first' % classifier_fn
         return 
-    
     model = SVM()
-    model.load('digits_svm.dat')
+    model.load(classifier_fn)
 
+    
     while True:
         ret, frame = cap.read()
         gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
 
+
         bin = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 31, 10)
         bin = cv2.medianBlur(bin, 3)
-        contours, heirs = cv2.findContours( bin.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
-        rects = map(cv2.boundingRect, contours)
-        valid_flags = [ 16 <= h <= 64  and w <= 1.2*h  for x, y, w, h in rects]
+        contours, heirs = cv2.findContours( bin.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
+        try: heirs = heirs[0]
+        except: heirs = []
 
-        for i, cnt in enumerate(contours):
-            if not valid_flags[i]:
+        for cnt, heir in zip(contours, heirs):
+            _, _, _, outer_i = heir
+            if outer_i >= 0:
                 continue
-            _, _, _, outer_i = heirs[0, i]
-            if outer_i >=0 and valid_flags[outer_i]:
+            x, y, w, h = cv2.boundingRect(cnt)
+            if not (16 <= h <= 64  and w <= 1.2*h):
                 continue
-            x, y, w, h = rects[i]
+            pad = max(h-w, 0)
+            x, w = x-pad/2, w+pad
             cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0))
-            sub = bin[y:,x:][:h,:w]
-            #sub = ~cv2.equalizeHist(sub)
-            #_, sub_bin = cv2.threshold(sub, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
 
-            s = 1.5*float(h)/SZ
-            m = cv2.moments(sub)
-            m00 = m['m00']
-            if m00/255 < 0.1*w*h or m00/255 > 0.9*w*h:
+            bin_roi = bin[y:,x:][:h,:w]
+            gray_roi = gray[y:,x:][:h,:w]
+
+            m = bin_roi != 0
+            if not 0.1 < m.mean() < 0.4:
                 continue
-
-            c1 = np.float32([m['m10'], m['m01']]) / m00
+            '''
+            v_in, v_out = gray_roi[m], gray_roi[~m]
+            if v_out.std() > 10.0:
+                continue
+            s = "%f, %f" % (abs(v_in.mean() - v_out.mean()), v_out.std())
+            cv2.putText(frame, s, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (200, 0, 0), thickness = 1)
+            '''
+            
+            s = 1.5*float(h)/SZ
+            m = cv2.moments(bin_roi)
+            c1 = np.float32([m['m10'], m['m01']]) / m['m00']
             c0 = np.float32([SZ/2, SZ/2])
             t = c1 - s*c0
             A = np.zeros((2, 3), np.float32)
             A[:,:2] = np.eye(2)*s
             A[:,2] = t
-            sub1 = cv2.warpAffine(sub, A, (SZ, SZ), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LINEAR)
-            sub1 = deskew(sub1)
+            bin_norm = cv2.warpAffine(bin_roi, A, (SZ, SZ), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LINEAR)
+            bin_norm = deskew(bin_norm)
             if x+w+SZ < frame.shape[1] and y+SZ < frame.shape[0]:
-                frame[y:,x+w:][:SZ, :SZ] = sub1[...,np.newaxis]
+                frame[y:,x+w:][:SZ, :SZ] = bin_norm[...,np.newaxis]
                 
-            sample = preprocess_hog([sub1])
+            sample = preprocess_hog([bin_norm])
             digit = model.predict(sample)[0]
             cv2.putText(frame, '%d'%digit, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (200, 0, 0), thickness = 1)
 
 
         cv2.imshow('frame', frame)
         cv2.imshow('bin', bin)
-        if cv2.waitKey(1) == 27:
+        ch = cv2.waitKey(1)
+        if ch == 27:
             break
 
 if __name__ == '__main__':
diff --git a/samples/python2/morphology.py b/samples/python2/morphology.py
index 68b95d862..b7f84fbb2 100644
--- a/samples/python2/morphology.py
+++ b/samples/python2/morphology.py
@@ -10,6 +10,7 @@ if __name__ == '__main__':
     try: fn = sys.argv[1]
     except: fn = '../cpp/baboon.jpg'
     img = cv2.imread(fn)
+    cv2.imshow('original', img)
 
     modes = cycle(['erode/dilate', 'open/close', 'blackhat/tophat', 'gradient'])
     str_modes = cycle(['ellipse', 'rect', 'cross'])