opencv/samples/python2/digits2.py

162 lines
5.4 KiB
Python
Raw Normal View History

import numpy as np
import cv2
from multiprocessing.pool import ThreadPool
SZ = 20 # size of each digit is SZ x SZ
CLASS_N = 10
def load_base(fn):
print 'loading "%s" ...' % fn
digits_img = cv2.imread(fn, 0)
h, w = digits_img.shape
digits = [np.hsplit(row, w/SZ) for row in np.vsplit(digits_img, h/SZ)]
digits = np.array(digits).reshape(-1, SZ, SZ)
digits = np.float32(digits).reshape(-1, SZ*SZ) / 255.0
labels = np.repeat(np.arange(CLASS_N), len(digits)/CLASS_N)
return digits, labels
def cross_validate(model_class, params, samples, labels, kfold = 4, pool = None):
n = len(samples)
folds = np.array_split(np.arange(n), kfold)
def f(i):
model = model_class(**params)
test_idx = folds[i]
train_idx = list(folds)
train_idx.pop(i)
train_idx = np.hstack(train_idx)
train_samples, train_labels = samples[train_idx], labels[train_idx]
test_samples, test_labels = samples[test_idx], labels[test_idx]
model.train(train_samples, train_labels)
resp = model.predict(test_samples)
score = (resp != test_labels).mean()
print ".",
return score
if pool is None:
scores = map(f, xrange(kfold))
else:
scores = pool.map(f, xrange(kfold))
return np.mean(scores)
class StatModel(object):
def load(self, fn):
self.model.load(fn)
def save(self, fn):
self.model.save(fn)
class KNearest(StatModel):
def __init__(self, k = 3):
self.k = k
@staticmethod
def adjust(samples, labels):
print 'adjusting KNearest ...'
best_err, best_k = np.inf, -1
for k in xrange(1, 11):
err = cross_validate(KNearest, dict(k=k), samples, labels)
if err < best_err:
best_err, best_k = err, k
print 'k = %d, error: %.2f %%' % (k, err*100)
best_params = dict(k=best_k)
print 'best params:', best_params
return best_params
def train(self, samples, responses):
self.model = cv2.KNearest()
self.model.train(samples, responses)
def predict(self, samples):
retval, results, neigh_resp, dists = self.model.find_nearest(samples, self.k)
return results.ravel()
class SVM(StatModel):
def __init__(self, C = 1, gamma = 0.5):
self.params = dict( kernel_type = cv2.SVM_RBF,
svm_type = cv2.SVM_C_SVC,
C = C,
gamma = gamma )
@staticmethod
def adjust(samples, labels):
Cs = np.logspace(0, 5, 10, base=2)
gammas = np.logspace(-7, -2, 10, base=2)
scores = np.zeros((len(Cs), len(gammas)))
scores[:] = np.nan
print 'adjusting SVM (may take a long time) ...'
def f(job):
i, j = job
params = dict(C = Cs[i], gamma=gammas[j])
score = cross_validate(SVM, params, samples, labels)
scores[i, j] = score
nready = np.isfinite(scores).sum()
print '%d / %d (best error: %.2f %%, last: %.2f %%)' % (nready, scores.size, np.nanmin(scores)*100, score*100)
pool = ThreadPool(processes=cv2.getNumberOfCPUs())
pool.map(f, np.ndindex(*scores.shape))
print scores
i, j = np.unravel_index(scores.argmin(), scores.shape)
best_params = dict(C = Cs[i], gamma=gammas[j])
print 'best params:', best_params
print 'best error: %.2f %%' % (scores.min()*100)
return best_params
def train(self, samples, responses):
self.model = cv2.SVM()
self.model.train(samples, responses, params = self.params)
def predict(self, samples):
return self.model.predict_all(samples).ravel()
def main_adjustSVM(samples, labels):
params = SVM.adjust(samples, labels)
print 'training SVM on all samples ...'
model = SVN(**params)
model.train(samples, labels)
print 'saving "digits_svm.dat" ...'
model.save('digits_svm.dat')
def main_adjustKNearest(samples, labels):
params = KNearest.adjust(samples, labels)
def main_showSVM(samples, labels):
from common import mosaic
train_n = int(0.9*len(samples))
digits_train, digits_test = np.split(samples[shuffle], [train_n])
labels_train, labels_test = np.split(labels[shuffle], [train_n])
print 'training SVM ...'
model = SVM(C=2.16, gamma=0.0536)
model.train(digits_train, labels_train)
train_err = (model.predict(digits_train) != labels_train).mean()
resp_test = model.predict(digits_test)
test_err = (resp_test != labels_test).mean()
print 'train errors: %.2f %%' % (train_err*100)
print 'test errors: %.2f %%' % (test_err*100)
# visualize test results
vis = []
for img, flag in zip(digits_test, resp_test == labels_test):
img = np.uint8(img*255).reshape(SZ, SZ)
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
if not flag:
img[...,:2] = 0
vis.append(img)
vis = mosaic(25, vis)
cv2.imshow('test', vis)
cv2.waitKey()
if __name__ == '__main__':
samples, labels = load_base('digits.png')
shuffle = np.random.permutation(len(samples))
samples, labels = samples[shuffle], labels[shuffle]
#main_adjustSVM(samples, labels)
#main_adjustKNearest(samples, labels)
main_showSVM(samples, labels)