wip: SizesV2

src/arcaea_offline_ocr/device/v2/ocr.py

@@ -7,7 +7,16 @@ import numpy as np
 from PIL import Image
 
 from ...crop import crop_xywh
-from ...mask import mask_byd, mask_ftr, mask_gray, mask_prs, mask_pst, mask_white
+from ...mask import (
+    mask_byd,
+    mask_ftr,
+    mask_gray,
+    mask_max_recall_purple,
+    mask_pfl_white,
+    mask_prs,
+    mask_pst,
+    mask_white,
+)
 from ...ocr import (
     FixRects,
     ocr_digit_samples_knn,
@@ -22,6 +31,7 @@ from ..shared import DeviceOcrResult
 from .preprocess import find_digits_preprocess
 from .rois import DeviceV2Rois
 from .shared import MAX_RECALL_CLOSE_KERNEL
+from .sizes import SizesV2
 
 
 class DeviceV2Ocr:
@@ -105,20 +115,33 @@ class DeviceV2Ocr:
                 roi = cv2.fillPoly(roi, [contour], [0])
         return ocr_digits_by_contour_knn(roi, self.knn_model)
 
+    def mask_pfl(self, pfl_roi: Mat, rois: DeviceV2Rois):
+        return (
+            mask_pfl_white(cv2.cvtColor(pfl_roi, cv2.COLOR_BGR2HSV))
+            if isinstance(rois.sizes, SizesV2)
+            else mask_gray(pfl_roi)
+        )
+
     def ocr_pure(self, rois: DeviceV2Rois):
-        roi = mask_gray(rois.pure)
+        roi = self.mask_pfl(rois.pure, rois)
         return self._base_ocr_pfl(roi, rois.sizes.factor)
 
     def ocr_far(self, rois: DeviceV2Rois):
-        roi = mask_gray(rois.far)
+        roi = self.mask_pfl(rois.far, rois)
         return self._base_ocr_pfl(roi, rois.sizes.factor)
 
     def ocr_lost(self, rois: DeviceV2Rois):
-        roi = mask_gray(rois.lost)
+        roi = self.mask_pfl(rois.lost, rois)
         return self._base_ocr_pfl(roi, rois.sizes.factor)
 
     def ocr_max_recall(self, rois: DeviceV2Rois):
-        roi = mask_gray(rois.max_recall_rating_class)
+        roi = (
+            mask_max_recall_purple(
+                cv2.cvtColor(rois.max_recall_rating_class, cv2.COLOR_BGR2HSV)
+            )
+            if isinstance(rois.sizes, SizesV2)
+            else mask_gray(rois.max_recall_rating_class)
+        )
         roi_closed = cv2.morphologyEx(roi, cv2.MORPH_CLOSE, MAX_RECALL_CLOSE_KERNEL)
         contours, _ = cv2.findContours(
             roi_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE

src/arcaea_offline_ocr/device/v2/sizes.py

@@ -1,10 +1,6 @@
 from typing import Tuple, Union
 
 
-def to_int(num: Union[int, float]) -> int:
-    return round(num)
-
-
 def apply_factor(num: Union[int, float], factor: float):
     return num * factor
 
@@ -172,3 +168,87 @@ class SizesV1(Sizes):
     @property
     def TITLE_WIDTH_RIGHT(self):
         return self.apply_factor(275)
+
+
+class SizesV2(Sizes):
+    def __init__(self, factor: float):
+        self.factor = factor
+
+    def apply_factor(self, num):
+        return apply_factor(num, self.factor)
+
+    @property
+    def TOP_BAR_HEIGHT(self):
+        return self.apply_factor(50)
+
+    @property
+    def SCORE_PANEL(self) -> Tuple[int, int]:
+        return tuple(self.apply_factor(num) for num in [447, 233])
+
+    @property
+    def PFL_TOP_FROM_VMID(self):
+        return self.apply_factor(142)
+
+    @property
+    def PFL_LEFT_FROM_HMID(self):
+        return self.apply_factor(10)
+
+    @property
+    def PFL_WIDTH(self):
+        return self.apply_factor(60)
+
+    @property
+    def PFL_FONT_PX(self):
+        return self.apply_factor(16)
+
+    @property
+    def PURE_FAR_GAP(self):
+        return self.apply_factor(20)
+
+    @property
+    def FAR_LOST_GAP(self):
+        return self.apply_factor(23)
+
+    @property
+    def SCORE_BOTTOM_FROM_VMID(self):
+        return self.apply_factor(-50)
+
+    @property
+    def SCORE_FONT_PX(self):
+        return self.apply_factor(45)
+
+    @property
+    def SCORE_WIDTH(self):
+        return self.apply_factor(280)
+
+    @property
+    def JACKET_RIGHT_FROM_HOR_MID(self):
+        return self.apply_factor(-235)
+
+    @property
+    def JACKET_WIDTH(self):
+        return self.apply_factor(375)
+
+    @property
+    def MR_RT_RIGHT_FROM_HMID(self):
+        return self.apply_factor(-330)
+
+    @property
+    def MR_RT_WIDTH(self):
+        return self.apply_factor(330)
+
+    @property
+    def MR_RT_HEIGHT(self):
+        return self.apply_factor(75)
+
+    @property
+    def TITLE_BOTTOM_FROM_VMID(self):
+        return self.apply_factor(-265)
+
+    @property
+    def TITLE_FONT_PX(self):
+        return self.apply_factor(40)
+
+    @property
+    def TITLE_WIDTH_RIGHT(self):
+        return self.apply_factor(275)

src/arcaea_offline_ocr/mask.py

@@ -8,6 +8,8 @@ __all__ = [
     "GRAY_MAX_HSV",
     "WHITE_MIN_HSV",
     "WHITE_MAX_HSV",
+    "PFL_WHITE_MIN_HSV",
+    "PFL_WHITE_MAX_HSV",
     "PST_MIN_HSV",
     "PST_MAX_HSV",
     "PRS_MIN_HSV",
@@ -16,13 +18,17 @@ __all__ = [
     "FTR_MAX_HSV",
     "BYD_MIN_HSV",
     "BYD_MAX_HSV",
+    "MAX_RECALL_PURPLE_MIN_HSV",
+    "MAX_RECALL_PURPLE_MAX_HSV",
     "mask_gray",
     "mask_white",
+    "mask_pfl_white",
     "mask_pst",
     "mask_prs",
     "mask_ftr",
     "mask_byd",
     "mask_rating_class",
+    "mask_max_recall_purple",
 ]
 
 GRAY_MIN_HSV = np.array([0, 0, 70], np.uint8)
@@ -34,6 +40,9 @@ GRAY_MAX_BGR = np.array([160] * 3, np.uint8)
 WHITE_MIN_HSV = np.array([0, 0, 240], np.uint8)
 WHITE_MAX_HSV = np.array([179, 10, 255], np.uint8)
 
+PFL_WHITE_MIN_HSV = np.array([0, 0, 248], np.uint8)
+PFL_WHITE_MAX_HSV = np.array([179, 10, 255], np.uint8)
+
 PST_MIN_HSV = np.array([100, 50, 80], np.uint8)
 PST_MAX_HSV = np.array([100, 255, 255], np.uint8)
 
@@ -46,6 +55,9 @@ FTR_MAX_HSV = np.array([155, 181, 150], np.uint8)
 BYD_MIN_HSV = np.array([170, 50, 50], np.uint8)
 BYD_MAX_HSV = np.array([179, 210, 198], np.uint8)
 
+MAX_RECALL_PURPLE_MIN_HSV = np.array([125, 0, 0], np.uint8)
+MAX_RECALL_PURPLE_MAX_HSV = np.array([130, 100, 150], np.uint8)
+
 
 def mask_gray(__img_bgr: Mat):
     # bgr_value_equal_mask = all(__img_bgr[:, 1:] == __img_bgr[:, :-1], axis=1)
@@ -63,6 +75,12 @@ def mask_white(img_hsv: Mat):
     return mask
 
 
+def mask_pfl_white(img_hsv: Mat):
+    mask = cv2.inRange(img_hsv, PFL_WHITE_MIN_HSV, PFL_WHITE_MAX_HSV)
+    mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2)))
+    return mask
+
+
 def mask_pst(img_hsv: Mat):
     mask = cv2.inRange(img_hsv, PST_MIN_HSV, PST_MAX_HSV)
     mask = cv2.dilate(mask, (1, 1))
@@ -93,3 +111,9 @@ def mask_rating_class(img_hsv: Mat):
     ftr = mask_ftr(img_hsv)
     byd = mask_byd(img_hsv)
     return cv2.bitwise_or(byd, cv2.bitwise_or(ftr, cv2.bitwise_or(pst, prs)))
+
+
+def mask_max_recall_purple(img_hsv: Mat):
+    mask = cv2.inRange(img_hsv, MAX_RECALL_PURPLE_MIN_HSV, MAX_RECALL_PURPLE_MAX_HSV)
+    mask = cv2.dilate(mask, (2, 2))
+    return mask