devbox
我家自动化
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

paddle ocr v4 2.6.1实战笔记_paddle:2.6.1-jupyter

作者：我家自动化 | 2024-05-19 15:50:44

赞

踩

paddle:2.6.1-jupyter

目录

效果图：

安装

模型权重是自动下载，如果提前下载会报错。

识别orc，并opencv可视化结果，支持中文可视化

官方原版预测可视化：

效果图：

安装

安装2.5.2识别结果为空

pip install paddlepaddle-gpu==2.6.1

模型权重是自动下载，如果提前下载会报错。

测试代码：

  1.  
  2. import os
  3. import time
  4. from paddleocr import PaddleOCR
  5.  
  6. filepath = r"weights/123.jpg"
  7.  
  8. ocr_model = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=True, show_log=1,
  9.                 det_db_box_thresh=0.1, use_dilation=True,
  10.                 det_model_dir='weight/ch_PP-OCRv4_det_server_infer.tar',
  11.                 cls_model_dir='weight/ch_ppocr_mobile_v2.0_cls_infer.tar',
  12.                 rec_model_dir='weight/ch_PP-OCRv4_rec_server_infer.tar')
  13.  
  14. t1 = time.time()
  15. for i in range(1):
  16.     result = ocr_model.ocr(img=filepath, det=True, rec=True, cls=True)[0]
  17. t2 = time.time()
  18. print((t2-t1) / 10)
  19.  
  20. for res_str in result:
  21.     print(res_str)

识别orc，并opencv可视化结果，支持中文可视化

  1. import codecs
  2. import os
  3. import time
  4.  
  5. import cv2
  6. import numpy as np
  7. from PIL import ImageFont
  8. from PIL import Image
  9. from PIL import ImageDraw
  10.  
  11. from paddleocr import PaddleOCR
  12.  
  13. filepath = r"weights/124.jpg"
  14.  
  15. ocr_model = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=True, show_log=1,
  16.                 det_db_box_thresh=0.1, use_dilation=True,
  17.                 det_model_dir='weight/ch_PP-OCRv4_det_server_infer.tar',
  18.                 cls_model_dir='weight/ch_ppocr_mobile_v2.0_cls_infer.tar',
  19.                 rec_model_dir='weight/ch_PP-OCRv4_rec_server_infer.tar')
  20.  
  21. t1 = time.time()
  22. for i in range(1):
  23.     result = ocr_model.ocr(img=filepath, det=True, rec=True, cls=True)[0]
  24. t2 = time.time()
  25. print((t2-t1) / 10)
  26.  
  27. font_path = 'simhei.ttf'  # 需要替换为你的中文字体路径
  28. font = ImageFont.truetype(font_path, 24)
  29. def cv2AddChineseText(img, text, position, textColor=(0, 255, 0), textSize=30):
  30.     img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
  31.     draw = ImageDraw.Draw(img)
  32.     draw.text(position, text, textColor, font=font)
  33.     return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
  34.  
  35. image=cv2.imread(filepath)
  36.  
  37. ocr_index=0
  38. for res_str in result:
  39.     if res_str[0][0][0]>36 and res_str[0][2][0]<84:
  40.         print(ocr_index,res_str)
  41.         points=res_str[0]
  42.         text = res_str[1][0]
  43.         points = np.array(points, dtype=np.int32).reshape((-1, 1, 2))
  44.         cv2.polylines(image, [points], isClosed=True, color=(255, 0, 0), thickness=2)
  45.         text_position = (int(points[0][0][0]), int(points[0][0][1] + 20))  # 微调文本位置
  46.  
  47.         # cv2.putText(image, '中文文本', (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 3)
  48.         image= cv2AddChineseText(image, text, text_position, textColor=(0, 255, 0), textSize=30)
  49.         print(ocr_index)
  50.     if res_str[0][0][0]>346 and res_str[0][2][0]<391:
  51.         print(ocr_index,res_str)
  52.         points=res_str[0]
  53.         text = res_str[1][0]
  54.         points = np.array(points, dtype=np.int32).reshape((-1, 1, 2))
  55.         cv2.polylines(image, [points], isClosed=True, color=(255, 0, 0), thickness=2)
  56.         text_position = (int(points[0][0][0]), int(points[0][0][1] + 20))  # 微调文本位置
  57.  
  58.         # cv2.putText(image, '中文文本', (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 3)
  59.         image= cv2AddChineseText(image, text, text_position, textColor=(0, 255, 0), textSize=30)
  60.     if res_str[0][0][0]>658 and res_str[0][2][0]<705:
  61.         print(ocr_index,res_str)
  62.         points=res_str[0]
  63.         text=res_str[1][0]
  64.         points=np.array(points,dtype=np.int32).reshape((-1, 1, 2))
  65.         cv2.polylines(image, [points], isClosed=True, color=(255, 0, 0), thickness=2)
  66.         text_position = (int(points[0][0][0]), int(points[0][0][1] + 20))  # 微调文本位置
  67.         image= cv2AddChineseText(image, text, text_position, textColor=(0, 255, 0), textSize=30)
  68.  
  69. cv2.imshow('Image with Rectangle and Text', image)
  70. cv2.waitKey(0)

官方原版预测可视化：

  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. #     http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14.  
  15. import os
  16. import sys
  17. import importlib
  18.  
  19. __dir__ = os.path.dirname(__file__)
  20.  
  21. import paddle
  22. from paddle.utils import try_import
  23.  
  24. sys.path.append(os.path.join(__dir__, ""))
  25.  
  26. import cv2
  27. import logging
  28. import numpy as np
  29. from pathlib import Path
  30. import base64
  31. from io import BytesIO
  32. from PIL import Image, ImageFont, ImageDraw
  33. from tools.infer import predict_system
  34.  
  35.  
  36. def _import_file(module_name, file_path, make_importable=False):
  37.     spec = importlib.util.spec_from_file_location(module_name, file_path)
  38.     module = importlib.util.module_from_spec(spec)
  39.     spec.loader.exec_module(module)
  40.     if make_importable:
  41.         sys.modules[module_name] = module
  42.     return module
  43.  
  44.  
  45. tools = _import_file("tools", os.path.join(__dir__, "tools/__init__.py"), make_importable=True)
  46. ppocr = importlib.import_module("ppocr", "paddleocr")
  47. ppstructure = importlib.import_module("ppstructure", "paddleocr")
  48. from ppocr.utils.logging import get_logger
  49.  
  50. logger = get_logger()
  51. from ppocr.utils.utility import (check_and_read, get_image_file_list, alpha_to_color, binarize_img, )
  52. from ppocr.utils.network import (maybe_download, download_with_progressbar, is_link, confirm_model_dir_url, )
  53. from tools.infer.utility import draw_ocr, str2bool, check_gpu
  54. from ppstructure.utility import init_args, draw_structure_result
  55. from ppstructure.predict_system import StructureSystem, save_structure_res, to_excel
  56.  
  57. logger = get_logger()
  58. __all__ = ["PaddleOCR", "PPStructure", "draw_ocr", "draw_structure_result", "save_structure_res", "download_with_progressbar", "to_excel", ]
  59.  
  60. SUPPORT_DET_MODEL = ["DB"]
  61. VERSION = "2.8.0"
  62. SUPPORT_REC_MODEL = ["CRNN", "SVTR_LCNet"]
  63. BASE_DIR = os.path.expanduser("~/.paddleocr/")
  64.  
  65. DEFAULT_OCR_MODEL_VERSION = "PP-OCRv4"
  66. SUPPORT_OCR_MODEL_VERSION = ["PP-OCR", "PP-OCRv2", "PP-OCRv3", "PP-OCRv4"]
  67. DEFAULT_STRUCTURE_MODEL_VERSION = "PP-StructureV2"
  68. SUPPORT_STRUCTURE_MODEL_VERSION = ["PP-Structure", "PP-StructureV2"]
  69. MODEL_URLS = {"OCR": {"PP-OCRv4": {"det": {"ch": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/chinese/ch_PP-OCRv4_det_infer.tar", }, "en": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_det_infer.tar", },
  70.     "ml": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/Multilingual_PP-OCRv3_det_infer.tar"}, },
  71.     "rec": {"ch": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/chinese/ch_PP-OCRv4_rec_infer.tar", "dict_path": "./ppocr/utils/ppocr_keys_v1.txt", }, "en": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/english/en_PP-OCRv4_rec_infer.tar", "dict_path": "./ppocr/utils/en_dict.txt", },
  72.         "korean": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/korean_PP-OCRv4_rec_infer.tar", "dict_path": "./ppocr/utils/dict/korean_dict.txt", },
  73.         "japan": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/japan_PP-OCRv4_rec_infer.tar", "dict_path": "./ppocr/utils/dict/japan_dict.txt", },
  74.         "chinese_cht": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/chinese_cht_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/chinese_cht_dict.txt", },
  75.         "ta": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/ta_PP-OCRv4_rec_infer.tar", "dict_path": "./ppocr/utils/dict/ta_dict.txt", },
  76.         "te": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/te_PP-OCRv4_rec_infer.tar", "dict_path": "./ppocr/utils/dict/te_dict.txt", },
  77.         "ka": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/ka_PP-OCRv4_rec_infer.tar", "dict_path": "./ppocr/utils/dict/ka_dict.txt", },
  78.         "latin": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/latin_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/latin_dict.txt", },
  79.         "arabic": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/arabic_PP-OCRv4_rec_infer.tar", "dict_path": "./ppocr/utils/dict/arabic_dict.txt", },
  80.         "cyrillic": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/cyrillic_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/cyrillic_dict.txt", },
  81.         "devanagari": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/devanagari_PP-OCRv4_rec_infer.tar", "dict_path": "./ppocr/utils/dict/devanagari_dict.txt", }, }, "cls": {"ch": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar", }}, },
  82.     "PP-OCRv3": {"det": {"ch": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/ch_PP-OCRv3_det_infer.tar", }, "en": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_det_infer.tar", },
  83.         "ml": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/Multilingual_PP-OCRv3_det_infer.tar"}, },
  84.         "rec": {"ch": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/ch_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/ppocr_keys_v1.txt", }, "en": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/en_dict.txt", },
  85.             "korean": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/korean_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/korean_dict.txt", },
  86.             "japan": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/japan_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/japan_dict.txt", },
  87.             "chinese_cht": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/chinese_cht_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/chinese_cht_dict.txt", },
  88.             "ta": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/ta_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/ta_dict.txt", },
  89.             "te": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/te_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/te_dict.txt", },
  90.             "ka": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/ka_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/ka_dict.txt", },
  91.             "latin": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/latin_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/latin_dict.txt", },
  92.             "arabic": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/arabic_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/arabic_dict.txt", },
  93.             "cyrillic": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/cyrillic_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/cyrillic_dict.txt", },
  94.             "devanagari": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/devanagari_PP-OCRv3_rec_infer.tar", "dict_path": "./ppocr/utils/dict/devanagari_dict.txt", }, }, "cls": {"ch": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar", }}, },
  95.     "PP-OCRv2": {"det": {"ch": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_det_infer.tar", }, }, "rec": {"ch": {"url": "https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_rec_infer.tar", "dict_path": "./ppocr/utils/ppocr_keys_v1.txt", }},
  96.         "cls": {"ch": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar", }}, }, "PP-OCR": {
  97.         "det": {"ch": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar", }, "en": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_ppocr_mobile_v2.0_det_infer.tar", },
  98.             "structure": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_det_infer.tar"}, }, "rec": {"ch": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/ppocr_keys_v1.txt", },
  99.             "en": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_number_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/en_dict.txt", },
  100.             "french": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/french_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/french_dict.txt", },
  101.             "german": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/german_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/german_dict.txt", },
  102.             "korean": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/korean_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/korean_dict.txt", },
  103.             "japan": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/japan_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/japan_dict.txt", },
  104.             "chinese_cht": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/chinese_cht_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/chinese_cht_dict.txt", },
  105.             "ta": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ta_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/ta_dict.txt", },
  106.             "te": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/te_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/te_dict.txt", },
  107.             "ka": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ka_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/ka_dict.txt", },
  108.             "latin": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/latin_ppocr_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/latin_dict.txt", },
  109.             "arabic": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/arabic_ppocr_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/arabic_dict.txt", },
  110.             "cyrillic": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/cyrillic_ppocr_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/cyrillic_dict.txt", },
  111.             "devanagari": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/devanagari_ppocr_mobile_v2.0_rec_infer.tar", "dict_path": "./ppocr/utils/dict/devanagari_dict.txt", },
  112.             "structure": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_rec_infer.tar", "dict_path": "ppocr/utils/dict/table_dict.txt", }, }, "cls": {"ch": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar", }}, }, },
  113.     "STRUCTURE": {"PP-Structure": {"table": {"en": {"url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_structure_infer.tar", "dict_path": "ppocr/utils/dict/table_structure_dict.txt", }}}, "PP-StructureV2": {
  114.         "table": {"en": {"url": "https://paddleocr.bj.bcebos.com/ppstructure/models/slanet/en_ppstructure_mobile_v2.0_SLANet_infer.tar", "dict_path": "ppocr/utils/dict/table_structure_dict.txt", },
  115.             "ch": {"url": "https://paddleocr.bj.bcebos.com/ppstructure/models/slanet/ch_ppstructure_mobile_v2.0_SLANet_infer.tar", "dict_path": "ppocr/utils/dict/table_structure_dict_ch.txt", }, },
  116.         "layout": {"en": {"url": "https://paddleocr.bj.bcebos.com/ppstructure/models/layout/picodet_lcnet_x1_0_fgd_layout_infer.tar", "dict_path": "ppocr/utils/dict/layout_dict/layout_publaynet_dict.txt", },
  117.             "ch": {"url": "https://paddleocr.bj.bcebos.com/ppstructure/models/layout/picodet_lcnet_x1_0_fgd_layout_cdla_infer.tar", "dict_path": "ppocr/utils/dict/layout_dict/layout_cdla_dict.txt", }, }, }, }, }
  118.  
  119.  
  120. def parse_args(mMain=True):
  121.     import argparse
  122.  
  123.     parser = init_args()
  124.     parser.add_help = mMain
  125.     parser.add_argument("--lang", type=str, default="ch")
  126.     parser.add_argument("--det", type=str2bool, default=True)
  127.     parser.add_argument("--rec", type=str2bool, default=True)
  128.     parser.add_argument("--type", type=str, default="ocr")
  129.     parser.add_argument("--savefile", type=str2bool, default=False)
  130.     parser.add_argument("--ocr_version", type=str, choices=SUPPORT_OCR_MODEL_VERSION, default="PP-OCRv4", help="OCR Model version, the current model support list is as follows: "
  131.                                                                                                                "1. PP-OCRv4/v3 Support Chinese and English detection and recognition model, and direction classifier model"
  132.                                                                                                                "2. PP-OCRv2 Support Chinese detection and recognition model. "
  133.                                                                                                                "3. PP-OCR support Chinese detection, recognition and direction classifier and multilingual recognition model.", )
  134.     parser.add_argument("--structure_version", type=str, choices=SUPPORT_STRUCTURE_MODEL_VERSION, default="PP-StructureV2", help="Model version, the current model support list is as follows:"
  135.                                                                                                                                  " 1. PP-Structure Support en table structure model."
  136.                                                                                                                                  " 2. PP-StructureV2 Support ch and en table structure model.", )
  137.  
  138.     for action in parser._actions:
  139.         if action.dest in ["rec_char_dict_path", "table_char_dict_path", "layout_dict_path", ]:
  140.             action.default = None
  141.     if mMain:
  142.         return parser.parse_args()
  143.     else:
  144.         inference_args_dict = {}
  145.         for action in parser._actions:
  146.             inference_args_dict[action.dest] = action.default
  147.         return argparse.Namespace(**inference_args_dict)
  148.  
  149.  
  150. def parse_lang(lang):
  151.     latin_lang = ["af", "az", "bs", "cs", "cy", "da", "de", "es", "et", "fr", "ga", "hr", "hu", "id", "is", "it", "ku", "la", "lt", "lv", "mi", "ms", "mt", "nl", "no", "oc", "pi", "pl", "pt", "ro", "rs_latin", "sk", "sl", "sq", "sv", "sw", "tl", "tr", "uz", "vi", "french", "german", ]
  152.     arabic_lang = ["ar", "fa", "ug", "ur"]
  153.     cyrillic_lang = ["ru", "rs_cyrillic", "be", "bg", "uk", "mn", "abq", "ady", "kbd", "ava", "dar", "inh", "che", "lbe", "lez", "tab", ]
  154.     devanagari_lang = ["hi", "mr", "ne", "bh", "mai", "ang", "bho", "mah", "sck", "new", "gom", "sa", "bgc", ]
  155.     if lang in latin_lang:
  156.         lang = "latin"
  157.     elif lang in arabic_lang:
  158.         lang = "arabic"
  159.     elif lang in cyrillic_lang:
  160.         lang = "cyrillic"
  161.     elif lang in devanagari_lang:
  162.         lang = "devanagari"
  163.     assert (lang in MODEL_URLS["OCR"][DEFAULT_OCR_MODEL_VERSION]["rec"]), "param lang must in {}, but got {}".format(MODEL_URLS["OCR"][DEFAULT_OCR_MODEL_VERSION]["rec"].keys(), lang)
  164.     if lang == "ch":
  165.         det_lang = "ch"
  166.     elif lang == "structure":
  167.         det_lang = "structure"
  168.     elif lang in ["en", "latin"]:
  169.         det_lang = "en"
  170.     else:
  171.         det_lang = "ml"
  172.     return lang, det_lang
  173.  
  174.  
  175. def get_model_config(type, version, model_type, lang):
  176.     if type == "OCR":
  177.         DEFAULT_MODEL_VERSION = DEFAULT_OCR_MODEL_VERSION
  178.     elif type == "STRUCTURE":
  179.         DEFAULT_MODEL_VERSION = DEFAULT_STRUCTURE_MODEL_VERSION
  180.     else:
  181.         raise NotImplementedError
  182.  
  183.     model_urls = MODEL_URLS[type]
  184.     if version not in model_urls:
  185.         version = DEFAULT_MODEL_VERSION
  186.     if model_type not in model_urls[version]:
  187.         if model_type in model_urls[DEFAULT_MODEL_VERSION]:
  188.             version = DEFAULT_MODEL_VERSION
  189.         else:
  190.             logger.error("{} models is not support, we only support {}".format(model_type, model_urls[DEFAULT_MODEL_VERSION].keys()))
  191.             sys.exit(-1)
  192.  
  193.     if lang not in model_urls[version][model_type]:
  194.         if lang in model_urls[DEFAULT_MODEL_VERSION][model_type]:
  195.             version = DEFAULT_MODEL_VERSION
  196.         else:
  197.             logger.error("lang {} is not support, we only support {} for {} models".format(lang, model_urls[DEFAULT_MODEL_VERSION][model_type].keys(), model_type, ))
  198.             sys.exit(-1)
  199.     return model_urls[version][model_type][lang]
  200.  
  201.  
  202. def img_decode(content: bytes):
  203.     np_arr = np.frombuffer(content, dtype=np.uint8)
  204.     return cv2.imdecode(np_arr, cv2.IMREAD_UNCHANGED)
  205.  
  206.  
  207. def check_img(img, alpha_color=(255, 255, 255)):
  208.     """
  209.     Check the image data. If it is another type of image file, try to decode it into a numpy array.
  210.     The inference network requires three-channel images, So the following channel conversions are done
  211.         single channel image: Gray to RGB R←Y,G←Y,B←Y
  212.         four channel image: alpha_to_color
  213.     args:
  214.         img: image data
  215.             file format: jpg, png and other image formats that opencv can decode, as well as gif and pdf formats
  216.             storage type: binary image, net image file, local image file
  217.         alpha_color: Background color in images in RGBA format
  218.         return: numpy.array (h, w, 3) or list (p, h, w, 3) (p: page of pdf), boolean, boolean
  219.     """
  220.     flag_gif, flag_pdf = False, False
  221.     if isinstance(img, bytes):
  222.         img = img_decode(img)
  223.     if isinstance(img, str):
  224.         # download net image
  225.         if is_link(img):
  226.             download_with_progressbar(img, "tmp.jpg")
  227.             img = "tmp.jpg"
  228.         image_file = img
  229.         img, flag_gif, flag_pdf = check_and_read(image_file)
  230.         if not flag_gif and not flag_pdf:
  231.             with open(image_file, "rb") as f:
  232.                 img_str = f.read()
  233.                 img = img_decode(img_str)
  234.             if img is None:
  235.                 try:
  236.                     buf = BytesIO()
  237.                     image = BytesIO(img_str)
  238.                     im = Image.open(image)
  239.                     rgb = im.convert("RGB")
  240.                     rgb.save(buf, "jpeg")
  241.                     buf.seek(0)
  242.                     image_bytes = buf.read()
  243.                     data_base64 = str(base64.b64encode(image_bytes), encoding="utf-8")
  244.                     image_decode = base64.b64decode(data_base64)
  245.                     img_array = np.frombuffer(image_decode, np.uint8)
  246.                     img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
  247.                 except:
  248.                     logger.error("error in loading image:{}".format(image_file))
  249.                     return None, flag_gif, flag_pdf
  250.         if img is None:
  251.             logger.error("error in loading image:{}".format(image_file))
  252.             return None, flag_gif, flag_pdf
  253.     # single channel image array.shape:h,w
  254.     if isinstance(img, np.ndarray) and len(img.shape) == 2:
  255.         img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
  256.     # four channel image array.shape:h,w,c
  257.     if isinstance(img, np.ndarray) and len(img.shape) == 3 and img.shape[2] == 4:
  258.         img = alpha_to_color(img, alpha_color)
  259.     return img, flag_gif, flag_pdf
  260.  
  261.  
  262. class PaddleOCR(predict_system.TextSystem):
  263.     def __init__(self, **kwargs):
  264.         """
  265.         paddleocr package
  266.         args:
  267.             **kwargs: other params show in paddleocr --help
  268.         """
  269.         params = parse_args(mMain=False)
  270.         params.__dict__.update(**kwargs)
  271.         assert (params.ocr_version in SUPPORT_OCR_MODEL_VERSION), "ocr_version must in {}, but get {}".format(SUPPORT_OCR_MODEL_VERSION, params.ocr_version)
  272.         params.use_gpu = check_gpu(params.use_gpu)
  273.  
  274.         if not params.show_log:
  275.             logger.setLevel(logging.INFO)
  276.         self.use_angle_cls = params.use_angle_cls
  277.         lang, det_lang = parse_lang(params.lang)
  278.  
  279.         # init model dir
  280.         det_model_config = get_model_config("OCR", params.ocr_version, "det", det_lang)
  281.         params.det_model_dir, det_url = confirm_model_dir_url(params.det_model_dir, os.path.join(BASE_DIR, "whl", "det", det_lang), det_model_config["url"], )
  282.         rec_model_config = get_model_config("OCR", params.ocr_version, "rec", lang)
  283.         params.rec_model_dir, rec_url = confirm_model_dir_url(params.rec_model_dir, os.path.join(BASE_DIR, "whl", "rec", lang), rec_model_config["url"], )
  284.         cls_model_config = get_model_config("OCR", params.ocr_version, "cls", "ch")
  285.         params.cls_model_dir, cls_url = confirm_model_dir_url(params.cls_model_dir, os.path.join(BASE_DIR, "whl", "cls"), cls_model_config["url"], )
  286.         if params.ocr_version in ["PP-OCRv3", "PP-OCRv4"]:
  287.             params.rec_image_shape = "3, 48, 320"
  288.         else:
  289.             params.rec_image_shape = "3, 32, 320"
  290.         # download model if using paddle infer
  291.         if not params.use_onnx:
  292.             maybe_download(params.det_model_dir, det_url)
  293.             maybe_download(params.rec_model_dir, rec_url)
  294.             maybe_download(params.cls_model_dir, cls_url)
  295.  
  296.         if params.det_algorithm not in SUPPORT_DET_MODEL:
  297.             logger.error("det_algorithm must in {}".format(SUPPORT_DET_MODEL))
  298.             sys.exit(0)
  299.         if params.rec_algorithm not in SUPPORT_REC_MODEL:
  300.             logger.error("rec_algorithm must in {}".format(SUPPORT_REC_MODEL))
  301.             sys.exit(0)
  302.  
  303.         if params.rec_char_dict_path is None:
  304.             params.rec_char_dict_path = str(Path(__file__).parent / rec_model_config["dict_path"])
  305.  
  306.         logger.debug(params)
  307.         # init det_model and rec_model
  308.         super().__init__(params)
  309.         self.page_num = params.page_num
  310.  
  311.     def ocr(self, img, det=True, rec=True, cls=True, bin=False, inv=False, alpha_color=(255, 255, 255), ):
  312.         """
  313.         OCR with PaddleOCR
  315.         args:
  316.             img: img for OCR, support ndarray, img_path and list or ndarray
  317.             det: use text detection or not. If False, only rec will be exec. Default is True
  318.             rec: use text recognition or not. If False, only det will be exec. Default is True
  319.             cls: use angle classifier or not. Default is True. If True, the text with rotation of 180 degrees can be recognized. If no text is rotated by 180 degrees, use cls=False to get better performance. Text with rotation of 90 or 270 degrees can be recognized even if cls=False.
  320.             bin: binarize image to black and white. Default is False.
  321.             inv: invert image colors. Default is False.
  322.             alpha_color: set RGB color Tuple for transparent parts replacement. Default is pure white.
  323.         """
  324.         assert isinstance(img, (np.ndarray, list, str, bytes))
  325.         if isinstance(img, list) and det == True:
  326.             logger.error("When input a list of images, det must be false")
  327.             exit(0)
  328.         if cls == True and self.use_angle_cls == False:
  329.             logger.warning("Since the angle classifier is not initialized, it will not be used during the forward process")
  330.  
  331.         img, flag_gif, flag_pdf = check_img(img, alpha_color)
  332.         # for infer pdf file
  333.         if isinstance(img, list) and flag_pdf:
  334.             if self.page_num > len(img) or self.page_num == 0:
  335.                 imgs = img
  336.             else:
  337.                 imgs = img[: self.page_num]
  338.         else:
  339.             imgs = [img]
  340.  
  341.         def preprocess_image(_image):
  342.             _image = alpha_to_color(_image, alpha_color)
  343.             if inv:
  344.                 _image = cv2.bitwise_not(_image)
  345.             if bin:
  346.                 _image = binarize_img(_image)
  347.             return _image
  348.  
  349.         if det and rec:
  350.             ocr_res = []
  351.             for idx, img in enumerate(imgs):
  352.                 img = preprocess_image(img)
  353.                 dt_boxes, rec_res, _ = self.__call__(img, cls)
  354.                 if not dt_boxes and not rec_res:
  355.                     ocr_res.append(None)
  356.                     continue
  357.                 tmp_res = [[box.tolist(), res] for box, res in zip(dt_boxes, rec_res)]
  358.                 ocr_res.append(tmp_res)
  359.             return ocr_res
  360.         elif det and not rec:
  361.             ocr_res = []
  362.             for idx, img in enumerate(imgs):
  363.                 img = preprocess_image(img)
  364.                 dt_boxes, elapse = self.text_detector(img)
  365.                 if dt_boxes.size == 0:
  366.                     ocr_res.append(None)
  367.                     continue
  368.                 tmp_res = [box.tolist() for box in dt_boxes]
  369.                 ocr_res.append(tmp_res)
  370.             return ocr_res
  371.         else:
  372.             ocr_res = []
  373.             cls_res = []
  374.             for idx, img in enumerate(imgs):
  375.                 if not isinstance(img, list):
  376.                     img = preprocess_image(img)
  377.                     img = [img]
  378.                 if self.use_angle_cls and cls:
  379.                     img, cls_res_tmp, elapse = self.text_classifier(img)
  380.                     if not rec:
  381.                         cls_res.append(cls_res_tmp)
  382.                 rec_res, elapse = self.text_recognizer(img)
  383.                 ocr_res.append(rec_res)
  384.             if not rec:
  385.                 return cls_res
  386.             return ocr_res
  387.  
  388.  
  389. class PPStructure(StructureSystem):
  390.     def __init__(self, **kwargs):
  391.         params = parse_args(mMain=False)
  392.         params.__dict__.update(**kwargs)
  393.         assert (params.structure_version in SUPPORT_STRUCTURE_MODEL_VERSION), "structure_version must in {}, but get {}".format(SUPPORT_STRUCTURE_MODEL_VERSION, params.structure_version)
  394.         params.use_gpu = check_gpu(params.use_gpu)
  395.         params.mode = "structure"
  396.  
  397.         if not params.show_log:
  398.             logger.setLevel(logging.INFO)
  399.         lang, det_lang = parse_lang(params.lang)
  400.         if lang == "ch":
  401.             table_lang = "ch"
  402.         else:
  403.             table_lang = "en"
  404.         if params.structure_version == "PP-Structure":
  405.             params.merge_no_span_structure = False
  406.  
  407.         # init model dir
  408.         det_model_config = get_model_config("OCR", params.ocr_version, "det", det_lang)
  409.         params.det_model_dir, det_url = confirm_model_dir_url(params.det_model_dir, os.path.join(BASE_DIR, "whl", "det", det_lang), det_model_config["url"], )
  410.         rec_model_config = get_model_config("OCR", params.ocr_version, "rec", lang)
  411.         params.rec_model_dir, rec_url = confirm_model_dir_url(params.rec_model_dir, os.path.join(BASE_DIR, "whl", "rec", lang), rec_model_config["url"], )
  412.         table_model_config = get_model_config("STRUCTURE", params.structure_version, "table", table_lang)
  413.         params.table_model_dir, table_url = confirm_model_dir_url(params.table_model_dir, os.path.join(BASE_DIR, "whl", "table"), table_model_config["url"], )
  414.         layout_model_config = get_model_config("STRUCTURE", params.structure_version, "layout", lang)
  415.         params.layout_model_dir, layout_url = confirm_model_dir_url(params.layout_model_dir, os.path.join(BASE_DIR, "whl", "layout"), layout_model_config["url"], )
  416.         # download model
  417.         if not params.use_onnx:
  418.             maybe_download(params.det_model_dir, det_url)
  419.             maybe_download(params.rec_model_dir, rec_url)
  420.             maybe_download(params.table_model_dir, table_url)
  421.             maybe_download(params.layout_model_dir, layout_url)
  422.  
  423.         if params.rec_char_dict_path is None:
  424.             params.rec_char_dict_path = str(Path(__file__).parent / rec_model_config["dict_path"])
  425.         if params.table_char_dict_path is None:
  426.             params.table_char_dict_path = str(Path(__file__).parent / table_model_config["dict_path"])
  427.         if params.layout_dict_path is None:
  428.             params.layout_dict_path = str(Path(__file__).parent / layout_model_config["dict_path"])
  429.         logger.debug(params)
  430.         super().__init__(params)
  431.  
  432.     def __call__(self, img, return_ocr_result_in_table=False, img_idx=0, alpha_color=(255, 255, 255), ):
  433.         img, flag_gif, flag_pdf = check_img(img, alpha_color)
  434.         if isinstance(img, list) and flag_pdf:
  435.             res_list = []
  436.             for index, pdf_img in enumerate(img):
  437.                 logger.info("processing {}/{} page:".format(index + 1, len(img)))
  438.                 res, _ = super().__call__(pdf_img, return_ocr_result_in_table, img_idx=index)
  439.                 res_list.append(res)
  440.             return res_list
  441.         res, _ = super().__call__(img, return_ocr_result_in_table, img_idx=img_idx)
  442.         return res
  443. def cv2AddChineseText(img, text, position, textColor=(0, 255, 0), textSize=30):
  444.     img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
  445.     draw = ImageDraw.Draw(img)
  446.     draw.text(position, text, textColor, font=font)
  447.     return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
  448.  
  449. if __name__ == '__main__':
  450.  
  451.     font_path = 'simhei.ttf'  # 需要替换为你的中文字体路径
  452.     font = ImageFont.truetype(font_path, 24)
  453.  
  454.  
  455.     # for cmd
  456.     args = parse_args(mMain=True)
  457.     image_dir = args.image_dir
  458.     image_file_list=['weights/123.jpg']
  459.     if args.type == "ocr":
  460.         engine = PaddleOCR(**(args.__dict__))
  461.     elif args.type == "structure":
  462.         engine = PPStructure(**(args.__dict__))
  463.     else:
  464.         raise NotImplementedError
  465.  
  466.     for img_path in image_file_list:
  467.         img_name = os.path.basename(img_path).split(".")[0]
  468.         logger.info("{}{}{}".format("*" * 10, img_path, "*" * 10))
  469.         if args.type == "ocr":
  470.             image=cv2.imread(img_path)
  471.             result = engine.ocr(img_path, det=args.det, rec=args.rec, cls=args.use_angle_cls, bin=args.binarize, inv=args.invert, alpha_color=args.alphacolor, )
  472.             if result is not None:
  473.                 lines = []
  474.                 for idx in range(len(result)):
  475.                     res = result[idx]
  476.                     for line in res:
  477.  
  478.                         points = line[0]
  479.                         text = line[1][0]
  480.                         points = np.array(points, dtype=np.int32).reshape((-1, 1, 2))
  481.                         cv2.polylines(image, [points], isClosed=True, color=(255, 0, 0), thickness=2)
  482.                         text_position = (int(points[0][0][0]), int(points[0][0][1] + 20))  # 微调文本位置
  483.  
  484.                         # cv2.putText(image, '中文文本', (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 3)
  485.                         image = cv2AddChineseText(image, text, text_position, textColor=(0, 255, 0), textSize=30)
  486.  
  487.                         logger.info(line)
  488.                         val = "["
  489.                         for box in line[0]:
  490.                             val += str(box[0]) + "," + str(box[1]) + ","
  491.  
  492.                         val = val[:-1]
  493.                         val += "]," + line[1][0] + "," + str(line[1][1]) + "\n"
  494.                         lines.append(val)
  495.                 if args.savefile:
  496.                     if os.path.exists(args.output) is False:
  497.                         os.mkdir(args.output)
  498.                     outfile = args.output + "/" + img_name + ".txt"
  499.                     with open(outfile, "w", encoding="utf-8") as f:
  500.                         f.writelines(lines)
  501.  
  502.         elif args.type == "structure":
  503.             img, flag_gif, flag_pdf = check_and_read(img_path)
  504.             if not flag_gif and not flag_pdf:
  505.                 img = cv2.imread(img_path)
  506.  
  507.             if not flag_pdf:
  508.                 if img is None:
  509.                     logger.error("error in loading image:{}".format(img_path))
  510.                     continue
  511.                 img_paths = [[img_path, img]]
  512.             else:
  513.                 img_paths = []
  514.                 for index, pdf_img in enumerate(img):
  515.                     os.makedirs(os.path.join(args.output, img_name), exist_ok=True)
  516.                     pdf_img_path = os.path.join(args.output, img_name, img_name + "_" + str(index) + ".jpg")
  517.                     cv2.imwrite(pdf_img_path, pdf_img)
  518.                     img_paths.append([pdf_img_path, pdf_img])
  519.  
  520.             all_res = []
  521.             for index, (new_img_path, img) in enumerate(img_paths):
  522.                 logger.info("processing {}/{} page:".format(index + 1, len(img_paths)))
  523.                 new_img_name = os.path.basename(new_img_path).split(".")[0]
  524.                 result = engine(img, img_idx=index)
  525.                 save_structure_res(result, args.output, img_name, index)
  526.  
  527.                 if args.recovery and result != []:
  528.                     from copy import deepcopy
  529.                     from ppstructure.recovery.recovery_to_doc import sorted_layout_boxes
  530.  
  531.                     h, w, _ = img.shape
  532.                     result_cp = deepcopy(result)
  533.                     result_sorted = sorted_layout_boxes(result_cp, w)
  534.                     all_res += result_sorted
  535.  
  536.             if args.recovery and all_res != []:
  537.                 try:
  538.                     from ppstructure.recovery.recovery_to_doc import convert_info_docx
  539.  
  540.                     convert_info_docx(img, all_res, args.output, img_name)
  541.                 except Exception as ex:
  542.                     logger.error("error in layout recovery image:{}, err msg: {}".format(img_name, ex))
  543.                     continue
  544.  
  545.             for item in all_res:
  546.                 item.pop("img")
  547.                 item.pop("res")
  548.                 logger.info(item)
  549.             logger.info("result save to {}".format(args.output))
  550.  
  551.         cv2.imshow('image', image)
  552.         cv2.waitKey(0)

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/我家自动化/article/detail/593603
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号