CoCalc -- test_digits.py

GitHub Repository: Tetragramm/opencv
Path: blob/master/modules/python/test/test_digits.py
¹⁶³³⁷ views
1
#!/usr/bin/env python
2

3
'''
4
SVM and KNearest digit recognition.
5

6
Sample loads a dataset of handwritten digits from '../data/digits.png'.
7
Then it trains a SVM and KNearest classifiers on it and evaluates
8
their accuracy.
9

10
Following preprocessing is applied to the dataset:
11
 - Moment-based image deskew (see deskew())
12
 - Digit images are split into 4 10x10 cells and 16-bin
13
   histogram of oriented gradients is computed for each
14
   cell
15
 - Transform histograms to space with Hellinger metric (see [1] (RootSIFT))
16

17

18
[1] R. Arandjelovic, A. Zisserman
19
    "Three things everyone should know to improve object retrieval"
20
    http://www.robots.ox.ac.uk/~vgg/publications/2012/Arandjelovic12/arandjelovic12.pdf
21

22
'''
23

24

25
# Python 2/3 compatibility
26
from __future__ import print_function
27

28
# built-in modules
29
from multiprocessing.pool import ThreadPool
30

31
import cv2 as cv
32

33
import numpy as np
34
from numpy.linalg import norm
35

36

37
SZ = 20 # size of each digit is SZ x SZ
38
CLASS_N = 10
39
DIGITS_FN = 'samples/data/digits.png'
40

41
def split2d(img, cell_size, flatten=True):
42
    h, w = img.shape[:2]
43
    sx, sy = cell_size
44
    cells = [np.hsplit(row, w//sx) for row in np.vsplit(img, h//sy)]
45
    cells = np.array(cells)
46
    if flatten:
47
        cells = cells.reshape(-1, sy, sx)
48
    return cells
49

50
def deskew(img):
51
    m = cv.moments(img)
52
    if abs(m['mu02']) < 1e-2:
53
        return img.copy()
54
    skew = m['mu11']/m['mu02']
55
    M = np.float32([[1, skew, -0.5*SZ*skew], [0, 1, 0]])
56
    img = cv.warpAffine(img, M, (SZ, SZ), flags=cv.WARP_INVERSE_MAP | cv.INTER_LINEAR)
57
    return img
58

59
class StatModel(object):
60
    def load(self, fn):
61
        self.model.load(fn)  # Known bug: https://github.com/opencv/opencv/issues/4969
62
    def save(self, fn):
63
        self.model.save(fn)
64

65
class KNearest(StatModel):
66
    def __init__(self, k = 3):
67
        self.k = k
68
        self.model = cv.ml.KNearest_create()
69

70
    def train(self, samples, responses):
71
        self.model.train(samples, cv.ml.ROW_SAMPLE, responses)
72

73
    def predict(self, samples):
74
        _retval, results, _neigh_resp, _dists = self.model.findNearest(samples, self.k)
75
        return results.ravel()
76

77
class SVM(StatModel):
78
    def __init__(self, C = 1, gamma = 0.5):
79
        self.model = cv.ml.SVM_create()
80
        self.model.setGamma(gamma)
81
        self.model.setC(C)
82
        self.model.setKernel(cv.ml.SVM_RBF)
83
        self.model.setType(cv.ml.SVM_C_SVC)
84

85
    def train(self, samples, responses):
86
        self.model.train(samples, cv.ml.ROW_SAMPLE, responses)
87

88
    def predict(self, samples):
89
        return self.model.predict(samples)[1].ravel()
90

91

92
def evaluate_model(model, digits, samples, labels):
93
    resp = model.predict(samples)
94
    err = (labels != resp).mean()
95

96
    confusion = np.zeros((10, 10), np.int32)
97
    for i, j in zip(labels, resp):
98
        confusion[int(i), int(j)] += 1
99

100
    return err, confusion
101

102
def preprocess_simple(digits):
103
    return np.float32(digits).reshape(-1, SZ*SZ) / 255.0
104

105
def preprocess_hog(digits):
106
    samples = []
107
    for img in digits:
108
        gx = cv.Sobel(img, cv.CV_32F, 1, 0)
109
        gy = cv.Sobel(img, cv.CV_32F, 0, 1)
110
        mag, ang = cv.cartToPolar(gx, gy)
111
        bin_n = 16
112
        bin = np.int32(bin_n*ang/(2*np.pi))
113
        bin_cells = bin[:10,:10], bin[10:,:10], bin[:10,10:], bin[10:,10:]
114
        mag_cells = mag[:10,:10], mag[10:,:10], mag[:10,10:], mag[10:,10:]
115
        hists = [np.bincount(b.ravel(), m.ravel(), bin_n) for b, m in zip(bin_cells, mag_cells)]
116
        hist = np.hstack(hists)
117

118
        # transform to Hellinger kernel
119
        eps = 1e-7
120
        hist /= hist.sum() + eps
121
        hist = np.sqrt(hist)
122
        hist /= norm(hist) + eps
123

124
        samples.append(hist)
125
    return np.float32(samples)
126

127
from tests_common import NewOpenCVTests
128

129
class digits_test(NewOpenCVTests):
130

131
    def load_digits(self, fn):
132
        digits_img = self.get_sample(fn, 0)
133
        digits = split2d(digits_img, (SZ, SZ))
134
        labels = np.repeat(np.arange(CLASS_N), len(digits)/CLASS_N)
135
        return digits, labels
136

137
    def test_digits(self):
138

139
        digits, labels = self.load_digits(DIGITS_FN)
140

141
        # shuffle digits
142
        rand = np.random.RandomState(321)
143
        shuffle = rand.permutation(len(digits))
144
        digits, labels = digits[shuffle], labels[shuffle]
145

146
        digits2 = list(map(deskew, digits))
147
        samples = preprocess_hog(digits2)
148

149
        train_n = int(0.9*len(samples))
150
        _digits_train, digits_test = np.split(digits2, [train_n])
151
        samples_train, samples_test = np.split(samples, [train_n])
152
        labels_train, labels_test = np.split(labels, [train_n])
153
        errors = list()
154
        confusionMatrixes = list()
155

156
        model = KNearest(k=4)
157
        model.train(samples_train, labels_train)
158
        error, confusion = evaluate_model(model, digits_test, samples_test, labels_test)
159
        errors.append(error)
160
        confusionMatrixes.append(confusion)
161

162
        model = SVM(C=2.67, gamma=5.383)
163
        model.train(samples_train, labels_train)
164
        error, confusion = evaluate_model(model, digits_test, samples_test, labels_test)
165
        errors.append(error)
166
        confusionMatrixes.append(confusion)
167

168
        eps = 0.001
169
        normEps = len(samples_test) * 0.02
170

171
        confusionKNN = [[45,  0,  0,  0,  0,  0,  0,  0,  0,  0],
172
         [ 0, 57,  0,  0,  0,  0,  0,  0,  0,  0],
173
         [ 0,  0, 59,  1,  0,  0,  0,  0,  1,  0],
174
         [ 0,  0,  0, 43,  0,  0,  0,  1,  0,  0],
175
         [ 0,  0,  0,  0, 38,  0,  2,  0,  0,  0],
176
         [ 0,  0,  0,  2,  0, 48,  0,  0,  1,  0],
177
         [ 0,  1,  0,  0,  0,  0, 51,  0,  0,  0],
178
         [ 0,  0,  1,  0,  0,  0,  0, 54,  0,  0],
179
         [ 0,  0,  0,  0,  0,  1,  0,  0, 46,  0],
180
         [ 1,  1,  0,  1,  1,  0,  0,  0,  2, 42]]
181

182
        confusionSVM = [[45,  0,  0,  0,  0,  0,  0,  0,  0,  0],
183
          [ 0, 57,  0,  0,  0,  0,  0,  0,  0,  0],
184
          [ 0,  0, 59,  2,  0,  0,  0,  0,  0,  0],
185
          [ 0,  0,  0, 43,  0,  0,  0,  1,  0,  0],
186
          [ 0,  0,  0,  0, 40,  0,  0,  0,  0,  0],
187
          [ 0,  0,  0,  1,  0, 50,  0,  0,  0,  0],
188
          [ 0,  0,  0,  0,  1,  0,  51, 0,  0,  0],
189
          [ 0,  0,  1,  0,  0,  0,  0,  54, 0,  0],
190
          [ 0,  0,  0,  0,  0,  0,  0,  0, 47,  0],
191
          [ 0,  1,  0,  1,  0,  0,  0,  0,  1, 45]]
192

193
        self.assertLess(cv.norm(confusionMatrixes[0] - confusionKNN, cv.NORM_L1), normEps)
194
        self.assertLess(cv.norm(confusionMatrixes[1] - confusionSVM, cv.NORM_L1), normEps)
195

196
        self.assertLess(errors[0] - 0.034, eps)
197
        self.assertLess(errors[1] - 0.018, eps)
198

199

200
if __name__ == '__main__':
201
    NewOpenCVTests.bootstrap()
202

203
Product

Resources

Company
