feat, wip: device V2

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])