-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
65 lines (52 loc) · 1.69 KB
/
main.py
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
65
import cv2 as cv
import histogram_segmentation as hist
import feature_extraction as fe
import classifier as cl
LETTER_SIZE = 40
DIGIT_SIZE = 20
EXTRACTION_STEP = 4
BIN_THRESHOLD = 175
def show_result(src, rects, result, use_letters):
i = 0
for word in rects:
t = ""
for _ in word[-1]:
if use_letters:
t += chr(int(result[i]) + 65)
else:
t += chr(int(result[i]) + 48)
i += 1
cv.rectangle(src, (word[2], word[0]), (word[3], word[1]), (255, 0, 0), 3)
cv.putText(src, t, (word[2], word[0]), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0))
cv.imshow("t", src)
cv.waitKey()
def predict(path, use_letters, used_classifier):
src = cv.imread("data\\" + path)
if use_letters:
dim = LETTER_SIZE
else:
dim = DIGIT_SIZE
histogram = hist.HistogramSegmentation(dim, dim, BIN_THRESHOLD, histogram_threshold=0)
ext = fe.FeatureExtraction(EXTRACTION_STEP)
words, b, rects = histogram.segment(src)
imgs = []
for w in words:
y0, y1, x0, x1, letters = w
for l in letters:
imgs.append(l)
feature_vectors = ext.extract(imgs)
classifier = cl.Classifier(used_classifier, 64, 32)
classifier.load(use_letters)
result = classifier.predict((feature_vectors, 0))
show_result(src, rects, result, use_letters)
def main():
path = input("Path to file:")
use_letters = input("Letters or digits?")
used_classifier = input("Choose classifier: svm, bayes, rtrees, ann")
if use_letters == "letters":
letters = True
else:
letters = False
predict(path, letters, used_classifier)
if __name__ == "__main__":
main()