cmaster/digit_recognizer_testing.py at master · rinalim/cmaster · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
#!/usr/bin/python
# -*- coding: utf8 -*-

#################################################################
# File name: digit_recognizer_testing.py                        #
# Description: Recognize and test digits on an image.           #
# Version: 0.0.1                                                #
# Author: Gökhan Sari                                           #
# E-mail: g-sari@g-sari.com                                     #
#################################################################

import cv2, sys
import numpy as np
from picture import Pic


class DigitRecognizerTesting:
    """Class used to test digits on an image"""

    def __init__(self):
        #self.image_to_test = Pic(pic_name="ocr_insurance_card_test_2.jpg", contour_dimension_from_h=21, contour_dimension_to_h=28)
        self.image_to_test = Pic(pic_name=sys.argv[1], contour_dimension_from_h=21, contour_dimension_to_h=28)
        self.load_training_data()
        self.model = cv2.ml.KNearest_create()
        self.model.train(self.samples, cv2.ml.ROW_SAMPLE, self.responses)

    def load_training_data(self):
        self.samples = np.loadtxt('ocr_training.data', np.float32)
        self.responses = np.loadtxt('ocr_responses.data', np.float32)
        self.responses = self.responses.reshape((self.responses.size, 1))

    def test(self):
        im = cv2.imread(self.image_to_test.pic_name)
        out = np.zeros(im.shape, np.uint8)
        gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
        thresh = cv2.adaptiveThreshold(gray, 255, 1, 1, 11, 2)
        insurance_card_number_output = "";
        _, contours, hierarchy = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
        #contours, hierarchy = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

        for cnt in contours:
            if cv2.contourArea(cnt) > self.image_to_test.contour_dimension_to_h:
                [x, y, w, h] = cv2.boundingRect(cnt)
                if h > self.image_to_test.contour_dimension_from_h and h < self.image_to_test.contour_dimension_to_h:
                    cv2.rectangle(im, (x - 1, y - 1), (x + 1 + w, y + 1 + h), (0, 255, 0), 1)
                    roi = thresh[y:y + h, x:x + w]
                    roismall = cv2.resize(roi, (10, 10))
                    roismall = roismall.reshape((1, 100))
                    roismall = np.float32(roismall)
                    retval, results, neigh_resp, dists = self.model.findNearest(roismall, k=1)
                    string = str(int((results[0][0])))
                    insurance_card_number_output += string
                    cv2.putText(out, string, (x, y + h), 0, 1, (0, 255, 0))

        reversed_insurance_card_number_output = insurance_card_number_output[::-1]
        print("Detected insurance card number: ")
        print "{:,}".format(long(reversed_insurance_card_number_output))
        #cv2.imshow('input', im)
        #cv2.imshow('output', out)
        #cv2.waitKey(0)

# Start the testing process
if __name__ == '__main__':
    DigitRecognizerTesting().test()