pytorch_model.py

#coding:utf-8
##添加文本方向 检测模型，自动检测文字方向，0、90、180、270
##pytorch版本的OCR识别
from ctpn.text_detect import text_detect
from angle.predict import predict as angle_detect##文字方向检测
from crnn.crnn import crnnOcr
from math import *
import numpy as np
import cv2
from PIL import Image
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