From 3ebb058cdfb0f99bfde464cdc863edb9c38ae4fc Mon Sep 17 00:00:00 2001 From: 283375 Date: Sat, 21 Jun 2025 16:06:49 +0800 Subject: [PATCH] refactor: `XYWHRect` and b30 ocr --- src/arcaea_offline_ocr/b30/chieri/v4/ocr.py | 36 +++++++++++------- src/arcaea_offline_ocr/b30/chieri/v4/rois.py | 40 ++++++++++---------- src/arcaea_offline_ocr/types.py | 35 +++++++++++------ src/arcaea_offline_ocr/utils.py | 21 ++++------ 4 files changed, 72 insertions(+), 60 deletions(-) diff --git a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py index 845145a..793adff 100644 --- a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py +++ b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py @@ -1,4 +1,3 @@ -from math import floor from typing import List, Optional, Tuple import cv2 @@ -13,9 +12,21 @@ from ....ocr import ( ) from ....phash_db import ImagePhashDatabase from ....types import Mat -from ....utils import construct_int_xywh_rect from ...shared import B30OcrResultItem -from .colors import * +from .colors import ( + BYD_MAX_HSV, + BYD_MIN_HSV, + FAR_BG_MAX_HSV, + FAR_BG_MIN_HSV, + FTR_MAX_HSV, + FTR_MIN_HSV, + LOST_BG_MAX_HSV, + LOST_BG_MIN_HSV, + PRS_MAX_HSV, + PRS_MIN_HSV, + PURE_BG_MAX_HSV, + PURE_BG_MIN_HSV, +) from .rois import ChieriBotV4Rois @@ -25,7 +36,7 @@ class ChieriBotV4Ocr: score_knn: cv2.ml.KNearest, pfl_knn: cv2.ml.KNearest, phash_db: ImagePhashDatabase, - factor: Optional[float] = 1.0, + factor: float = 1.0, ): self.__score_knn = score_knn self.__pfl_knn = pfl_knn @@ -72,9 +83,8 @@ class ChieriBotV4Ocr: self.factor = img.shape[0] / 4400 def ocr_component_rating_class(self, component_bgr: Mat) -> int: - rating_class_rect = construct_int_xywh_rect( - self.rois.component_rois.rating_class_rect - ) + rating_class_rect = self.rois.component_rois.rating_class_rect.rounded() + rating_class_roi = crop_xywh(component_bgr, rating_class_rect) rating_class_roi = cv2.cvtColor(rating_class_roi, cv2.COLOR_BGR2HSV) rating_class_masks = [ @@ -89,9 +99,7 @@ class ChieriBotV4Ocr: return max(enumerate(rating_class_results), key=lambda i: i[1])[0] + 1 def ocr_component_song_id(self, component_bgr: Mat): - jacket_rect = construct_int_xywh_rect( - self.rois.component_rois.jacket_rect, floor - ) + jacket_rect = self.rois.component_rois.jacket_rect.floored() jacket_roi = cv2.cvtColor( crop_xywh(component_bgr, jacket_rect), cv2.COLOR_BGR2GRAY ) @@ -99,7 +107,7 @@ class ChieriBotV4Ocr: def ocr_component_score_knn(self, component_bgr: Mat) -> int: # sourcery skip: inline-immediately-returned-variable - score_rect = construct_int_xywh_rect(self.rois.component_rois.score_rect) + score_rect = self.rois.component_rois.score_rect.rounded() score_roi = cv2.cvtColor( crop_xywh(component_bgr, score_rect), cv2.COLOR_BGR2GRAY ) @@ -119,7 +127,9 @@ class ChieriBotV4Ocr: score_roi = cv2.fillPoly(score_roi, [contour], 0) return ocr_digits_by_contour_knn(score_roi, self.score_knn) - def find_pfl_rects(self, component_pfl_processed: Mat) -> List[List[int]]: + def find_pfl_rects( + self, component_pfl_processed: Mat + ) -> List[Tuple[int, int, int, int]]: # sourcery skip: inline-immediately-returned-variable pfl_roi_find = cv2.morphologyEx( component_pfl_processed, @@ -146,7 +156,7 @@ class ChieriBotV4Ocr: return pfl_rects_adjusted def preprocess_component_pfl(self, component_bgr: Mat) -> Mat: - pfl_rect = construct_int_xywh_rect(self.rois.component_rois.pfl_rect) + pfl_rect = self.rois.component_rois.pfl_rect.rounded() pfl_roi = crop_xywh(component_bgr, pfl_rect) pfl_roi_hsv = cv2.cvtColor(pfl_roi, cv2.COLOR_BGR2HSV) diff --git a/src/arcaea_offline_ocr/b30/chieri/v4/rois.py b/src/arcaea_offline_ocr/b30/chieri/v4/rois.py index 9926b8a..2c6ce1d 100644 --- a/src/arcaea_offline_ocr/b30/chieri/v4/rois.py +++ b/src/arcaea_offline_ocr/b30/chieri/v4/rois.py @@ -1,12 +1,12 @@ -from typing import List, Optional +from typing import List from ....crop import crop_xywh from ....types import Mat, XYWHRect -from ....utils import apply_factor, construct_int_xywh_rect +from ....utils import apply_factor class ChieriBotV4ComponentRois: - def __init__(self, factor: Optional[float] = 1.0): + def __init__(self, factor: float = 1.0): self.__factor = factor @property @@ -19,11 +19,11 @@ class ChieriBotV4ComponentRois: @property def top_font_color_detect(self): - return apply_factor((35, 10, 120, 100), self.factor) + return apply_factor(XYWHRect(35, 10, 120, 100), self.factor) @property def bottom_font_color_detect(self): - return apply_factor((30, 125, 175, 110), self.factor) + return apply_factor(XYWHRect(30, 125, 175, 110), self.factor) @property def bg_point(self): @@ -31,31 +31,31 @@ class ChieriBotV4ComponentRois: @property def rating_class_rect(self): - return apply_factor((21, 40, 7, 20), self.factor) + return apply_factor(XYWHRect(21, 40, 7, 20), self.factor) @property def title_rect(self): - return apply_factor((35, 10, 430, 50), self.factor) + return apply_factor(XYWHRect(35, 10, 430, 50), self.factor) @property def jacket_rect(self): - return apply_factor((263, 0, 239, 239), self.factor) + return apply_factor(XYWHRect(263, 0, 239, 239), self.factor) @property def score_rect(self): - return apply_factor((30, 60, 270, 55), self.factor) + return apply_factor(XYWHRect(30, 60, 270, 55), self.factor) @property def pfl_rect(self): - return apply_factor((50, 125, 80, 100), self.factor) + return apply_factor(XYWHRect(50, 125, 80, 100), self.factor) @property def date_rect(self): - return apply_factor((205, 200, 225, 25), self.factor) + return apply_factor(XYWHRect(205, 200, 225, 25), self.factor) class ChieriBotV4Rois: - def __init__(self, factor: Optional[float] = 1.0): + def __init__(self, factor: float = 1.0): self.__factor = factor self.__component_rois = ChieriBotV4ComponentRois(factor) @@ -100,9 +100,7 @@ class ChieriBotV4Rois: def horizontal_items(self): return 3 - @property - def vertical_items(self): - return 10 + vertical_items = 10 @property def b33_vertical_gap(self): @@ -112,16 +110,17 @@ class ChieriBotV4Rois: first_rect = XYWHRect(x=self.left, y=self.top, w=self.width, h=self.height) results = [] + last_rect = first_rect for vi in range(self.vertical_items): rect = XYWHRect(*first_rect) rect += (0, (self.vertical_gap + self.height) * vi, 0, 0) for hi in range(self.horizontal_items): if hi > 0: rect += ((self.width + self.horizontal_gap), 0, 0, 0) - int_rect = construct_int_xywh_rect(rect) - results.append(crop_xywh(img_bgr, int_rect)) + results.append(crop_xywh(img_bgr, rect.rounded())) + last_rect = rect - rect += ( + last_rect += ( -(self.width + self.horizontal_gap) * 2, self.height + self.b33_vertical_gap, 0, @@ -129,8 +128,7 @@ class ChieriBotV4Rois: ) for hi in range(self.horizontal_items): if hi > 0: - rect += ((self.width + self.horizontal_gap), 0, 0, 0) - int_rect = construct_int_xywh_rect(rect) - results.append(crop_xywh(img_bgr, int_rect)) + last_rect += ((self.width + self.horizontal_gap), 0, 0, 0) + results.append(crop_xywh(img_bgr, last_rect.rounded())) return results diff --git a/src/arcaea_offline_ocr/types.py b/src/arcaea_offline_ocr/types.py index 7f1bc5b..d1bafa8 100644 --- a/src/arcaea_offline_ocr/types.py +++ b/src/arcaea_offline_ocr/types.py @@ -1,25 +1,36 @@ -from collections.abc import Iterable -from typing import NamedTuple, Tuple, Union +from math import floor +from typing import Callable, NamedTuple, Union import numpy as np Mat = np.ndarray +_IntOrFloat = Union[int, float] + class XYWHRect(NamedTuple): - x: int - y: int - w: int - h: int + x: _IntOrFloat + y: _IntOrFloat + w: _IntOrFloat + h: _IntOrFloat - def __add__(self, other: Union["XYWHRect", Tuple[int, int, int, int]]): - if not isinstance(other, Iterable) or len(other) != 4: - raise ValueError() + def _to_int(self, func: Callable[[_IntOrFloat], int]): + return (func(self.x), func(self.y), func(self.w), func(self.h)) + + def rounded(self): + return self._to_int(round) + + def floored(self): + return self._to_int(floor) + + def __add__(self, other): + if not isinstance(other, (list, tuple)) or len(other) != 4: + raise TypeError() return self.__class__(*[a + b for a, b in zip(self, other)]) - def __sub__(self, other: Union["XYWHRect", Tuple[int, int, int, int]]): - if not isinstance(other, Iterable) or len(other) != 4: - raise ValueError() + def __sub__(self, other): + if not isinstance(other, (list, tuple)) or len(other) != 4: + raise TypeError() return self.__class__(*[a - b for a, b in zip(self, other)]) diff --git a/src/arcaea_offline_ocr/utils.py b/src/arcaea_offline_ocr/utils.py index a6238c6..2ca6458 100644 --- a/src/arcaea_offline_ocr/utils.py +++ b/src/arcaea_offline_ocr/utils.py @@ -1,5 +1,5 @@ from collections.abc import Iterable -from typing import Callable, TypeVar, Union, overload +from typing import TypeVar, overload import cv2 import numpy as np @@ -15,32 +15,25 @@ def imread_unicode(filepath: str, flags: int = cv2.IMREAD_UNCHANGED): return cv2.imdecode(np.fromfile(filepath, dtype=np.uint8), flags) -def construct_int_xywh_rect( - rect: XYWHRect, func: Callable[[Union[int, float]], int] = round -): - return XYWHRect(*[func(num) for num in rect]) +@overload +def apply_factor(item: int, factor: float) -> float: ... @overload -def apply_factor(item: int, factor: float) -> float: - ... - - -@overload -def apply_factor(item: float, factor: float) -> float: - ... +def apply_factor(item: float, factor: float) -> float: ... T = TypeVar("T", bound=Iterable) @overload -def apply_factor(item: T, factor: float) -> T: - ... +def apply_factor(item: T, factor: float) -> T: ... def apply_factor(item, factor: float): if isinstance(item, (int, float)): return item * factor + if isinstance(item, XYWHRect): + return item.__class__(*[i * factor for i in item]) if isinstance(item, Iterable): return item.__class__([i * factor for i in item])