diff --git a/module/ocr/keyword.py b/module/ocr/keyword.py index 5b8ad91..2452ab4 100644 --- a/module/ocr/keyword.py +++ b/module/ocr/keyword.py @@ -201,3 +201,18 @@ class Keyword: # Not found raise ScriptError(f'Cannot find a {cls.__name__} instance that matches "{name}"') + + +class KeywordDigitCounter(Keyword): + """ + A fake Keyword class to filter digit counters in ocr results + OcrResultButton.match_keyword will be a str + """ + + @classmethod + def find(cls, name, lang: str = None, ignore_punctuation=True): + from module.ocr.ocr import DigitCounter + if DigitCounter.is_format_matched(name): + return name + else: + raise ScriptError diff --git a/module/ocr/ocr.py b/module/ocr/ocr.py index a099d26..82fb32d 100644 --- a/module/ocr/ocr.py +++ b/module/ocr/ocr.py @@ -1,7 +1,6 @@ import re import time from datetime import timedelta -from typing import Optional import cv2 import numpy as np @@ -13,13 +12,12 @@ from module.base.decorator import cached_property from module.base.utils import area_pad, corner2area, crop, extract_white_letters, float2str from module.exception import ScriptError from module.logger import logger -from module.ocr.keyword import Keyword from module.ocr.models import OCR_MODEL, TextSystem from module.ocr.utils import merge_buttons class OcrResultButton: - def __init__(self, boxed_result: BoxedResult, matched_keyword: Optional[Keyword]): + def __init__(self, boxed_result: BoxedResult, matched_keyword): """ Args: boxed_result: BoxedResult from ppocr-onnx diff --git a/module/ocr/utils.py b/module/ocr/utils.py index 357836f..d7a1584 100644 --- a/module/ocr/utils.py +++ b/module/ocr/utils.py @@ -74,8 +74,8 @@ def pair_buttons(group1, group2, relative_area): Pair buttons in group1 with those in group2 in the relative_area. Args: - group1 (list[OcrResultButton]): - group2 (list[OcrResultButton]): + group1 (list[OcrResultButton], Iterable[OcrResultButton]): + group2 (list[OcrResultButton], Iterable[OcrResultButton]): relative_area (tuple[int, int, int, int]): Yields: