1 Star 0 Fork 0

灬花落丶莫相离/CHINESE-OCR

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
该仓库未声明开源许可证文件(LICENSE),使用请关注具体项目描述及其代码上游依赖。
克隆/下载
pytorch_model.py 3.62 KB
一键复制 编辑 原始数据 按行查看 历史
xiaofengalg 提交于 2018-04-14 22:33 . push all files
# 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
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
1
https://gitee.com/bjt.zhao/CHINESE-OCR.git
git@gitee.com:bjt.zhao/CHINESE-OCR.git
bjt.zhao
CHINESE-OCR
CHINESE-OCR
master

搜索帮助

0d507c66 1850385 C8b1a773 1850385