使用paddlepaddle做OCR识别

Posted 2023-03-02 白水baishui

1. 环境配置

Python： 使用paddlepaddle需要3.6以上的Python版本

依赖库：

• pandas，pip install pandas；
• cv2，pip install opencv-python；
• paddlepaddle
  CPU 版本：pip install paddlepaddle -i https://mirror.baidu.com/pypi/simple；
  GPU 版本，支持CUDA 10.1：python -m pip install paddlepaddle-gpu==2.1.2.post101 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html；
  GPU 版本，支持CUDA 10.2：python -m pip install paddlepaddle-gpu -i https://mirror.baidu.com/pypi/simple；
  GPU 版本，支持CUDA 11.0：python -m pip install paddlepaddle-gpu==2.1.2.post110 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html；
  GPU 版本，支持CUDA 11.2：python -m pip install paddlepaddle-gpu==2.1.2.post112 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html；
• paddlehub，pip install paddlehub；
• shapely，pip install shapely；
• pyclipper，pip install pyclipper；

2. 代码

import os
import paddlehub as hub
import cv2
from operator import itemgetter

# ------ CUDA 设置 ------ #
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # 获取GPU设备
os.environ["CUDA_VISIBLE_DEVICES"] = "0" # 设置GPU编号

# ------ 过滤ocr结果 ------ #
def filter_ocr_results(test_img_path, ocr_results):
    for img_path, result in zip(test_img_path, ocr_results):  # zip 合并两个list
        data = result['data']
        data = sorted(data, key=itemgetter('confidence'), reverse=True) # 按照置信度对list中的dict进行排序，置信度高的排在前面
        for information in data: # 遍历全部识别结果、置信度（即识别结果正确的可能性）、识别框的位置（四个角点的坐标）
            print(information)

if __name__ == "__main__":
    # ------ 更改测试数据 ------ #
    data_path = '1.jpg' # 获取图像文件的相对路径，可以是单张图片，也可以是图片数据集

    # 接将path下所有图片加载
    img_path_list = []
    basedir = os.path.abspath(os.path.dirname(__file__))  # 获取当前文件的绝对路径
    input_path_dir = basedir + '/' + data_path  # 文件夹路径
    img_path_list.append(input_path_dir)
    # 读取测试文件夹中的图片路径
    np_images = [cv2.imread(image_path) for image_path in img_path_list]

    # OCR图片的保存路径，默认设为
    output_path = basedir

    # 输出图像文件路径和结果输出路径
    print("input_path_dir: " + input_path_dir)
    print("output_path: " + output_path)

    # 加载预训练模型
    ocr = hub.Module(name="chinese_ocr_db_crnn_server")

    # 进行OCR识别
    ocr_results = ocr.recognize_text(
        images=np_images,  # 图片数据，ndarray.shape 为 [H, W, C]，BGR格式；
        use_gpu=False,  # 是否使用 GPU；若使用GPU，请先设置CUDA_VISIBLE_DEVICES环境变量
        output_dir=output_path,  # 识别结果图片文件的保存路径；
        visualization=True,  # 是否将识别结果保存为图片文件；
        box_thresh=0.0001,  # 检测文本框置信度的阈值，越小识别的内容越多，但也更可能识别出不相关的内容，不可设置为0；
        text_thresh=0.5)  # 识别中文文本置信度的阈值，是指将内容识别为中文字符的阈值参数，由于也需要对英文和数字进行识别，该值不宜过小；
    # 过滤识别结果
    filter_ocr_results(img_path_list, ocr_results)

3. 识别结果

我们用一张图片来进行展示，原图：

识别结果：

程序输出：

'text': '12', 'confidence': 0.9999337792396545, 'text_box_position': [[276, 211], [320, 211], [320, 254], [276, 254]]
'text': 'DAWNRAYS', 'confidence': 0.9993633031845093, 'text_box_position': [[682, 560], [763, 557], [764, 576], [683, 579]]
'text': '宜立舒', 'confidence': 0.9987321496009827, 'text_box_position': [[122, 203], [220, 202], [220, 234], [122, 235]]
'text': '与适宜抗生素合用', 'confidence': 0.9979811906814575, 'text_box_position': [[127, 544], [284, 540], [284, 562], [127, 566]]
'text': '复方雷尼替丁胶囊', 'confidence': 0.9938433170318604, 'text_box_position': [[101, 268], [474, 268], [474, 317], [101, 317]]
'text': '批准文号：国药准字H20030953', 'confidence': 0.9904718995094299, 'text_box_position': [[358, 184], [585, 187], [585, 205], [358, 202]]
'text': '治疗幽门螺旋杆菌感染', 'confidence': 0.9837646484375, 'text_box_position': [[128, 574], [320, 568], [320, 589], [128, 595]]
'text': '十二指肠溃疡', 'confidence': 0.9615536332130432, 'text_box_position': [[126, 515], [246, 512], [246, 534], [126, 537]]
'text': 'Compound RanitidineCapsules', 'confidence': 0.9529160857200623, 'text_box_position': [[99, 329], [475, 326], [475, 350], [99, 353]]
'text': '每粒含盐酸雷尼替丁（以CHNO，S计）', 'confidence': 0.9517455101013184, 'text_box_position': [[373, 221], [581, 223], [581, 241], [373, 239]]
'text': '良性胃溃疡', 'confidence': 0.9417906999588013, 'text_box_position': [[126, 485], [226, 482], [227, 507], [127, 510]]
'text': '适应症：', 'confidence': 0.9281517267227173, 'text_box_position': [[127, 451], [230, 451], [230, 479], [127, 479]]
'text': 'Capsules150mg、枸橡酸秘钾（以餐计）110mg', 'confidence': 0.9192306995391846, 'text_box_position': [[315, 237], [587, 238], [587, 259], [315, 258]]
'text': '东瑞制藥', 'confidence': 0.81662917137146, 'text_box_position': [[683, 535], [763, 533], [764, 559], [684, 561]]