尝试统一NumberRecognizer(未替换各处number)

mower/utils/number.py

@@ -0,0 +1,73 @@
+from typing import Literal
+
+import cv2
+
+from mower.models import Digtal
+
+from . import typealias as tp
+from .image import cropimg
+
+
+class NumberRecognizer:
+    def segment(self, img: tp.GrayImage) -> list[tuple[int, int, int, int]]:
+        """
+        传入二值化的图像，分割图像中的连通区域
+        """
+        contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        rect = [cv2.boundingRect(c) for c in contours]
+        rect.sort(key=lambda c: c[0])  # 按照x坐标排序
+        return rect
+
+    def filter_rectangles(
+        self,
+        rect: list[tuple[int, int, int, int]],
+        st: int | None = None,
+        ed: int | None = None,
+        limits: list[dict] = [],
+    ) -> list:
+        """
+        过滤并替换不满足条件的矩形
+        limits: **大的数字写前面**,如[{"w":20,"h":10,"char":"."},{"w":10,"h":3,"char":""}]则将w<10,h<3的矩形替换为"."，w<20,h<10的矩形替换为""
+        """
+        rect = rect[st:ed]
+        if limits:
+            for i in range(len(rect)):
+                x, y, w, h = rect[i]
+                for lim in limits:
+                    if w < lim["w"] and h < lim["h"]:
+                        rect[i] = lim["char"]
+                        break
+        return rect
+
+    def number_match(
+        self,
+        img,
+        rect: list,
+        font: Literal["riic_base", "noto", "secret_front"] = "noto",
+        target_range: list = range(10),
+    ) -> str:
+        if font == "riic_base":
+            templates = Digtal().riic_base_digits
+        elif font == "noto":
+            templates = Digtal().noto_sans
+        elif font == "secret_front":
+            templates = Digtal().secret_front
+        value = ""
+        for i in range(len(rect)):
+            if len(rect[i]) != 4:
+                value += rect[i]
+                continue
+            x, y, w, h = rect[i]
+            digit = cropimg(img, ((x, y), (x + w, y + h)))
+            digit = cv2.copyMakeBorder(
+                digit, 10, 10, 10, 10, cv2.BORDER_CONSTANT, None, (0,)
+            )
+            score = []
+            for i in target_range:
+                im = templates[i]
+                result = cv2.matchTemplate(digit, im, cv2.TM_SQDIFF_NORMED)
+                min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
+                score.append(min_val)
+            value += str(score.index(min(score)))
+
+        return value