len(alphabet) + 1,
nh,
n_rnn=2,
leakyRelu=False,
lstmFlag=LSTMFLAG)
elif crnn_type == "lite_lstm" or crnn_type == "lite_dense":
crnn_net = LiteCrnn(32,
1,
len(alphabet) + 1,
nh,
n_rnn=2,
leakyRelu=False,
lstmFlag=LSTMFLAG)
assert crnn_type is not None
crnn_handle = CRNNHandle(crnn_model_path, crnn_net, gpu_id=GPU_ID)
crnn_vertical_handle = None
if crnn_vertical_model_path is not None:
crnn_vertical_net = LiteCrnn(32,
1,
len(alphabet) + 1,
nh,
n_rnn=2,
leakyRelu=False,
lstmFlag=True)
crnn_vertical_handle = CRNNHandle(crnn_vertical_model_path,
crnn_vertical_net,
gpu_id=GPU_ID)
assert angle_type in ["shufflenetv2_05"]
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, 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
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)
len(alphabet) + 1,
nh,
n_rnn=2,
leakyRelu=False,
lstmFlag=LSTMFLAG)
elif crnn_type == "lite_lstm" or crnn_type == "lite_dense":
crnn_net = LiteCrnn(32,
1,
len(alphabet) + 1,
nh,
n_rnn=2,
leakyRelu=False,
lstmFlag=LSTMFLAG)
assert crnn_type is not None
crnn_handle = CRNNHandle(crnn_model_path, crnn_net, gpu_id=GPU_ID)
assert angle_type in ["shufflenetv2_05"]
if angle_type == "shufflenetv2_05":
angle_net = shufflenet_v2_x0_5(num_classes=len(lable_map_dict),
pretrained=False)
angle_handle = AangleClassHandle(angle_model_path, angle_net, gpu_id=GPU_ID)
def crnnRec(im,
rects_re,
leftAdjust=False,
rightAdjust=False,
alph=0.2,