Merge pull request #7 from 283375/device-v2

device-v2 temporary merge

src/arcaea_offline_ocr/crop.py

@@ -1,53 +1,56 @@
-from typing import Any, Tuple
+from math import floor
+from typing import Tuple
+
+from numpy import all, array, count_nonzero
 
-from .device import Device
 from .types import Mat
 
-__all__ = [
-    "crop_img",
-    "crop_from_device_attr",
-    "crop_to_pure",
-    "crop_to_far",
-    "crop_to_lost",
-    "crop_to_max_recall",
-    "crop_to_rating_class",
-    "crop_to_score",
-    "crop_to_title",
-]
+__all__ = ["crop_xywh", "crop_black_edges"]
 
 
-def crop_img(img: Mat, *, top: int, left: int, bottom: int, right: int):
-    return img[top:bottom, left:right]
-
-
-def crop_from_device_attr(img: Mat, rect: Tuple[int, int, int, int]):
+def crop_xywh(mat: Mat, rect: Tuple[int, int, int, int]):
     x, y, w, h = rect
-    return crop_img(img, top=y, left=x, bottom=y + h, right=x + w)
+    return mat[y : y + h, x : x + w]
 
 
-def crop_to_pure(screenshot: Mat, device: Device):
-    return crop_from_device_attr(screenshot, device.pure)
+def is_black_edge(list_of_pixels: Mat, black_pixel=None):
+    if black_pixel is None:
+        black_pixel = array([0, 0, 0], list_of_pixels.dtype)
+    pixels = list_of_pixels.reshape([-1, 3])
+    return count_nonzero(all(pixels < black_pixel, axis=1)) > floor(len(pixels) * 0.6)
 
 
-def crop_to_far(screenshot: Mat, device: Device):
-    return crop_from_device_attr(screenshot, device.far)
+def crop_black_edges(screenshot: Mat):
+    cropped = screenshot.copy()
+    black_pixel = array([50, 50, 50], screenshot.dtype)
+    height, width = screenshot.shape[:2]
+    left = 0
+    right = width
+    top = 0
+    bottom = height
 
+    for i in range(width):
+        column = cropped[:, i]
+        if not is_black_edge(column, black_pixel):
+            break
+        left += 1
 
-def crop_to_lost(screenshot: Mat, device: Device):
-    return crop_from_device_attr(screenshot, device.lost)
+    for i in sorted(range(width), reverse=True):
+        column = cropped[:, i]
+        if i <= left + 1 or not is_black_edge(column, black_pixel):
+            break
+        right -= 1
 
+    for i in range(height):
+        row = cropped[i]
+        if not is_black_edge(row, black_pixel):
+            break
+        top += 1
 
-def crop_to_max_recall(screenshot: Mat, device: Device):
-    return crop_from_device_attr(screenshot, device.max_recall)
+    for i in sorted(range(height), reverse=True):
+        row = cropped[i]
+        if i <= top + 1 or not is_black_edge(row, black_pixel):
+            break
+        bottom -= 1
 
-
-def crop_to_rating_class(screenshot: Mat, device: Device):
-    return crop_from_device_attr(screenshot, device.rating_class)
-
-
-def crop_to_score(screenshot: Mat, device: Device):
-    return crop_from_device_attr(screenshot, device.score)
-
-
-def crop_to_title(screenshot: Mat, device: Device):
-    return crop_from_device_attr(screenshot, device.title)
+    return cropped[top:bottom, left:right]

src/arcaea_offline_ocr/device/v1/crop.py

@@ -0,0 +1,64 @@
+from math import floor
+from typing import Any, Tuple
+
+from numpy import all, array, count_nonzero
+
+from ...types import Mat
+from .definition import Device
+
+__all__ = [
+    "crop_img",
+    "crop_from_device_attr",
+    "crop_to_pure",
+    "crop_to_far",
+    "crop_to_lost",
+    "crop_to_max_recall",
+    "crop_to_rating_class",
+    "crop_to_score",
+    "crop_to_title",
+    "crop_black_edges",
+]
+
+
+def crop_img(img: Mat, *, top: int, left: int, bottom: int, right: int):
+    return img[top:bottom, left:right]
+
+
+def crop_from_device_attr(img: Mat, rect: Tuple[int, int, int, int]):
+    x, y, w, h = rect
+    return crop_img(img, top=y, left=x, bottom=y + h, right=x + w)
+
+
+def crop_to_pure(screenshot: Mat, device: Device):
+    return crop_from_device_attr(screenshot, device.pure)
+
+
+def crop_to_far(screenshot: Mat, device: Device):
+    return crop_from_device_attr(screenshot, device.far)
+
+
+def crop_to_lost(screenshot: Mat, device: Device):
+    return crop_from_device_attr(screenshot, device.lost)
+
+
+def crop_to_max_recall(screenshot: Mat, device: Device):
+    return crop_from_device_attr(screenshot, device.max_recall)
+
+
+def crop_to_rating_class(screenshot: Mat, device: Device):
+    return crop_from_device_attr(screenshot, device.rating_class)
+
+
+def crop_to_score(screenshot: Mat, device: Device):
+    return crop_from_device_attr(screenshot, device.score)
+
+
+def crop_to_title(screenshot: Mat, device: Device):
+    return crop_from_device_attr(screenshot, device.title)
+
+
+def is_black_edge(list_of_pixels: Mat, black_pixel=None):
+    if black_pixel is None:
+        black_pixel = array([0, 0, 0], list_of_pixels.dtype)
+    pixels = list_of_pixels.reshape([-1, 3])
+    return count_nonzero(all(pixels < black_pixel, axis=1)) > floor(len(pixels) * 0.6)

src/arcaea_offline_ocr/device/v1/definition.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+from typing import Any, Dict, Tuple
+
+__all__ = ["Device"]
+
+
+@dataclass(kw_only=True)
+class Device:
+    version: int
+    uuid: str
+    name: str
+    pure: Tuple[int, int, int, int]
+    far: Tuple[int, int, int, int]
+    lost: Tuple[int, int, int, int]
+    max_recall: Tuple[int, int, int, int]
+    rating_class: Tuple[int, int, int, int]
+    score: Tuple[int, int, int, int]
+    title: Tuple[int, int, int, int]
+
+    @classmethod
+    def from_json_object(cls, json_dict: Dict[str, Any]):
+        if json_dict["version"] == 1:
+            return cls(
+                version=1,
+                uuid=json_dict["uuid"],
+                name=json_dict["name"],
+                pure=json_dict["pure"],
+                far=json_dict["far"],
+                lost=json_dict["lost"],
+                max_recall=json_dict["max_recall"],
+                rating_class=json_dict["rating_class"],
+                score=json_dict["score"],
+                title=json_dict["title"],
+            )
+
+    def repr_info(self):
+        return f"Device(version={self.version}, uuid={repr(self.uuid)}, name={repr(self.name)})"

src/arcaea_offline_ocr/device/v2/definition.py

@@ -0,0 +1,26 @@
+from typing import Iterable
+
+from attrs import define, field
+
+from ...types import XYWHRect
+
+
+def iterable_to_xywh_rect(__iter: Iterable) -> XYWHRect:
+    return XYWHRect(*__iter)
+
+
+@define(kw_only=True)
+class DeviceV2:
+    version = field(type=int)
+    uuid = field(type=str)
+    name = field(type=str)
+    crop_black_edges = field(type=bool)
+    factor = field(type=float)
+    pure = field(converter=iterable_to_xywh_rect, default=[0, 0, 0, 0])
+    far = field(converter=iterable_to_xywh_rect, default=[0, 0, 0, 0])
+    lost = field(converter=iterable_to_xywh_rect, default=[0, 0, 0, 0])
+    score = field(converter=iterable_to_xywh_rect, default=[0, 0, 0, 0])
+    max_recall_rating_class = field(
+        converter=iterable_to_xywh_rect, default=[0, 0, 0, 0]
+    )
+    title = field(converter=iterable_to_xywh_rect, default=[0, 0, 0, 0])

src/arcaea_offline_ocr/device/v2/find.py

@@ -0,0 +1,202 @@
+from typing import List, Tuple
+
+import attrs
+import cv2
+import numpy as np
+
+from ...crop import crop_xywh
+from ...mask import mask_gray
+from ...types import Mat, XYWHRect
+from .definition import DeviceV2
+from .shared import *
+
+
+@attrs.define(kw_only=True)
+class FindOcrBoundingRectsResult:
+    pure: XYWHRect
+    far: XYWHRect
+    lost: XYWHRect
+    max_recall: XYWHRect
+    gray_masked_image: Mat
+
+
+def find_ocr_bounding_rects(__img_bgr: Mat, device: DeviceV2):
+    """
+    [DEPRECATED]
+    ---
+    Deprecated since new method supports directly calculate rois.
+    """
+
+    img_masked = mask_gray(__img_bgr)
+
+    # process pure/far/lost
+    pfl_roi = crop_xywh(img_masked, device.pure_far_lost)
+    # close small gaps in fonts
+    # pfl_roi = cv2.GaussianBlur(pfl_roi, [5, 5], 0, 0)
+    # cv2.imshow("test2", pfl_roi)
+    # cv2.waitKey(0)
+
+    pfl_roi = cv2.morphologyEx(pfl_roi, cv2.MORPH_OPEN, PFL_DENOISE_KERNEL)
+    pfl_roi = cv2.morphologyEx(pfl_roi, cv2.MORPH_CLOSE, PFL_CLOSE_HORIZONTAL_KERNEL)
+
+    pfl_contours, _ = cv2.findContours(
+        pfl_roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
+    )
+    pfl_contours = sorted(pfl_contours, key=cv2.contourArea)
+
+    # pfl_roi_cnt = cv2.drawContours(pfl_roi, pfl_contours, -1, [50], 2)
+    # cv2.imshow("test2", pfl_roi_cnt)
+    # cv2.waitKey(0)
+
+    pfl_rects = [list(cv2.boundingRect(c)) for c in pfl_contours]
+
+    # for r in pfl_rects:
+    #     img = pfl_roi.copy()
+    #     cv2.imshow("test2", cv2.rectangle(img, r, [80] * 3, 2))
+    #     cv2.waitKey(0)
+
+    # only keep those rect.height > mask.height * 0.15
+    pfl_rects = list(filter(lambda rect: rect[3] > pfl_roi.shape[0] * 0.15, pfl_rects))
+    # choose the first 3 rects by rect.x value
+    pfl_rects = sorted(pfl_rects, key=lambda rect: rect[0])[:3]
+    # and sort them by rect.y
+    # ensure it is pure -> far -> lost roi.
+    pure_rect, far_rect, lost_rect = sorted(pfl_rects, key=lambda rect: rect[1])
+
+    # for r in [pure_rect, far_rect, lost_rect]:
+    #     img = pfl_roi.copy()
+    #     cv2.imshow("test2", cv2.rectangle(img, r, [80] * 3, 2))
+    #     cv2.waitKey(0)
+
+    # process max recall
+    max_recall_roi = crop_xywh(img_masked, device.max_recall_rating_class)
+    max_recall_roi = cv2.morphologyEx(
+        max_recall_roi, cv2.MORPH_OPEN, MAX_RECALL_DENOISE_KERNEL
+    )
+    max_recall_roi = cv2.erode(max_recall_roi, MAX_RECALL_ERODE_KERNEL)
+    max_recall_roi = cv2.morphologyEx(
+        max_recall_roi, cv2.MORPH_CLOSE, MAX_RECALL_CLOSE_KERNEL
+    )
+    max_recall_contours, _ = cv2.findContours(
+        max_recall_roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
+    )
+    max_recall_rects = [list(cv2.boundingRect(c)) for c in max_recall_contours]
+    # only keep those rect.height > mask.height * 0.1
+    max_recall_rects = list(
+        filter(lambda rect: rect[3] > max_recall_roi.shape[0] * 0.1, max_recall_rects)
+    )
+    # select the 2nd rect by rect.x
+    max_recall_rect = max_recall_rects[1]
+
+    # img = max_recall_roi.copy()
+    # cv2.imshow("test2", cv2.rectangle(img, max_recall_rect, [80] * 3, 2))
+    # cv2.waitKey(0)
+
+    # finally, map rect geometries to the original image
+    for rect in [pure_rect, far_rect, lost_rect]:
+        rect[0] += device.pure_far_lost[0]
+        rect[1] += device.pure_far_lost[1]
+
+    for rect in [max_recall_rect]:
+        rect[0] += device.max_recall_rating_class[0]
+        rect[1] += device.max_recall_rating_class[1]
+
+    # add a 2px border to every rect
+    for rect in [pure_rect, far_rect, lost_rect, max_recall_rect]:
+        # width += 2, height += 2
+        rect[2] += 4
+        rect[3] += 4
+        # top -= 1, left -= 1
+        rect[0] -= 2
+        rect[1] -= 2
+
+    return FindOcrBoundingRectsResult(
+        pure=XYWHRect(*pure_rect),
+        far=XYWHRect(*far_rect),
+        lost=XYWHRect(*lost_rect),
+        max_recall=XYWHRect(*max_recall_rect),
+        gray_masked_image=img_masked,
+    )
+
+
+def find_digits_preprocess(__img_masked: Mat) -> Mat:
+    img = __img_masked.copy()
+    img_denoised = cv2.morphologyEx(img, cv2.MORPH_OPEN, PFL_DENOISE_KERNEL)
+    # img_denoised = cv2.bitwise_and(img, img_denoised)
+
+    denoise_contours, _ = cv2.findContours(
+        img_denoised, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+    )
+    # cv2.drawContours(img_denoised, contours, -1, [128], 2)
+
+    # fill all contour.area < max(contour.area) * ratio with black pixels
+    # for denoise purposes
+
+    # define threshold contour area
+    # we assume the smallest digit "1", is 80% height of the image,
+    # and at least 1.5 pixel wide, considering cv2.contourArea always
+    # returns a smaller value than the actual contour area.
+    max_contour_area = __img_masked.shape[0] * 0.8 * 1.5
+    filtered_contours = list(
+        filter(lambda c: cv2.contourArea(c) >= max_contour_area, denoise_contours)
+    )
+
+    filtered_contours_flattened = {tuple(c.flatten()) for c in filtered_contours}
+
+    for contour in denoise_contours:
+        if tuple(contour.flatten()) not in filtered_contours_flattened:
+            img_denoised = cv2.fillPoly(img_denoised, [contour], [0])
+
+    # old algorithm, finding the largest contour area
+    ## contour_area_tuples = [(contour, cv2.contourArea(contour)) for contour in contours]
+    ## contour_area_tuples = sorted(
+    ##     contour_area_tuples, key=lambda item: item[1], reverse=True
+    ## )
+    ## max_contour_area = contour_area_tuples[0][1]
+    ## print(max_contour_area, [item[1] for item in contour_area_tuples])
+    ## contours_filter_end_index = len(contours)
+    ## for i, item in enumerate(contour_area_tuples):
+    ##     contour, area = item
+    ##     if area < max_contour_area * 0.15:
+    ##         contours_filter_end_index = i
+    ##         break
+    ## contours = [item[0] for item in contour_area_tuples]
+    ## for contour in contours[-contours_filter_end_index - 1:]:
+    ##     img = cv2.fillPoly(img, [contour], [0])
+    ##     img_denoised = cv2.fillPoly(img_denoised, [contour], [0])
+    ## contours = contours[:contours_filter_end_index]
+
+    return img_denoised
+
+
+def find_digits(__img_masked: Mat) -> List[Mat]:
+    img_denoised = find_digits_preprocess(__img_masked)
+
+    cv2.imshow("den", img_denoised)
+    cv2.waitKey(0)
+
+    contours, _ = cv2.findContours(
+        img_denoised, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+    )
+
+    img_x_roi = []  # type: List[Tuple[int, Mat]]
+    # img_x_roi = list[tuple[int, Mat]] - list[tuple[rect.x, roi_denoised]]
+    for contour in contours:
+        rect = cv2.boundingRect(contour)
+        # filter out rect.height < img.height * factor
+        if rect[3] < img_denoised.shape[0] * 0.8:
+            continue
+        contour -= (rect[0], rect[1])
+        img_denoised_roi = crop_xywh(img_denoised, rect)
+        # make a same size black image
+        contour_mask = np.zeros(img_denoised_roi.shape, img_denoised_roi.dtype)
+        # fill the contour area with white pixels
+        contour_mask = cv2.fillPoly(contour_mask, [contour], [255])
+        # apply mask to cropped images
+        img_denoised_roi_masked = cv2.bitwise_and(contour_mask, img_denoised_roi)
+        img_x_roi.append((rect[0], img_denoised_roi_masked))
+
+    # sort by rect.x
+    img_x_roi = sorted(img_x_roi, key=lambda item: item[0])
+
+    return [item[1] for item in img_x_roi]

src/arcaea_offline_ocr/device/v2/ocr.py

@@ -0,0 +1,89 @@
+from typing import Optional
+
+import attrs
+import cv2
+import numpy as np
+
+from ...mask import mask_byd, mask_ftr, mask_gray, mask_prs, mask_pst, mask_white
+from ...ocr import ocr_digits_knn_model
+from ...types import Mat, cv2_ml_KNearest
+from .find import find_digits
+from .rois import DeviceV2Rois
+
+
+@attrs.define
+class DeviceV2OcrResult:
+    pure: int
+    far: int
+    lost: int
+    score: int
+    rating_class: int
+    max_recall: int
+    title: Optional[str]
+
+
+class DeviceV2Ocr:
+    def __init__(self):
+        self.__rois = None
+        self.__knn_model = None
+
+    @property
+    def rois(self):
+        if not self.__rois:
+            raise ValueError("`rois` unset.")
+        return self.__rois
+
+    @rois.setter
+    def rois(self, rois: DeviceV2Rois):
+        self.__rois = rois
+
+    @property
+    def knn_model(self):
+        if not self.__knn_model:
+            raise ValueError("`knn_model` unset.")
+        return self.__knn_model
+
+    @knn_model.setter
+    def knn_model(self, model: cv2_ml_KNearest):
+        self.__knn_model = model
+
+    def _base_ocr_digits(self, roi_processed: Mat):
+        digits = find_digits(roi_processed)
+        result = ""
+        for digit in digits:
+            roi_result = ocr_digits_knn_model(digit, self.knn_model)
+            if roi_result is not None:
+                result += str(roi_result)
+        return int(result, base=10)
+
+    @property
+    def pure(self):
+        roi = mask_gray(self.rois.pure)
+        return self._base_ocr_digits(roi)
+
+    @property
+    def far(self):
+        roi = mask_gray(self.rois.far)
+        return self._base_ocr_digits(roi)
+
+    @property
+    def lost(self):
+        roi = mask_gray(self.rois.lost)
+        return self._base_ocr_digits(roi)
+
+    @property
+    def score(self):
+        roi = cv2.cvtColor(self.rois.score, cv2.COLOR_BGR2HSV)
+        roi = mask_white(roi)
+        return self._base_ocr_digits(roi)
+
+    @property
+    def rating_class(self):
+        roi = cv2.cvtColor(self.rois.max_recall_rating_class, cv2.COLOR_BGR2HSV)
+        results = [
+            mask_pst(roi),
+            mask_prs(roi),
+            mask_ftr(roi),
+            mask_byd(roi),
+        ]
+        return max(enumerate(results), key=lambda e: np.count_nonzero(e[1]))[0]

src/arcaea_offline_ocr/device/v2/rois.py

@@ -0,0 +1,265 @@
+from typing import Tuple, Union
+
+from ...crop import crop_black_edges, crop_xywh
+from ...types import Mat, XYWHRect
+from .definition import DeviceV2
+
+
+def to_int(num: Union[int, float]) -> int:
+    return round(num)
+
+
+def apply_factor(num: Union[int, float], factor: float):
+    return num * factor
+
+
+class 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 [485, 239])
+
+    @property
+    def PFL_TOP_FROM_VER_MID(self):
+        return self.apply_factor(135)
+
+    @property
+    def PFL_LEFT_FROM_HOR_MID(self):
+        return self.apply_factor(5)
+
+    @property
+    def PFL_WIDTH(self):
+        return self.apply_factor(150)
+
+    @property
+    def PFL_FONT_PX(self):
+        return self.apply_factor(26)
+
+    @property
+    def PURE_FAR_GAP(self):
+        return self.apply_factor(12)
+
+    @property
+    def FAR_LOST_GAP(self):
+        return self.apply_factor(10)
+
+    @property
+    def SCORE_BOTTOM_FROM_VER_MID(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 COVER_RIGHT_FROM_HOR_MID(self):
+        return self.apply_factor(-235)
+
+    @property
+    def COVER_WIDTH(self):
+        return self.apply_factor(375)
+
+    @property
+    def MAX_RECALL_RATING_CLASS_RIGHT_FROM_HOR_MID(self):
+        return self.apply_factor(-300)
+
+    @property
+    def MAX_RECALL_RATING_CLASS_WIDTH(self):
+        return self.apply_factor(275)
+
+    @property
+    def MAX_RECALL_RATING_CLASS_HEIGHT(self):
+        return self.apply_factor(75)
+
+    @property
+    def TITLE_BOTTOM_FROM_VER_MID(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)
+
+
+class DeviceV2Rois:
+    def __init__(self, device: DeviceV2):
+        self.device = device
+        self.sizes = Sizes(self.device.factor)
+        self.__img = None
+
+    @staticmethod
+    def construct_int_xywh_rect(x, y, w, h) -> XYWHRect:
+        return XYWHRect(*[to_int(item) for item in [x, y, w, h]])
+
+    @property
+    def img(self):
+        return self.__img
+
+    @img.setter
+    def img(self, img: Mat):
+        self.__img = (
+            crop_black_edges(img) if self.device.crop_black_edges else img.copy()
+        )
+
+    @property
+    def h(self):
+        return self.img.shape[0]
+
+    @property
+    def ver_mid(self):
+        return self.h / 2
+
+    @property
+    def w(self):
+        return self.img.shape[1]
+
+    @property
+    def hor_mid(self):
+        return self.w / 2
+
+    @property
+    def h_fixed(self):
+        """img_height -= top_bar_height"""
+        return self.h - self.sizes.TOP_BAR_HEIGHT
+
+    @property
+    def h_fixed_mid(self):
+        return self.sizes.TOP_BAR_HEIGHT + self.h_fixed / 2
+
+    @property
+    def pfl_top(self):
+        return self.h_fixed_mid + self.sizes.PFL_TOP_FROM_VER_MID
+
+    @property
+    def pfl_left(self):
+        return self.hor_mid + self.sizes.PFL_LEFT_FROM_HOR_MID
+
+    @property
+    def pure_rect(self):
+        return self.construct_int_xywh_rect(
+            x=self.pfl_left,
+            y=self.pfl_top,
+            w=self.sizes.PFL_WIDTH,
+            h=self.sizes.PFL_FONT_PX,
+        )
+
+    @property
+    def pure(self):
+        return crop_xywh(self.img, self.pure_rect)
+
+    @property
+    def far_rect(self):
+        return self.construct_int_xywh_rect(
+            x=self.pfl_left,
+            y=self.pfl_top + self.sizes.PFL_FONT_PX + self.sizes.PURE_FAR_GAP,
+            w=self.sizes.PFL_WIDTH,
+            h=self.sizes.PFL_FONT_PX,
+        )
+
+    @property
+    def far(self):
+        return crop_xywh(self.img, self.far_rect)
+
+    @property
+    def lost_rect(self):
+        return self.construct_int_xywh_rect(
+            x=self.pfl_left,
+            y=(
+                self.pfl_top
+                + self.sizes.PFL_FONT_PX * 2
+                + self.sizes.PURE_FAR_GAP
+                + self.sizes.FAR_LOST_GAP
+            ),
+            w=self.sizes.PFL_WIDTH,
+            h=self.sizes.PFL_FONT_PX,
+        )
+
+    @property
+    def lost(self):
+        return crop_xywh(self.img, self.lost_rect)
+
+    @property
+    def score_rect(self):
+        return self.construct_int_xywh_rect(
+            x=self.hor_mid - (self.sizes.SCORE_WIDTH / 2),
+            y=(
+                self.h_fixed_mid
+                + self.sizes.SCORE_BOTTOM_FROM_VER_MID
+                - self.sizes.SCORE_FONT_PX
+            ),
+            w=self.sizes.SCORE_WIDTH,
+            h=self.sizes.SCORE_FONT_PX,
+        )
+
+    @property
+    def score(self):
+        return crop_xywh(self.img, self.score_rect)
+
+    @property
+    def max_recall_rating_class_rect(self):
+        x = (
+            self.hor_mid
+            + self.sizes.COVER_RIGHT_FROM_HOR_MID
+            - self.sizes.COVER_WIDTH
+            - 25
+        )
+        return self.construct_int_xywh_rect(
+            x=x,
+            y=(
+                self.h_fixed_mid
+                - self.sizes.SCORE_PANEL[1] / 2
+                - self.sizes.MAX_RECALL_RATING_CLASS_HEIGHT
+            ),
+            w=self.sizes.MAX_RECALL_RATING_CLASS_WIDTH,
+            h=self.sizes.MAX_RECALL_RATING_CLASS_HEIGHT,
+        )
+
+    @property
+    def max_recall_rating_class(self):
+        return crop_xywh(self.img, self.max_recall_rating_class_rect)
+
+    @property
+    def title_rect(self):
+        return self.construct_int_xywh_rect(
+            x=0,
+            y=self.h_fixed_mid
+            + self.sizes.TITLE_BOTTOM_FROM_VER_MID
+            - self.sizes.TITLE_FONT_PX,
+            w=self.hor_mid + self.sizes.TITLE_WIDTH_RIGHT,
+            h=self.sizes.TITLE_FONT_PX,
+        )
+
+    @property
+    def title(self):
+        return crop_xywh(self.img, self.title_rect)
+
+    @property
+    def cover_rect(self):
+        return self.construct_int_xywh_rect(
+            x=self.hor_mid
+            + self.sizes.COVER_RIGHT_FROM_HOR_MID
+            - self.sizes.COVER_WIDTH,
+            y=self.h_fixed_mid - self.sizes.SCORE_PANEL[1] / 2,
+            w=self.sizes.COVER_WIDTH,
+            h=self.sizes.COVER_WIDTH,
+        )
+
+    @property
+    def cover(self):
+        return crop_xywh(self.img, self.cover_rect)

src/arcaea_offline_ocr/device/v2/shared.py

@@ -0,0 +1,9 @@
+from cv2 import MORPH_CROSS, MORPH_ELLIPSE, MORPH_RECT, getStructuringElement
+
+PFL_DENOISE_KERNEL = getStructuringElement(MORPH_RECT, [2, 2])
+PFL_ERODE_KERNEL = getStructuringElement(MORPH_RECT, [3, 3])
+PFL_CLOSE_HORIZONTAL_KERNEL = getStructuringElement(MORPH_RECT, [10, 1])
+
+MAX_RECALL_DENOISE_KERNEL = getStructuringElement(MORPH_RECT, [3, 3])
+MAX_RECALL_ERODE_KERNEL = getStructuringElement(MORPH_RECT, [2, 2])
+MAX_RECALL_CLOSE_KERNEL = getStructuringElement(MORPH_RECT, [20, 1])

src/arcaea_offline_ocr/mask.py

@@ -1,5 +1,5 @@
-from cv2 import BORDER_CONSTANT, BORDER_ISOLATED, bitwise_or, dilate, inRange
-from numpy import array, uint8
+import cv2
+from numpy import array, max, min, uint8
 
 from .types import Mat
 
@@ -26,7 +26,10 @@ __all__ = [
 ]
 
 GRAY_MIN_HSV = array([0, 0, 70], uint8)
-GRAY_MAX_HSV = array([0, 70, 200], uint8)
+GRAY_MAX_HSV = array([0, 0, 200], uint8)
+
+GRAY_MIN_BGR = array([50] * 3, uint8)
+GRAY_MAX_BGR = array([160] * 3, uint8)
 
 WHITE_MIN_HSV = array([0, 0, 240], uint8)
 WHITE_MAX_HSV = array([179, 10, 255], uint8)
@@ -44,39 +47,43 @@ BYD_MIN_HSV = array([170, 50, 50], uint8)
 BYD_MAX_HSV = array([179, 210, 198], uint8)
 
 
-def mask_gray(img_hsv: Mat):
-    mask = inRange(img_hsv, GRAY_MIN_HSV, GRAY_MAX_HSV)
-    mask = dilate(mask, (2, 2))
-    return mask
+def mask_gray(__img_bgr: Mat):
+    # bgr_value_equal_mask = all(__img_bgr[:, 1:] == __img_bgr[:, :-1], axis=1)
+    bgr_value_equal_mask = max(__img_bgr, axis=2) - min(__img_bgr, axis=2) <= 5
+    img_bgr = __img_bgr.copy()
+    img_bgr[~bgr_value_equal_mask] = array([0, 0, 0], __img_bgr.dtype)
+    img_bgr = cv2.erode(img_bgr, cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2)))
+    img_bgr = cv2.dilate(img_bgr, cv2.getStructuringElement(cv2.MORPH_RECT, (1, 1)))
+    return cv2.inRange(img_bgr, GRAY_MIN_BGR, GRAY_MAX_BGR)
 
 
 def mask_white(img_hsv: Mat):
-    mask = inRange(img_hsv, WHITE_MIN_HSV, WHITE_MAX_HSV)
-    mask = dilate(mask, (5, 5), borderType=BORDER_CONSTANT | BORDER_ISOLATED)
+    mask = cv2.inRange(img_hsv, WHITE_MIN_HSV, WHITE_MAX_HSV)
+    mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2)))
     return mask
 
 
 def mask_pst(img_hsv: Mat):
-    mask = inRange(img_hsv, PST_MIN_HSV, PST_MAX_HSV)
-    mask = dilate(mask, (1, 1))
+    mask = cv2.inRange(img_hsv, PST_MIN_HSV, PST_MAX_HSV)
+    mask = cv2.dilate(mask, (1, 1))
     return mask
 
 
 def mask_prs(img_hsv: Mat):
-    mask = inRange(img_hsv, PRS_MIN_HSV, PRS_MAX_HSV)
-    mask = dilate(mask, (1, 1))
+    mask = cv2.inRange(img_hsv, PRS_MIN_HSV, PRS_MAX_HSV)
+    mask = cv2.dilate(mask, (1, 1))
     return mask
 
 
 def mask_ftr(img_hsv: Mat):
-    mask = inRange(img_hsv, FTR_MIN_HSV, FTR_MAX_HSV)
-    mask = dilate(mask, (1, 1))
+    mask = cv2.inRange(img_hsv, FTR_MIN_HSV, FTR_MAX_HSV)
+    mask = cv2.dilate(mask, (1, 1))
     return mask
 
 
 def mask_byd(img_hsv: Mat):
-    mask = inRange(img_hsv, BYD_MIN_HSV, BYD_MAX_HSV)
-    mask = dilate(mask, (2, 2))
+    mask = cv2.inRange(img_hsv, BYD_MIN_HSV, BYD_MAX_HSV)
+    mask = cv2.dilate(mask, (2, 2))
     return mask
 
 
@@ -85,4 +92,4 @@ def mask_rating_class(img_hsv: Mat):
     prs = mask_prs(img_hsv)
     ftr = mask_ftr(img_hsv)
     byd = mask_byd(img_hsv)
-    return bitwise_or(byd, bitwise_or(ftr, bitwise_or(pst, prs)))
+    return cv2.bitwise_or(byd, cv2.bitwise_or(ftr, cv2.bitwise_or(pst, prs)))

src/arcaea_offline_ocr/ocr.py

@@ -1,19 +1,8 @@
 import re
 from typing import Dict, List
 
-from cv2 import (
-    CHAIN_APPROX_SIMPLE,
-    RETR_EXTERNAL,
-    TM_CCOEFF_NORMED,
-    boundingRect,
-    findContours,
-    imshow,
-    matchTemplate,
-    minMaxLoc,
-    rectangle,
-    resize,
-    waitKey,
-)
+import cv2
+import numpy as np
 from imutils import grab_contours
 from imutils import resize as imutils_resize
 from pytesseract import image_to_string
@@ -24,7 +13,7 @@ from .template import (
     load_builtin_digit_template,
     matchTemplateMultiple,
 )
-from .types import Mat
+from .types import Mat, cv2_ml_KNearest
 
 __all__ = [
     "group_numbers",
@@ -155,6 +144,18 @@ def ocr_digits(
     return int(joined_str) if joined_str else None
 
 
+def ocr_digits_knn_model(img_gray: Mat, knn_model: cv2_ml_KNearest):
+    if img_gray.shape[:2] != (20, 20):
+        img = cv2.resize(img_gray, [20, 20])
+    else:
+        img = img_gray.copy()
+
+    img = img.astype(np.float32)
+    img = img.reshape([1, -1])
+    retval, _, _, _ = knn_model.findNearest(img, 10)
+    return int(retval)
+
+
 def ocr_pure(img_masked: Mat):
     template = load_builtin_digit_template("default")
     return ocr_digits(
@@ -173,9 +174,11 @@ def ocr_score(img_cropped: Mat):
     templates = load_builtin_digit_template("default").regular
     templates_dict = dict(enumerate(templates[:10]))
 
-    cnts = findContours(img_cropped.copy(), RETR_EXTERNAL, CHAIN_APPROX_SIMPLE)
+    cnts = cv2.findContours(
+        img_cropped.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+    )
     cnts = grab_contours(cnts)
-    rects = [boundingRect(cnt) for cnt in cnts]
+    rects = [cv2.boundingRect(cnt) for cnt in cnts]
     rects = sorted(rects, key=lambda r: r[0])
 
     # debug
@@ -190,9 +193,9 @@ def ocr_score(img_cropped: Mat):
 
         digit_results: Dict[int, float] = {}
         for digit, template in templates_dict.items():
-            template = resize(template, roi.shape[::-1])
-            template_result = matchTemplate(roi, template, TM_CCOEFF_NORMED)
-            min_val, max_val, min_loc, max_loc = minMaxLoc(template_result)
+            template = cv2.resize(template, roi.shape[::-1])
+            template_result = cv2.matchTemplate(roi, template, cv2.TM_CCOEFF_NORMED)
+            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(template_result)
             digit_results[digit] = max_val
 
         digit_results = {k: v for k, v in digit_results.items() if v > 0.5}

src/arcaea_offline_ocr/recognize.py

@@ -4,12 +4,15 @@ from typing import Callable, Optional
 from cv2 import COLOR_BGR2HSV, GaussianBlur, cvtColor, imread
 
 from .crop import *
-from .device import Device
+
+# from .device import Device
 from .mask import *
 from .ocr import *
 from .types import Mat
 from .utils import imread_unicode
 
+Device = None
+
 __all__ = [
     "process_digits_ocr_img",
     "process_tesseract_ocr_img",

src/arcaea_offline_ocr/types.py

@@ -1,4 +1,42 @@
+from collections.abc import Iterable
+from typing import Any, NamedTuple, Protocol, Tuple, Union
+
 import numpy as np
 
 # from pylance
 Mat = np.ndarray[int, np.dtype[np.generic]]
+
+
+class XYWHRect(NamedTuple):
+    x: int
+    y: int
+    w: int
+    h: int
+
+    def __add__(self, other: Union["XYWHRect", Tuple[int, int, int, int]]):
+        if not isinstance(other, Iterable) or len(other) != 4:
+            raise ValueError()
+
+        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()
+
+        return self.__class__(*[a - b for a, b in zip(self, other)])
+
+
+class cv2_ml_StatModel(Protocol):
+    def predict(self, samples: np.ndarray, results: np.ndarray, flags: int = 0):
+        ...
+
+    def train(self, samples: np.ndarray, layout: int, responses: np.ndarray):
+        ...
+
+
+class cv2_ml_KNearest(cv2_ml_StatModel, Protocol):
+    def findNearest(
+        self, samples: np.ndarray, k: int
+    ) -> Tuple[Any, np.ndarray, np.ndarray, np.ndarray]:
+        """cv.ml.KNearest.findNearest(samples, k[, results[, neighborResponses[, dist]]]) -> retval, results, neighborResponses, dist"""
+        ...