diff --git a/src/arcaea_offline_ocr/device/v2/ocr.py b/src/arcaea_offline_ocr/device/v2/ocr.py index cff2f0a..bff9c6b 100644 --- a/src/arcaea_offline_ocr/device/v2/ocr.py +++ b/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 diff --git a/src/arcaea_offline_ocr/device/v2/sizes.py b/src/arcaea_offline_ocr/device/v2/sizes.py index 254695a..3347cb2 100644 --- a/src/arcaea_offline_ocr/device/v2/sizes.py +++ b/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) diff --git a/src/arcaea_offline_ocr/mask.py b/src/arcaea_offline_ocr/mask.py index 6466d2d..085c99d 100644 --- a/src/arcaea_offline_ocr/mask.py +++ b/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