129 lines
4.2 KiB
Python
129 lines
4.2 KiB
Python
from typing import Literal
|
|
|
|
import cv2
|
|
|
|
from mower.models import Digtal
|
|
from mower.utils import config
|
|
from mower.utils import typealias as tp
|
|
from mower.utils.image import cropimg, thres2
|
|
from mower.utils.log import logger
|
|
|
|
|
|
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
|
|
|
|
def number(
|
|
self,
|
|
font: Literal["riic_base", "noto", "secret_front"] = "noto",
|
|
scope: tp.Scope | None = None,
|
|
height: int = 25,
|
|
thres: int = 127,
|
|
img: tp.GrayImage | None = None,
|
|
rect_st: int | None = None,
|
|
rect_ed: int | None = None,
|
|
rect_limits: list[dict] = [],
|
|
target_range: list = range(10),
|
|
res_if_error: str = "",
|
|
) -> str:
|
|
try:
|
|
if img is None:
|
|
img = cropimg(config.recog.gray, scope)
|
|
default_height = 25
|
|
if height != default_height:
|
|
scale = 25 / height
|
|
img = cv2.resize(img, None, None, scale, scale)
|
|
img = thres2(img, thres)
|
|
rect = self.segment(img)
|
|
rect = self.filter_rectangles(rect, rect_st, rect_ed, rect_limits)
|
|
return self.number_match(img, rect, font, target_range)
|
|
except Exception as e:
|
|
logger.exception(e)
|
|
return res_if_error
|
|
|
|
def number_int(
|
|
self,
|
|
font: Literal["riic_base", "noto", "secret_front"] = "noto",
|
|
scope: tp.Scope | None = None,
|
|
height: int = 25,
|
|
thres: int = 127,
|
|
img: tp.GrayImage | None = None,
|
|
rect_st: int | None = None,
|
|
rect_ed: int | None = None,
|
|
rect_limits: list[dict] = [],
|
|
target_range: list = range(10),
|
|
res_if_error: str = "0",
|
|
) -> int:
|
|
res = self.number(
|
|
font,
|
|
scope,
|
|
height,
|
|
thres,
|
|
img,
|
|
rect_st,
|
|
rect_ed,
|
|
rect_limits,
|
|
target_range,
|
|
res_if_error,
|
|
)
|
|
return int(res) if res else int(res_if_error)
|