301 lines
9.0 KiB
Python
301 lines
9.0 KiB
Python
|
# Copyright (c) Opendatalab. All rights reserved.
|
|||
|
|
import copy
|
|||
|
|
import os
|
|||
|
|
import warnings
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
import yaml
|
|||
|
|
from loguru import logger
|
|||
|
|
|
|||
|
|
from mineru.utils.config_reader import get_device
|
|||
|
|
from mineru.utils.enum_class import ModelPath
|
|||
|
|
from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
|
|||
|
|
from mineru.utils.ocr_utils import check_img, preprocess_image, sorted_boxes, merge_det_boxes, update_det_boxes, get_rotate_crop_image
|
|||
|
|
from mineru.model.utils.tools.infer.predict_system import TextSystem
|
|||
|
|
from mineru.model.utils.tools.infer import pytorchocr_utility as utility
|
|||
|
|
import argparse
|
|||
|
|
|
|||
|
|
|
|||
|
|
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",
|
|||
|
|
"fi",
|
|||
|
|
"eu",
|
|||
|
|
"gl",
|
|||
|
|
"lb",
|
|||
|
|
"rm",
|
|||
|
|
"ca",
|
|||
|
|
"qu",
|
|||
|
|
]
|
|||
|
|
arabic_lang = ["ar", "fa", "ug", "ur", "ps", "ku", "sd", "bal"]
|
|||
|
|
cyrillic_lang = [
|
|||
|
|
"ru",
|
|||
|
|
"rs_cyrillic",
|
|||
|
|
"be",
|
|||
|
|
"bg",
|
|||
|
|
"uk",
|
|||
|
|
"mn",
|
|||
|
|
"abq",
|
|||
|
|
"ady",
|
|||
|
|
"kbd",
|
|||
|
|
"ava",
|
|||
|
|
"dar",
|
|||
|
|
"inh",
|
|||
|
|
"che",
|
|||
|
|
"lbe",
|
|||
|
|
"lez",
|
|||
|
|
"tab",
|
|||
|
|
"kk",
|
|||
|
|
"ky",
|
|||
|
|
"tg",
|
|||
|
|
"mk",
|
|||
|
|
"tt",
|
|||
|
|
"cv",
|
|||
|
|
"ba",
|
|||
|
|
"mhr",
|
|||
|
|
"mo",
|
|||
|
|
"udm",
|
|||
|
|
"kv",
|
|||
|
|
"os",
|
|||
|
|
"bua",
|
|||
|
|
"xal",
|
|||
|
|
"tyv",
|
|||
|
|
"sah",
|
|||
|
|
"kaa",
|
|||
|
|
]
|
|||
|
|
east_slavic_lang = ["ru", "be", "uk"]
|
|||
|
|
devanagari_lang = [
|
|||
|
|
"hi",
|
|||
|
|
"mr",
|
|||
|
|
"ne",
|
|||
|
|
"bh",
|
|||
|
|
"mai",
|
|||
|
|
"ang",
|
|||
|
|
"bho",
|
|||
|
|
"mah",
|
|||
|
|
"sck",
|
|||
|
|
"new",
|
|||
|
|
"gom",
|
|||
|
|
"sa",
|
|||
|
|
"bgc",
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_model_params(lang, config):
|
|||
|
|
if lang in config['lang']:
|
|||
|
|
params = config['lang'][lang]
|
|||
|
|
det = params.get('det')
|
|||
|
|
rec = params.get('rec')
|
|||
|
|
dict_file = params.get('dict')
|
|||
|
|
return det, rec, dict_file
|
|||
|
|
else:
|
|||
|
|
raise Exception (f'Language {lang} not supported')
|
|||
|
|
|
|||
|
|
|
|||
|
|
root_dir = os.path.join(Path(__file__).resolve().parent.parent, 'utils')
|
|||
|
|
|
|||
|
|
|
|||
|
|
class PytorchPaddleOCR(TextSystem):
|
|||
|
|
def __init__(self, *args, **kwargs):
|
|||
|
|
parser = utility.init_args()
|
|||
|
|
args = parser.parse_args(args)
|
|||
|
|
|
|||
|
|
self.lang = kwargs.get('lang', 'ch')
|
|||
|
|
self.enable_merge_det_boxes = kwargs.get("enable_merge_det_boxes", True)
|
|||
|
|
|
|||
|
|
device = get_device()
|
|||
|
|
if device == 'cpu' and self.lang in ['ch', 'ch_server', 'japan', 'chinese_cht']:
|
|||
|
|
# logger.warning("The current device in use is CPU. To ensure the speed of parsing, the language is automatically switched to ch_lite.")
|
|||
|
|
self.lang = 'ch_lite'
|
|||
|
|
|
|||
|
|
if self.lang in latin_lang:
|
|||
|
|
self.lang = 'latin'
|
|||
|
|
elif self.lang in east_slavic_lang:
|
|||
|
|
self.lang = 'east_slavic'
|
|||
|
|
elif self.lang in arabic_lang:
|
|||
|
|
self.lang = 'arabic'
|
|||
|
|
elif self.lang in cyrillic_lang:
|
|||
|
|
self.lang = 'cyrillic'
|
|||
|
|
elif self.lang in devanagari_lang:
|
|||
|
|
self.lang = 'devanagari'
|
|||
|
|
else:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
models_config_path = os.path.join(root_dir, 'pytorchocr', 'utils', 'resources', 'models_config.yml')
|
|||
|
|
with open(models_config_path) as file:
|
|||
|
|
config = yaml.safe_load(file)
|
|||
|
|
det, rec, dict_file = get_model_params(self.lang, config)
|
|||
|
|
ocr_models_dir = ModelPath.pytorch_paddle
|
|||
|
|
|
|||
|
|
det_model_path = f"{ocr_models_dir}/{det}"
|
|||
|
|
det_model_path = os.path.join(auto_download_and_get_model_root_path(det_model_path), det_model_path)
|
|||
|
|
rec_model_path = f"{ocr_models_dir}/{rec}"
|
|||
|
|
rec_model_path = os.path.join(auto_download_and_get_model_root_path(rec_model_path), rec_model_path)
|
|||
|
|
kwargs['det_model_path'] = det_model_path
|
|||
|
|
kwargs['rec_model_path'] = rec_model_path
|
|||
|
|
kwargs['rec_char_dict_path'] = os.path.join(root_dir, 'pytorchocr', 'utils', 'resources', 'dict', dict_file)
|
|||
|
|
kwargs['rec_batch_num'] = 6
|
|||
|
|
|
|||
|
|
kwargs['device'] = device
|
|||
|
|
|
|||
|
|
default_args = vars(args)
|
|||
|
|
default_args.update(kwargs)
|
|||
|
|
args = argparse.Namespace(**default_args)
|
|||
|
|
|
|||
|
|
super().__init__(args)
|
|||
|
|
|
|||
|
|
def ocr(self,
|
|||
|
|
img,
|
|||
|
|
det=True,
|
|||
|
|
rec=True,
|
|||
|
|
mfd_res=None,
|
|||
|
|
tqdm_enable=False,
|
|||
|
|
tqdm_desc="OCR-rec Predict",
|
|||
|
|
):
|
|||
|
|
assert isinstance(img, (np.ndarray, list, str, bytes))
|
|||
|
|
if isinstance(img, list) and det == True:
|
|||
|
|
logger.error('When input a list of images, det must be false')
|
|||
|
|
exit(0)
|
|||
|
|
img = check_img(img)
|
|||
|
|
imgs = [img]
|
|||
|
|
with warnings.catch_warnings():
|
|||
|
|
warnings.simplefilter("ignore", category=RuntimeWarning)
|
|||
|
|
if det and rec:
|
|||
|
|
ocr_res = []
|
|||
|
|
for img in imgs:
|
|||
|
|
img = preprocess_image(img)
|
|||
|
|
dt_boxes, rec_res = self.__call__(img, mfd_res=mfd_res)
|
|||
|
|
if not dt_boxes and not rec_res:
|
|||
|
|
ocr_res.append(None)
|
|||
|
|
continue
|
|||
|
|
tmp_res = [[box.tolist(), res] for box, res in zip(dt_boxes, rec_res)]
|
|||
|
|
ocr_res.append(tmp_res)
|
|||
|
|
return ocr_res
|
|||
|
|
elif det and not rec:
|
|||
|
|
ocr_res = []
|
|||
|
|
for img in imgs:
|
|||
|
|
img = preprocess_image(img)
|
|||
|
|
dt_boxes, elapse = self.text_detector(img)
|
|||
|
|
# logger.debug("dt_boxes num : {}, elapsed : {}".format(len(dt_boxes), elapse))
|
|||
|
|
if dt_boxes is None:
|
|||
|
|
ocr_res.append(None)
|
|||
|
|
continue
|
|||
|
|
dt_boxes = sorted_boxes(dt_boxes)
|
|||
|
|
# merge_det_boxes 和 update_det_boxes 都会把poly转成bbox再转回poly，因此需要过滤所有倾斜程度较大的文本框
|
|||
|
|
if self.enable_merge_det_boxes:
|
|||
|
|
dt_boxes = merge_det_boxes(dt_boxes)
|
|||
|
|
if mfd_res:
|
|||
|
|
dt_boxes = update_det_boxes(dt_boxes, mfd_res)
|
|||
|
|
tmp_res = [box.tolist() for box in dt_boxes]
|
|||
|
|
ocr_res.append(tmp_res)
|
|||
|
|
return ocr_res
|
|||
|
|
elif not det and rec:
|
|||
|
|
ocr_res = []
|
|||
|
|
for img in imgs:
|
|||
|
|
if not isinstance(img, list):
|
|||
|
|
img = preprocess_image(img)
|
|||
|
|
img = [img]
|
|||
|
|
rec_res, elapse = self.text_recognizer(img, tqdm_enable=tqdm_enable, tqdm_desc=tqdm_desc)
|
|||
|
|
# logger.debug("rec_res num : {}, elapsed : {}".format(len(rec_res), elapse))
|
|||
|
|
ocr_res.append(rec_res)
|
|||
|
|
return ocr_res
|
|||
|
|
|
|||
|
|
def __call__(self, img, mfd_res=None):
|
|||
|
|
|
|||
|
|
if img is None:
|
|||
|
|
logger.debug("no valid image provided")
|
|||
|
|
return None, None
|
|||
|
|
|
|||
|
|
ori_im = img.copy()
|
|||
|
|
dt_boxes, elapse = self.text_detector(img)
|
|||
|
|
|
|||
|
|
if dt_boxes is None:
|
|||
|
|
logger.debug("no dt_boxes found, elapsed : {}".format(elapse))
|
|||
|
|
return None, None
|
|||
|
|
else:
|
|||
|
|
pass
|
|||
|
|
# logger.debug("dt_boxes num : {}, elapsed : {}".format(len(dt_boxes), elapse))
|
|||
|
|
img_crop_list = []
|
|||
|
|
|
|||
|
|
dt_boxes = sorted_boxes(dt_boxes)
|
|||
|
|
|
|||
|
|
# merge_det_boxes 和 update_det_boxes 都会把poly转成bbox再转回poly，因此需要过滤所有倾斜程度较大的文本框
|
|||
|
|
if self.enable_merge_det_boxes:
|
|||
|
|
dt_boxes = merge_det_boxes(dt_boxes)
|
|||
|
|
|
|||
|
|
if mfd_res:
|
|||
|
|
dt_boxes = update_det_boxes(dt_boxes, mfd_res)
|
|||
|
|
|
|||
|
|
for bno in range(len(dt_boxes)):
|
|||
|
|
tmp_box = copy.deepcopy(dt_boxes[bno])
|
|||
|
|
img_crop = get_rotate_crop_image(ori_im, tmp_box)
|
|||
|
|
img_crop_list.append(img_crop)
|
|||
|
|
|
|||
|
|
rec_res, elapse = self.text_recognizer(img_crop_list)
|
|||
|
|
# logger.debug("rec_res num : {}, elapsed : {}".format(len(rec_res), elapse))
|
|||
|
|
|
|||
|
|
filter_boxes, filter_rec_res = [], []
|
|||
|
|
for box, rec_result in zip(dt_boxes, rec_res):
|
|||
|
|
text, score = rec_result
|
|||
|
|
if score >= self.drop_score:
|
|||
|
|
filter_boxes.append(box)
|
|||
|
|
filter_rec_res.append(rec_result)
|
|||
|
|
|
|||
|
|
return filter_boxes, filter_rec_res
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
pytorch_paddle_ocr = PytorchPaddleOCR()
|
|||
|
|
img = cv2.imread("/Users/myhloli/Downloads/screenshot-20250326-194348.png")
|
|||
|
|
dt_boxes, rec_res = pytorch_paddle_ocr(img)
|
|||
|
|
ocr_res = []
|
|||
|
|
if not dt_boxes and not rec_res:
|
|||
|
|
ocr_res.append(None)
|
|||
|
|
else:
|
|||
|
|
tmp_res = [[box.tolist(), res] for box, res in zip(dt_boxes, rec_res)]
|
|||
|
|
ocr_res.append(tmp_res)
|
|||
|
|
print(ocr_res)
|
|||
|
|
|
|||
|
|
|