OCR/hmOCR/core.py

import logging
from .parts import *
from .utility import *
from time import sleep
from numpy import ndarray
from copy import deepcopy
from .argument import ArgType
from threading import Thread, Lock

__all__ = ["Engine", "HuiMvOCR"]
logger = logging.getLogger("hm-ocr")
logger.setLevel(logging.INFO)


class Engine:
    __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 __call__(self, img, cls: "bool" = True, 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


class HuiMvOCR:
    __lock = Lock()
    __tasks = []  # item: [img: "ndarray", ocr_args: "dict", callback: "fn", callback_args: "dict"]

    def __init__(self, args: "ArgType"):
        self.interval = args.interval

        for i in range(args.workers):
            Thread(target=self.__processor, args=(Engine(args), i), daemon=True).start()

    @staticmethod
    def __processor(ocr: "Engine", eid: "int"):
        logger.info(f"================ Engine[{eid}] initialized ================")
        while True:
            if HuiMvOCR.__tasks:
                HuiMvOCR.__lock.acquire()
                img, ocr_args, callback, callback_args = HuiMvOCR.__tasks.pop(0)
                HuiMvOCR.__lock.release()
                res = ocr(img)
                callback(res, **callback_args)
            sleep(0.1)

    def rec_one(self, img: "ndarray", cls: "bool" = True, use_space: "bool" = True):
        def callback(res):
            foo[1] = res
            foo[0] = 1

        foo = [0, None]  # finish_count, result
        args = {"cls": cls, "use_space": use_space}
        HuiMvOCR.__tasks.append([img, args, callback, {}])
        while foo[0] < 1:
            sleep(self.interval)
        return foo[1]

    def rec_multi(self, images: "list[ndarray]", cls: "bool" = False, use_space: "bool" = False):
        def callback(res, index):
            foo[1][index] = res
            foo[0] += 1

        size, args = len(images), {"cls": cls, "use_space": use_space}
        foo = [0, [...] * size]  # finish_count, result
        for i in range(size):
            HuiMvOCR.__tasks.append([images[i], args, callback, {"index": i}])
        while foo[0] < size:
            sleep(self.interval)
        return foo[1]