当前位置:
首页
>
文章详情
利用modelscope自主搭建OCR图像文字识别的方法
2023-09-20

ocr文字识别

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import cv2

from PIL import Image
import numpy as np
import math
from PIL import ImageDraw
from urllib import request
ocr_detection = pipeline(Tasks.ocr_detection, model='damo/cv_resnet18_ocr-detection-line-level_damo')
ocr_recognition = pipeline(Tasks.ocr_recognition, model='damo/cv_convnextTiny_ocr-recognition-general_damo')

### 使用url
def crop_image(img, position):
    def distance(x1,y1,x2,y2):
        return math.sqrt(pow(x1 - x2, 2) + pow(y1 - y2, 2))    
    position = position.tolist()
    for i in range(4):
        for j in range(i+1, 4):
            if(position[i][0] > position[j][0]):
                tmp = position[j]
                position[j] = position[i]
                position[i] = tmp
    if position[0][1] > position[1][1]:
        tmp = position[0]
        position[0] = position[1]
        position[1] = tmp

    if position[2][1] > position[3][1]:
        tmp = position[2]
        position[2] = position[3]
        position[3] = tmp

    x1, y1 = position[0][0], position[0][1]
    x2, y2 = position[2][0], position[2][1]
    x3, y3 = position[3][0], position[3][1]
    x4, y4 = position[1][0], position[1][1]

    corners = np.zeros((4,2), np.float32)
    corners[0] = [x1, y1]
    corners[1] = [x2, y2]
    corners[2] = [x4, y4]
    corners[3] = [x3, y3]

    img_width = distance((x1+x4)/2, (y1+y4)/2, (x2+x3)/2, (y2+y3)/2)
    img_height = distance((x1+x2)/2, (y1+y2)/2, (x4+x3)/2, (y4+y3)/2)

    corners_trans = np.zeros((4,2), np.float32)
    corners_trans[0] = [0, 0]
    corners_trans[1] = [img_width - 1, 0]
    corners_trans[2] = [0, img_height - 1]
    corners_trans[3] = [img_width - 1, img_height - 1]

    transform = cv2.getPerspectiveTransform(corners, corners_trans)
    dst = cv2.warpPerspective(img, transform, (int(img_width), int(img_height)))
    return dst
def order_point(coor):
    arr = np.array(coor).reshape([4, 2])
    sum_ = np.sum(arr, 0)
    centroid = sum_ / arr.shape[0]
    theta = np.arctan2(arr[:, 1] - centroid[1], arr[:, 0] - centroid[0])
    sort_points = arr[np.argsort(theta)]
    sort_points = sort_points.reshape([4, -1])
    if sort_points[0][0] > centroid[0]:
        sort_points = np.concatenate([sort_points[3:], sort_points[:3]])
    sort_points = sort_points.reshape([4, 2]).astype('float32')
    return sort_points

def text_detection(image_full, ocr_detection):
    det_result = ocr_detection(image_full)
    det_result = det_result['polygons']
    # sort detection result with coord
    det_result_list = det_result.tolist()
    det_result_list = sorted(det_result_list, key=lambda x: 0.01*sum(x[::2])/4+sum(x[1::2])/4)     
    return np.array(det_result_list)

def text_recognition(det_result, image_full, ocr_recognition):
    output = []
    for i in range(det_result.shape[0]):
        pts = order_point(det_result[i])
        image_crop = crop_image(image_full, pts)
        result = ocr_recognition(image_crop)
        output.append(",".join(result['text']))
    
    return output

def text_ocr(image_full):
    det_result = text_detection(image_full, ocr_detection)
    ocr_result = text_recognition(det_result, image_full, ocr_recognition)        
    
    return ocr_result

imgurl = 'http://aiapi.deituicms.com/attach/2023/09/20/861e3442c36cb45b79ef3cd2c62f3ece.jpg'
image_path = './static/ocr.png'
request.urlretrieve(imgurl, image_path)
image = cv2.imread(image_path)
result = text_ocr(image)
print(result)
content="\r\n".join(result)
print(content)




评论列表
写跟帖
取消
评论
{{item.nickname}}
{{item.ip_city}}
{{item.timeago}}
{{item.content}}
加载更多