forked from DayBreak-u/chineseocr_lite
-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
85 lines (64 loc) · 2.53 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
from config import *
from crnn import CRNNHandle
from angnet import AngleNetHandle
from utils import draw_bbox, crop_rect, sorted_boxes, get_rotate_crop_image
from PIL import Image
import numpy as np
import cv2
import copy
from dbnet.dbnet_infer import DBNET
import time
import traceback
class OcrHandle(object):
def __init__(self):
self.text_handle = DBNET(model_path)
self.crnn_handle = CRNNHandle(crnn_model_path)
if angle_detect:
self.angle_handle = AngleNetHandle(angle_net_path)
def crnnRecWithBox(self,im, boxes_list,score_list):
"""
crnn模型，ocr识别
@@model,
@@converter,
@@im:Array
@@text_recs:text box
@@ifIm:是否输出box对应的img
"""
results = []
boxes_list = sorted_boxes(np.array(boxes_list))
line_imgs = []
for index, (box, score) in enumerate(zip(boxes_list[:angle_detect_num], score_list[:angle_detect_num])):
tmp_box = copy.deepcopy(box)
partImg_array = get_rotate_crop_image(im, tmp_box.astype(np.float32))
partImg = Image.fromarray(partImg_array).convert("RGB")
line_imgs.append(partImg)
angle_res = False
if angle_detect:
angle_res = self.angle_handle.predict_rbgs(line_imgs)
count = 1
for index, (box ,score) in enumerate(zip(boxes_list,score_list)):
tmp_box = copy.deepcopy(box)
partImg_array = get_rotate_crop_image(im, tmp_box.astype(np.float32))
partImg = Image.fromarray(partImg_array).convert("RGB")
if angle_detect and angle_res:
partImg = partImg.rotate(180)
if not is_rgb:
partImg = partImg.convert('L')
try:
if is_rgb:
simPred = self.crnn_handle.predict_rbg(partImg) ##识别的文本
else:
simPred = self.crnn_handle.predict(partImg) ##识别的文本
except Exception as e:
print(traceback.format_exc())
continue
if simPred.strip() != '':
results.append([tmp_box,"{}、 ".format(count)+ simPred,score])
count += 1
return results
def text_predict(self,img,short_size):
boxes_list, score_list = self.text_handle.process(np.asarray(img).astype(np.uint8),short_size=short_size)
result = self.crnnRecWithBox(np.array(img), boxes_list,score_list)
return result
if __name__ == "__main__":
pass