work on digits_video.py

This commit is contained in:
Alexander Mordvintsev
2012-07-12 11:51:27 +00:00
parent 9f0a16c844
commit 65e2350606
2 changed files with 40 additions and 26 deletions

View File

@@ -1,70 +1,83 @@
import numpy as np import numpy as np
import cv2 import cv2
import os import os
import sys
import video import video
from common import mosaic from common import mosaic
from digits import * from digits import *
def main(): def main():
cap = video.create_capture() try: src = sys.argv[1]
except: src = 0
cap = video.create_capture(src)
classifier_fn = 'digits_svm.dat' classifier_fn = 'digits_svm.dat'
if not os.path.exists(classifier_fn): if not os.path.exists(classifier_fn):
print '"%s" not found, run digits.py first' % classifier_fn print '"%s" not found, run digits.py first' % classifier_fn
return return
model = SVM() model = SVM()
model.load('digits_svm.dat') model.load(classifier_fn)
while True: while True:
ret, frame = cap.read() ret, frame = cap.read()
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) 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.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 31, 10)
bin = cv2.medianBlur(bin, 3) bin = cv2.medianBlur(bin, 3)
contours, heirs = cv2.findContours( bin.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) contours, heirs = cv2.findContours( bin.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
rects = map(cv2.boundingRect, contours) try: heirs = heirs[0]
valid_flags = [ 16 <= h <= 64 and w <= 1.2*h for x, y, w, h in rects] except: heirs = []
for i, cnt in enumerate(contours): for cnt, heir in zip(contours, heirs):
if not valid_flags[i]: _, _, _, outer_i = heir
if outer_i >= 0:
continue continue
_, _, _, outer_i = heirs[0, i] x, y, w, h = cv2.boundingRect(cnt)
if outer_i >=0 and valid_flags[outer_i]: if not (16 <= h <= 64 and w <= 1.2*h):
continue 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)) 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 bin_roi = bin[y:,x:][:h,:w]
m = cv2.moments(sub) gray_roi = gray[y:,x:][:h,:w]
m00 = m['m00']
if m00/255 < 0.1*w*h or m00/255 > 0.9*w*h: m = bin_roi != 0
if not 0.1 < m.mean() < 0.4:
continue 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]) c0 = np.float32([SZ/2, SZ/2])
t = c1 - s*c0 t = c1 - s*c0
A = np.zeros((2, 3), np.float32) A = np.zeros((2, 3), np.float32)
A[:,:2] = np.eye(2)*s A[:,:2] = np.eye(2)*s
A[:,2] = t A[:,2] = t
sub1 = cv2.warpAffine(sub, A, (SZ, SZ), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LINEAR) bin_norm = cv2.warpAffine(bin_roi, A, (SZ, SZ), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LINEAR)
sub1 = deskew(sub1) bin_norm = deskew(bin_norm)
if x+w+SZ < frame.shape[1] and y+SZ < frame.shape[0]: 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] digit = model.predict(sample)[0]
cv2.putText(frame, '%d'%digit, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (200, 0, 0), thickness = 1) cv2.putText(frame, '%d'%digit, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (200, 0, 0), thickness = 1)
cv2.imshow('frame', frame) cv2.imshow('frame', frame)
cv2.imshow('bin', bin) cv2.imshow('bin', bin)
if cv2.waitKey(1) == 27: ch = cv2.waitKey(1)
if ch == 27:
break break
if __name__ == '__main__': if __name__ == '__main__':

View File

@@ -10,6 +10,7 @@ if __name__ == '__main__':
try: fn = sys.argv[1] try: fn = sys.argv[1]
except: fn = '../cpp/baboon.jpg' except: fn = '../cpp/baboon.jpg'
img = cv2.imread(fn) img = cv2.imread(fn)
cv2.imshow('original', img)
modes = cycle(['erode/dilate', 'open/close', 'blackhat/tophat', 'gradient']) modes = cycle(['erode/dilate', 'open/close', 'blackhat/tophat', 'gradient'])
str_modes = cycle(['ellipse', 'rect', 'cross']) str_modes = cycle(['ellipse', 'rect', 'cross'])