-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
135 lines (103 loc) · 3.17 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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
# -*- coding: utf-8 -*-
"""
@author: Dani, richard
"""
#import net_training as nt
import get_chars as gc
import string
import net_training as net
import numpy as np
import utils
reload(gc)
import matplotlib.pyplot as plt
from skimage import io
import cv2
import rotation
reload(rotation)
import im_treatement
reload(im_treatement)
import morf
reload(morf)
import sys
lines = []
trans_table = list(string.digits)+list(string.ascii_uppercase)+list(string.ascii_lowercase)
global_params = {
"TEST_MODE": {
"line_detect":False,
"word_detect":False,
"char_detect":False,
"im_treatement":False,
"rotation":False,
},
"new_net":False,
"image_name": "sign.jpg",
"contour" : 1200,
"status_msg" : True
}
def neural_predict(im,model):
"""
Funcio que passa la imatge a la xarxa neuronal i retorna la lletra predita.
* Inputs:
- im = skimage.io image
- model = model neuronal
*Outputs:
- letter = caracter predit
"""
im = im / np.max(im)
try:
im = utils.rescale(im.astype(float),(64,64),4)
im = np.reshape(im,(1,64,64,1))
letter = trans_table[net.net_predict(im,model)]
return letter
except:
return "?"
def process_img(im, params=global_params):
"""
Funció que fa tot el procediment per obtenir el text de la imatge.
* Inputs:
- im = skimage.io image
- params = diccionari de parametres
*Outputs:
- text = String amb el text trobat a la imatge
"""
plt.close('all')
if(len(im.shape) != 2):
im = cv2.cvtColor(im, cv2.COLOR_RGB2GRAY)
if(params["status_msg"]): print "Starting new request"
im = im_treatement.treatement(im, params)
im = morf.apply_morf(im, params)
im = rotation.fix_rotation(im, params)
lines = gc.get_all(im, params)
if(params["status_msg"]): print "Predicting results..."
if(params["status_msg"]): print "====================="
model = net.loadmodel('model')
text = ""
for line in lines:
for word in line:
for letter in word:
result = neural_predict(letter,model)
if(params["status_msg"]): sys.stdout.write(result)
text += result
if(params["status_msg"]): sys.stdout.write(" ")
text += " "
if(params["status_msg"]): sys.stdout.write("\n")
text += "\n"
if(params["status_msg"]): print "========================"
if(params["status_msg"]): print "Finished all operations."
return text
def main():
"""
Benvingut a l'OCR!
Canvia la imatge dessitgada al parametre "image_name" de la variable
global_param i assegurat que aquesta imatge esta a la carpeta ./img
Aquesta funcio et guardara a text.txt el resultat de la compilacio. Vagi
de gust!
"""
global lines
im = io.imread("./img/"+global_params["image_name"])
text = process_img(im, global_params)
f = open("test.txt","w")
f.write(text)
f.close()
if __name__ == "__main__":
main()