From ee0943ef0bcd3fd37c9cd1ecd2db200c60db45fb Mon Sep 17 00:00:00 2001 From: 283375 Date: Sat, 5 Aug 2023 02:28:55 +0800 Subject: [PATCH] feat, wip: device V2 --- .../device/v2/definition.py | 26 ++ src/arcaea_offline_ocr/device/v2/find.py | 202 +++++++++++++ src/arcaea_offline_ocr/device/v2/ocr.py | 89 ++++++ src/arcaea_offline_ocr/device/v2/rois.py | 265 ++++++++++++++++++ src/arcaea_offline_ocr/device/v2/shared.py | 9 + 5 files changed, 591 insertions(+) create mode 100644 src/arcaea_offline_ocr/device/v2/definition.py create mode 100644 src/arcaea_offline_ocr/device/v2/find.py create mode 100644 src/arcaea_offline_ocr/device/v2/ocr.py create mode 100644 src/arcaea_offline_ocr/device/v2/rois.py create mode 100644 src/arcaea_offline_ocr/device/v2/shared.py 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])