Commit b987154e authored by Alexander Mordvintsev's avatar Alexander Mordvintsev

digits_video.py prints warning if trained classifier (should be created by digits.py) not found

parent 3804ca3e
import numpy as np import numpy as np
import cv2 import cv2
#import video
import digits import digits
import os
import video
from common import mosaic from common import mosaic
#cap = video.create_capture()
cap = cv2.VideoCapture(0)
def main():
model = digits.SVM() cap = video.create_capture()
model.load('digits_svm.dat')
classifier_fn = 'digits_svm.dat'
SZ = 20 if not os.path.exists(classifier_fn):
print '"%s" not found, run digits.py first' % classifier_fn
while True: return
ret, frame = cap.read()
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) model = digits.SVM()
model.load('digits_svm.dat')
bin = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 31, 10)
bin = cv2.medianBlur(bin, 3) SZ = 20
contours, _ = cv2.findContours( bin.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
while True:
boxes = [] ret, frame = cap.read()
for cnt in contours: gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
x, y, w, h = cv2.boundingRect(cnt)
if h < 20 or h > 60 or 1.2*h < w: bin = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 31, 10)
continue bin = cv2.medianBlur(bin, 3)
cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0)) contours, _ = cv2.findContours( bin.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
sub = bin[y:,x:][:h,:w]
#sub = ~cv2.equalizeHist(sub) boxes = []
#_, sub_bin = cv2.threshold(sub, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU) for cnt in contours:
x, y, w, h = cv2.boundingRect(cnt)
s = 1.1*h/SZ if h < 20 or h > 60 or 1.2*h < w:
m = cv2.moments(sub) continue
m00 = m['m00'] cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0))
if m00/255 < 0.1*w*h or m00/255 > 0.9*w*h: sub = bin[y:,x:][:h,:w]
continue #sub = ~cv2.equalizeHist(sub)
#_, sub_bin = cv2.threshold(sub, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
#frame[y:,x:][:h,:w] = sub[...,np.newaxis]
c1 = np.float32([m['m10'], m['m01']]) / m00 s = 1.1*h/SZ
c0 = np.float32([SZ/2, SZ/2]) m = cv2.moments(sub)
t = c1 - s*c0 m00 = m['m00']
A = np.zeros((2, 3), np.float32) if m00/255 < 0.1*w*h or m00/255 > 0.9*w*h:
A[:,:2] = np.eye(2)*2 continue
A[:,2] = t
sub1 = cv2.warpAffine(sub, A, (SZ, SZ), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LINEAR) #frame[y:,x:][:h,:w] = sub[...,np.newaxis]
sub1 = digits.deskew(sub1) c1 = np.float32([m['m10'], m['m01']]) / m00
sample = np.float32(sub1).reshape(1,SZ*SZ) / 255.0 c0 = np.float32([SZ/2, SZ/2])
digit = model.predict(sample)[0] t = c1 - s*c0
A = np.zeros((2, 3), np.float32)
cv2.putText(frame, '%d'%digit, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (200, 0, 0), thickness = 1) A[:,:2] = np.eye(2)*2
A[:,2] = t
boxes.append(sub1) sub1 = cv2.warpAffine(sub, A, (SZ, SZ), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LINEAR)
sub1 = digits.deskew(sub1)
sample = np.float32(sub1).reshape(1,SZ*SZ) / 255.0
if len(boxes) > 0: digit = model.predict(sample)[0]
cv2.imshow('box', mosaic(10, boxes))
cv2.putText(frame, '%d'%digit, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (200, 0, 0), thickness = 1)
cv2.imshow('frame', frame) boxes.append(sub1)
cv2.imshow('bin', bin)
if cv2.waitKey(1) == 27:
break if len(boxes) > 0:
cv2.imshow('box', mosaic(10, boxes))
cv2.imshow('frame', frame)
cv2.imshow('bin', bin)
if cv2.waitKey(1) == 27:
break
if __name__ == '__main__':
main()
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment