forked from xiaofengShi/CHINESE-OCR
-
Notifications
You must be signed in to change notification settings - Fork 2
/
model.py
136 lines (115 loc) · 4.33 KB
/
model.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
136
# coding:utf-8
##添加文本方向 检测模型，自动检测文字方向，0、90、180、270
from math import *
import cv2
import numpy as np
from PIL import Image
import sys
sys.path.append("ocr")
from angle.predict import predict as angle_detect ##文字方向检测
from crnn.crnn import crnnOcr
from ctpn.text_detect import text_detect
from ocr.model import predict as ocr
def crnnRec(im, text_recs, ocrMode='keras', adjust=False):
"""
crnn模型，ocr识别
@@model,
@@converter,
@@im:Array
@@text_recs:text box
"""
index = 0
results = {}
xDim, yDim = im.shape[1], im.shape[0]
for index, rec in enumerate(text_recs):
results[index] = [
rec,
]
xlength = int((rec[6] - rec[0]) * 0.1)
ylength = int((rec[7] - rec[1]) * 0.2)
if adjust:
pt1 = (max(1, rec[0] - xlength), max(1, rec[1] - ylength))
pt2 = (rec[2], rec[3])
pt3 = (min(rec[6] + xlength, xDim - 2),
min(yDim - 2, rec[7] + ylength))
pt4 = (rec[4], rec[5])
else:
pt1 = (max(1, rec[0]), max(1, rec[1]))
pt2 = (rec[2], rec[3])
pt3 = (min(rec[6], xDim - 2), min(yDim - 2, rec[7]))
pt4 = (rec[4], rec[5])
degree = degrees(atan2(pt2[1] - pt1[1], pt2[0] - pt1[0])) ##图像倾斜角度
partImg = dumpRotateImage(im, degree, pt1, pt2, pt3, pt4)
# 根据ctpn进行识别出的文字区域，进行不同文字区域的crnn识别
image = Image.fromarray(partImg).convert('L')
# 进行识别出的文字识别
if ocrMode == 'keras':
sim_pred = ocr(image)
else:
sim_pred = crnnOcr(image)
results[index].append(sim_pred) ##识别文字
return results
def dumpRotateImage(img, degree, pt1, pt2, pt3, pt4):
height, width = img.shape[:2]
heightNew = int(width * fabs(sin(radians(degree))) +
height * fabs(cos(radians(degree))))
widthNew = int(height * fabs(sin(radians(degree))) +
width * fabs(cos(radians(degree))))
matRotation = cv2.getRotationMatrix2D((width / 2, height / 2), degree, 1)
matRotation[0, 2] += (widthNew - width) / 2
matRotation[1, 2] += (heightNew - height) / 2
imgRotation = cv2.warpAffine(
img, matRotation, (widthNew, heightNew), borderValue=(255, 255, 255))
pt1 = list(pt1)
pt3 = list(pt3)
[[pt1[0]], [pt1[1]]] = np.dot(matRotation,
np.array([[pt1[0]], [pt1[1]], [1]]))
[[pt3[0]], [pt3[1]]] = np.dot(matRotation,
np.array([[pt3[0]], [pt3[1]], [1]]))
ydim, xdim = imgRotation.shape[:2]
imgOut = imgRotation[max(1, int(pt1[1])):min(ydim - 1, int(pt3[1])),
max(1, int(pt1[0])):min(xdim - 1, int(pt3[0]))]
# height,width=imgOut.shape[:2]
return imgOut
def model(img, model='keras', adjust=False, detectAngle=False):
"""
@@param:img,
@@param:model,选择的ocr模型，支持keras\pytorch版本
@@param:adjust 调整文字识别结果
@@param:detectAngle,是否检测文字朝向
"""
angle = 0
if detectAngle:
# 进行文字旋转方向检测，分为[0, 90, 180, 270]四种情况
angle = angle_detect(img=np.copy(img)) ##文字朝向检测
print('The angel of this character is:', angle)
im = Image.fromarray(img)
print('Rotate the array of this img!')
if angle == 90:
im = im.transpose(Image.ROTATE_90)
elif angle == 180:
im = im.transpose(Image.ROTATE_180)
elif angle == 270:
im = im.transpose(Image.ROTATE_270)
img = np.array(im)
# 进行图像中的文字区域的识别
text_recs, tmp, img=text_detect(img)
# 识别区域排列
text_recs = sort_box(text_recs)
#
result = crnnRec(img, text_recs, model, adjust=adjust)
return result, tmp, angle
def sort_box(box):
"""
对box排序,及页面进行排版
text_recs[index, 0] = x1
text_recs[index, 1] = y1
text_recs[index, 2] = x2
text_recs[index, 3] = y2
text_recs[index, 4] = x3
text_recs[index, 5] = y3
text_recs[index, 6] = x4
text_recs[index, 7] = y4
"""
box = sorted(box, key=lambda x: sum([x[1], x[3], x[5], x[7]]))
return box