diff --git a/src/arcaea_offline_ocr/crop.py b/src/arcaea_offline_ocr/crop.py index 95b95c9..7af0b34 100644 --- a/src/arcaea_offline_ocr/crop.py +++ b/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] diff --git a/src/arcaea_offline_ocr/device/__init__.py b/src/arcaea_offline_ocr/device/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/arcaea_offline_ocr/device/v1/crop.py b/src/arcaea_offline_ocr/device/v1/crop.py new file mode 100644 index 0000000..5977d15 --- /dev/null +++ b/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) diff --git a/src/arcaea_offline_ocr/device/v1/definition.py b/src/arcaea_offline_ocr/device/v1/definition.py new file mode 100644 index 0000000..7f12de1 --- /dev/null +++ b/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)})" diff --git a/src/arcaea_offline_ocr/device/v2/__init__.py b/src/arcaea_offline_ocr/device/v2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/arcaea_offline_ocr/device/v2/definition.py b/src/arcaea_offline_ocr/device/v2/definition.py new file mode 100644 index 0000000..31dd17a --- /dev/null +++ b/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]) diff --git a/src/arcaea_offline_ocr/device/v2/find.py b/src/arcaea_offline_ocr/device/v2/find.py new file mode 100644 index 0000000..d0ce402 --- /dev/null +++ b/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] diff --git a/src/arcaea_offline_ocr/device/v2/ocr.py b/src/arcaea_offline_ocr/device/v2/ocr.py new file mode 100644 index 0000000..5bc4634 --- /dev/null +++ b/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] diff --git a/src/arcaea_offline_ocr/device/v2/rois.py b/src/arcaea_offline_ocr/device/v2/rois.py new file mode 100644 index 0000000..7f0e285 --- /dev/null +++ b/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) diff --git a/src/arcaea_offline_ocr/device/v2/shared.py b/src/arcaea_offline_ocr/device/v2/shared.py new file mode 100644 index 0000000..bc68755 --- /dev/null +++ b/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]) diff --git a/src/arcaea_offline_ocr/mask.py b/src/arcaea_offline_ocr/mask.py index 9704082..23c9454 100644 --- a/src/arcaea_offline_ocr/mask.py +++ b/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))) diff --git a/src/arcaea_offline_ocr/ocr.py b/src/arcaea_offline_ocr/ocr.py index edffa20..ee24c88 100644 --- a/src/arcaea_offline_ocr/ocr.py +++ b/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} diff --git a/src/arcaea_offline_ocr/recognize.py b/src/arcaea_offline_ocr/recognize.py index 20a3e8a..8c666e3 100644 --- a/src/arcaea_offline_ocr/recognize.py +++ b/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", diff --git a/src/arcaea_offline_ocr/types.py b/src/arcaea_offline_ocr/types.py index 68ff8c4..3a3dc92 100644 --- a/src/arcaea_offline_ocr/types.py +++ b/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""" + ...