| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556 |
- from .parts import *
- from .utility import *
- from copy import deepcopy
- from .argument import ArgType
- from concurrent.futures import ThreadPoolExecutor
- __all__ = ["HuiMvOcr"]
- class HuiMvOcr:
- __worker_count = 1
- def __init__(self, args: "ArgType"):
- self.det = Detector(args)
- self.cls = Classifier(args)
- self.rec = Recognizer(args)
- self.det_box_type = args.det_box_type
- self.drop_score = args.drop_score
- self.crop_image_res_index = 0
- def ocr_one(self, img, cls: "bool" = False, use_space: "bool" = True):
- ori_im = img.copy()
- dt_boxes = self.det(img)
- if dt_boxes is None:
- return None
- dt_boxes = sorted_boxes(dt_boxes)
- size = len(dt_boxes)
- img_crop_list = [...] * size
- for i in range(size):
- box = deepcopy(dt_boxes[i])
- if self.det_box_type == "quad":
- img_crop_list[i] = get_rotate_crop_image(ori_im, box)
- else:
- img_crop_list[i] = get_min_area_rect_crop(ori_im, box)
- if cls:
- img_crop_list = self.cls(img_crop_list)
- rec_res = self.rec(img_crop_list, use_space=use_space)
- filter_rec_res = []
- for text, score in rec_res:
- if score > self.drop_score:
- filter_rec_res.append((text, score))
- return filter_rec_res
- def ocr_multi(self, img_list, cls: "bool" = False, use_space: "bool" = True):
- pool = ThreadPoolExecutor(HuiMvOcr.__worker_count)
- loop = range(len(img_list))
- tasks = [pool.submit(self.ocr_one, img_list[i], cls, use_space) for i in loop]
- return [tasks[i].result() for i in loop]
|
