-
Notifications
You must be signed in to change notification settings - Fork 964
/
pytorch_model.py
113 lines (93 loc) · 3.62 KB
/
pytorch_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
# coding:utf-8
##添加文本方向 检测模型，自动检测文字方向，0、90、180、270
##pytorch版本的OCR识别
from math import *
import cv2
import numpy as np
from PIL import Image
from angle.predict import predict as angle_detect ##文字方向检测
from crnn.crnn import crnnOcr
from ctpn.text_detect import text_detect
def crnnRec(im, text_recs, 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)
image = Image.fromarray(partImg).convert('L')
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, adjust=False, detectAngle=False):
"""
@@param:img,
@@param:model,选择的ocr模型，支持keras\pytorch版本
@@param:adjust 调整文字识别结果
@@param:detectAngle,是否检测文字朝向
"""
angle = 0
if detectAngle:
angle = angle_detect(img=np.copy(img)) ##文字朝向检测
im = Image.fromarray(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