refactor: XYWHRect and b30 ocr

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(
+        rating_class_rect = self.rois.component_rois.rating_class_rect.rounded()
-            self.rois.component_rois.rating_class_rect
+
-        )
         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(
+        jacket_rect = self.rois.component_rois.jacket_rect.floored()
-            self.rois.component_rois.jacket_rect, floor
-        )
         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)
 

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
+    vertical_items = 10
-    def vertical_items(self):
-        return 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, rect.rounded()))
-                results.append(crop_xywh(img_bgr, int_rect))
+            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)
+                last_rect += ((self.width + self.horizontal_gap), 0, 0, 0)
-            int_rect = construct_int_xywh_rect(rect)
+            results.append(crop_xywh(img_bgr, last_rect.rounded()))
-            results.append(crop_xywh(img_bgr, int_rect))
 
         return results

src/arcaea_offline_ocr/types.py

@@ -1,25 +1,36 @@
-from collections.abc import Iterable
+from math import floor
-from typing import NamedTuple, Tuple, Union
+from typing import Callable, NamedTuple, Union
 
 import numpy as np
 
 Mat = np.ndarray
 
+_IntOrFloat = Union[int, float]
+
 
 class XYWHRect(NamedTuple):
-    x: int
+    x: _IntOrFloat
-    y: int
+    y: _IntOrFloat
-    w: int
+    w: _IntOrFloat
-    h: int
+    h: _IntOrFloat
 
-    def __add__(self, other: Union["XYWHRect", Tuple[int, int, int, int]]):
+    def _to_int(self, func: Callable[[_IntOrFloat], int]):
-        if not isinstance(other, Iterable) or len(other) != 4:
+        return (func(self.x), func(self.y), func(self.w), func(self.h))
-            raise ValueError()
+
+    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]]):
+    def __sub__(self, other):
-        if not isinstance(other, Iterable) or len(other) != 4:
+        if not isinstance(other, (list, tuple)) or len(other) != 4:
-            raise ValueError()
+            raise TypeError()
 
         return self.__class__(*[a - b for a, b in zip(self, other)])

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(
+@overload
-    rect: XYWHRect, func: Callable[[Union[int, float]], int] = round
+def apply_factor(item: int, factor: float) -> float: ...
-):
-    return XYWHRect(*[func(num) for num in rect])
 
 
 @overload
-def apply_factor(item: int, factor: float) -> float:
+def apply_factor(item: float, factor: float) -> float: ...
-    ...
-
-
-@overload
-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])