From 151b1fd6d866077253d7c32ead3c17c1d771a25c Mon Sep 17 00:00:00 2001 From: 283375 Date: Fri, 29 Sep 2023 02:26:17 +0800 Subject: [PATCH 01/26] refactor: basic structure --- src/arcaea_offline_ocr/extractor/__init__.py | 0 src/arcaea_offline_ocr/extractor/common.py | 5 ++ .../extractor/sizes/__init__.py | 0 .../extractor/sizes/common.py | 15 +++++ src/arcaea_offline_ocr/masker/__init__.py | 0 src/arcaea_offline_ocr/masker/common.py | 55 +++++++++++++++++++ 6 files changed, 75 insertions(+) create mode 100644 src/arcaea_offline_ocr/extractor/__init__.py create mode 100644 src/arcaea_offline_ocr/extractor/common.py create mode 100644 src/arcaea_offline_ocr/extractor/sizes/__init__.py create mode 100644 src/arcaea_offline_ocr/extractor/sizes/common.py create mode 100644 src/arcaea_offline_ocr/masker/__init__.py create mode 100644 src/arcaea_offline_ocr/masker/common.py diff --git a/src/arcaea_offline_ocr/extractor/__init__.py b/src/arcaea_offline_ocr/extractor/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/arcaea_offline_ocr/extractor/common.py b/src/arcaea_offline_ocr/extractor/common.py new file mode 100644 index 0000000..64d5e49 --- /dev/null +++ b/src/arcaea_offline_ocr/extractor/common.py @@ -0,0 +1,5 @@ +from .sizes.common import Sizes + + +class Extractor: + sizes: Sizes diff --git a/src/arcaea_offline_ocr/extractor/sizes/__init__.py b/src/arcaea_offline_ocr/extractor/sizes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/arcaea_offline_ocr/extractor/sizes/common.py b/src/arcaea_offline_ocr/extractor/sizes/common.py new file mode 100644 index 0000000..fa85a20 --- /dev/null +++ b/src/arcaea_offline_ocr/extractor/sizes/common.py @@ -0,0 +1,15 @@ +from typing import Tuple + +Rect = Tuple[int, int, int, int] + + +class Sizes: + pure: Rect + far: Rect + lost: Rect + score: Rect + rating_class: Rect + max_recall: Rect + jacket: Rect + clear_status: Rect + partner_icon: Rect diff --git a/src/arcaea_offline_ocr/masker/__init__.py b/src/arcaea_offline_ocr/masker/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/arcaea_offline_ocr/masker/common.py b/src/arcaea_offline_ocr/masker/common.py new file mode 100644 index 0000000..8b8c867 --- /dev/null +++ b/src/arcaea_offline_ocr/masker/common.py @@ -0,0 +1,55 @@ +import cv2 + + +class Masker: + @staticmethod + def pure(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def far(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def lost(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def score(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def rating_class_pst(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def rating_class_prs(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def rating_class_ftr(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def rating_class_byd(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def max_recall(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def clear_status_track_lost(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def clear_status_track_complete(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def clear_status_full_recall(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() + + @staticmethod + def clear_status_pure_memory(roi_bgr: cv2.Mat) -> cv2.Mat: + raise NotImplementedError() From da42699ac823da83d34d45ee46c41a529e2b8af8 Mon Sep 17 00:00:00 2001 From: 283375 Date: Fri, 29 Sep 2023 21:28:43 +0800 Subject: [PATCH 02/26] refactor: auto sizes --- .../extractor/sizes/__init__.py | 1 + .../extractor/sizes/auto/__init__.py | 3 + .../extractor/sizes/auto/common.py | 7 + .../extractor/sizes/auto/t1.py | 123 +++++++++++++++++ .../extractor/sizes/auto/t2.py | 125 ++++++++++++++++++ .../extractor/sizes/custom.py | 0 6 files changed, 259 insertions(+) create mode 100644 src/arcaea_offline_ocr/extractor/sizes/auto/__init__.py create mode 100644 src/arcaea_offline_ocr/extractor/sizes/auto/common.py create mode 100644 src/arcaea_offline_ocr/extractor/sizes/auto/t1.py create mode 100644 src/arcaea_offline_ocr/extractor/sizes/auto/t2.py create mode 100644 src/arcaea_offline_ocr/extractor/sizes/custom.py diff --git a/src/arcaea_offline_ocr/extractor/sizes/__init__.py b/src/arcaea_offline_ocr/extractor/sizes/__init__.py index e69de29..4083875 100644 --- a/src/arcaea_offline_ocr/extractor/sizes/__init__.py +++ b/src/arcaea_offline_ocr/extractor/sizes/__init__.py @@ -0,0 +1 @@ +from .auto import * diff --git a/src/arcaea_offline_ocr/extractor/sizes/auto/__init__.py b/src/arcaea_offline_ocr/extractor/sizes/auto/__init__.py new file mode 100644 index 0000000..681ed11 --- /dev/null +++ b/src/arcaea_offline_ocr/extractor/sizes/auto/__init__.py @@ -0,0 +1,3 @@ +from .common import AutoSizes +from .t1 import AutoSizesT1 +from .t2 import AutoSizesT2 diff --git a/src/arcaea_offline_ocr/extractor/sizes/auto/common.py b/src/arcaea_offline_ocr/extractor/sizes/auto/common.py new file mode 100644 index 0000000..66d42e6 --- /dev/null +++ b/src/arcaea_offline_ocr/extractor/sizes/auto/common.py @@ -0,0 +1,7 @@ +from ..common import Sizes + + +class AutoSizes(Sizes): + def __init__(self, w: int, h: int): + self.w = w + self.h = h diff --git a/src/arcaea_offline_ocr/extractor/sizes/auto/t1.py b/src/arcaea_offline_ocr/extractor/sizes/auto/t1.py new file mode 100644 index 0000000..415e065 --- /dev/null +++ b/src/arcaea_offline_ocr/extractor/sizes/auto/t1.py @@ -0,0 +1,123 @@ +from .common import AutoSizes + + +class AutoSizesT1(AutoSizes): + @property + def factor(self): + return ( + ((self.w / 16) * 9) / 720 if (self.w / self.h) < (16 / 9) else self.h / 720 + ) + + @property + def w_mid(self): + return self.w / 2 + + @property + def h_mid(self): + return self.h / 2 + + @property + def top_bar(self): + return (0, 0, self.w, 50 * self.factor) + + @property + def layout_area_h_mid(self): + return self.h / 2 + self.top_bar[3] + + @property + def pfl_left_from_w_mid(self): + return 5 * self.factor + + @property + def pfl_x(self): + return self.w_mid + self.pfl_left_from_w_mid + + @property + def pfl_w(self): + return 76 * self.factor + + @property + def pfl_h(self): + return 26 * self.factor + + @property + def pure(self): + return ( + self.pfl_x, + self.layout_area_h_mid + 110 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def far(self): + return ( + self.pfl_x, + self.pure[1] + self.pure[3] + 12 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def lost(self): + return ( + self.pfl_x, + self.far[1] + self.far[3] + 10 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def score(self): + w = 280 * self.factor + h = 45 * self.factor + return ( + self.w_mid - w / 2, + self.layout_area_h_mid - 75 * self.factor - h, + w, + h, + ) + + @property + def rating_class(self): + return ( + self.w_mid - 610 * self.factor, + self.layout_area_h_mid - 180 * self.factor, + 265 * self.factor, + 35 * self.factor, + ) + + @property + def max_recall(self): + return ( + self.w_mid - 465 * self.factor, + self.layout_area_h_mid - 215 * self.factor, + 150 * self.factor, + 35 * self.factor, + ) + + @property + def jacket(self): + return ( + self.w_mid - 610 * self.factor, + self.layout_area_h_mid - 143 * self.factor, + 375 * self.factor, + 375 * self.factor, + ) + + @property + def clear_status(self): + w = 550 * self.factor + h = 60 * self.factor + return ( + self.w_mid - w / 2, + self.layout_area_h_mid - 155 * self.factor - h, + w, + h, + ) + + @property + def partner_icon(self): + w = 90 * self.factor + h = 75 * self.factor + return (self.w_mid - w / 2, 0, w, h) diff --git a/src/arcaea_offline_ocr/extractor/sizes/auto/t2.py b/src/arcaea_offline_ocr/extractor/sizes/auto/t2.py new file mode 100644 index 0000000..a99107c --- /dev/null +++ b/src/arcaea_offline_ocr/extractor/sizes/auto/t2.py @@ -0,0 +1,125 @@ +from .common import AutoSizes + + +class AutoSizesT2(AutoSizes): + @property + def factor(self): + return ( + ((self.w / 16) * 9) / 1080 + if (self.w / self.h) < (16 / 9) + else self.h / 1080 + ) + + @property + def w_mid(self): + return self.w / 2 + + @property + def h_mid(self): + return self.h / 2 + + @property + def top_bar(self): + return (0, 0, self.w, 75 * self.factor) + + @property + def layout_area_h_mid(self): + return self.h / 2 + self.top_bar[3] + + @property + def pfl_mid_from_w_mid(self): + return 60 * self.factor + + @property + def pfl_x(self): + return self.w_mid + 10 * self.factor + + @property + def pfl_w(self): + return 100 * self.factor + + @property + def pfl_h(self): + return 24 * self.factor + + @property + def pure(self): + return ( + self.pfl_x, + self.layout_area_h_mid + 175 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def far(self): + return ( + self.pfl_x, + self.pure[1] + self.pure[3] + 30 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def lost(self): + return ( + self.pfl_x, + self.far[1] + self.far[3] + 35 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def score(self): + w = 420 * self.factor + h = 70 * self.factor + return ( + self.w_mid - w / 2, + self.layout_area_h_mid - 110 * self.factor - h, + w, + h, + ) + + @property + def rating_class(self): + return ( + max(0, self.w_mid - 965 * self.factor), + self.layout_area_h_mid - 330 * self.factor, + 350 * self.factor, + 110 * self.factor, + ) + + @property + def max_recall(self): + return ( + self.w_mid - 625 * self.factor, + self.layout_area_h_mid - 275 * self.factor, + 150 * self.factor, + 50 * self.factor, + ) + + @property + def jacket(self): + return ( + self.w_mid - 915 * self.factor, + self.layout_area_h_mid - 215 * self.factor, + 565 * self.factor, + 565 * self.factor, + ) + + @property + def clear_status(self): + w = 825 * self.factor + h = 90 * self.factor + return ( + self.w_mid - w / 2, + self.layout_area_h_mid - 235 * self.factor - h, + w, + h, + ) + + @property + def partner_icon(self): + w = 135 * self.factor + h = 110 * self.factor + return (self.w_mid - w / 2, 0, w, h) diff --git a/src/arcaea_offline_ocr/extractor/sizes/custom.py b/src/arcaea_offline_ocr/extractor/sizes/custom.py new file mode 100644 index 0000000..e69de29 From 580744b641db2fe068176f25b99c8020a6ae80c0 Mon Sep 17 00:00:00 2001 From: 283375 Date: Fri, 29 Sep 2023 21:40:09 +0800 Subject: [PATCH 03/26] refactor: extractor --- src/arcaea_offline_ocr/extractor/__init__.py | 2 + src/arcaea_offline_ocr/extractor/common.py | 42 +++++++++++++++++++- 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/src/arcaea_offline_ocr/extractor/__init__.py b/src/arcaea_offline_ocr/extractor/__init__.py index e69de29..c288e7f 100644 --- a/src/arcaea_offline_ocr/extractor/__init__.py +++ b/src/arcaea_offline_ocr/extractor/__init__.py @@ -0,0 +1,2 @@ +from .common import Extractor +from .sizes import * diff --git a/src/arcaea_offline_ocr/extractor/common.py b/src/arcaea_offline_ocr/extractor/common.py index 64d5e49..f723aff 100644 --- a/src/arcaea_offline_ocr/extractor/common.py +++ b/src/arcaea_offline_ocr/extractor/common.py @@ -1,5 +1,45 @@ +import cv2 + +from ..crop import crop_xywh from .sizes.common import Sizes class Extractor: - sizes: Sizes + def __init__(self, img: cv2.Mat, sizes: Sizes): + self.img = img + self.sizes = sizes + + def __construct_int_rect(self, rect): + return tuple(round(r) for r in rect) + + @property + def pure(self): + return crop_xywh(self.img, self.__construct_int_rect(self.sizes.pure)) + + @property + def far(self): + return crop_xywh(self.img, self.__construct_int_rect(self.sizes.far)) + + @property + def lost(self): + return crop_xywh(self.img, self.__construct_int_rect(self.sizes.lost)) + + @property + def score(self): + return crop_xywh(self.img, self.__construct_int_rect(self.sizes.score)) + + @property + def rating_class(self): + return crop_xywh(self.img, self.__construct_int_rect(self.sizes.rating_class)) + + @property + def max_recall(self): + return crop_xywh(self.img, self.__construct_int_rect(self.sizes.max_recall)) + + @property + def clear_status(self): + return crop_xywh(self.img, self.__construct_int_rect(self.sizes.clear_status)) + + @property + def partner_icon(self): + return crop_xywh(self.img, self.__construct_int_rect(self.sizes.partner_icon)) From 8d33491d9b2b2901ad201195edf5678a9c731efa Mon Sep 17 00:00:00 2001 From: 283375 Date: Sun, 1 Oct 2023 02:48:45 +0800 Subject: [PATCH 04/26] refactor: masker --- src/arcaea_offline_ocr/masker/__init__.py | 2 + .../masker/auto/__init__.py | 3 + src/arcaea_offline_ocr/masker/auto/common.py | 5 + src/arcaea_offline_ocr/masker/auto/t1.py | 123 +++++++++++++++++ src/arcaea_offline_ocr/masker/auto/t2.py | 128 ++++++++++++++++++ src/arcaea_offline_ocr/masker/common.py | 52 +++---- 6 files changed, 287 insertions(+), 26 deletions(-) create mode 100644 src/arcaea_offline_ocr/masker/auto/__init__.py create mode 100644 src/arcaea_offline_ocr/masker/auto/common.py create mode 100644 src/arcaea_offline_ocr/masker/auto/t1.py create mode 100644 src/arcaea_offline_ocr/masker/auto/t2.py diff --git a/src/arcaea_offline_ocr/masker/__init__.py b/src/arcaea_offline_ocr/masker/__init__.py index e69de29..a21c811 100644 --- a/src/arcaea_offline_ocr/masker/__init__.py +++ b/src/arcaea_offline_ocr/masker/__init__.py @@ -0,0 +1,2 @@ +from .auto import AutoMasker, AutoMaskerT1, AutoMaskerT2 +from .common import Masker diff --git a/src/arcaea_offline_ocr/masker/auto/__init__.py b/src/arcaea_offline_ocr/masker/auto/__init__.py new file mode 100644 index 0000000..edc1e5a --- /dev/null +++ b/src/arcaea_offline_ocr/masker/auto/__init__.py @@ -0,0 +1,3 @@ +from .common import AutoMasker +from .t1 import AutoMaskerT1 +from .t2 import AutoMaskerT2 diff --git a/src/arcaea_offline_ocr/masker/auto/common.py b/src/arcaea_offline_ocr/masker/auto/common.py new file mode 100644 index 0000000..8ead91f --- /dev/null +++ b/src/arcaea_offline_ocr/masker/auto/common.py @@ -0,0 +1,5 @@ +from ..common import Masker + + +class AutoMasker(Masker): + ... diff --git a/src/arcaea_offline_ocr/masker/auto/t1.py b/src/arcaea_offline_ocr/masker/auto/t1.py new file mode 100644 index 0000000..ee26ecc --- /dev/null +++ b/src/arcaea_offline_ocr/masker/auto/t1.py @@ -0,0 +1,123 @@ +import cv2 +import numpy as np + +from .common import AutoMasker + +GRAY_BGR_MIN = np.array([50] * 3, np.uint8) +GRAY_BGR_MAX = np.array([160] * 3, np.uint8) + +WHITE_HSV_MIN = np.array([0, 0, 240], np.uint8) +WHITE_HSV_MAX = np.array([179, 10, 255], np.uint8) + +PST_HSV_MIN = np.array([100, 50, 80], np.uint8) +PST_HSV_MAX = np.array([100, 255, 255], np.uint8) + +PRS_HSV_MIN = np.array([43, 40, 75], np.uint8) +PRS_HSV_MAX = np.array([50, 155, 190], np.uint8) + +FTR_HSV_MIN = np.array([149, 30, 0], np.uint8) +FTR_HSV_MAX = np.array([155, 181, 150], np.uint8) + +BYD_HSV_MIN = np.array([170, 50, 50], np.uint8) +BYD_HSV_MAX = np.array([179, 210, 198], np.uint8) + +TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8) +TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8) + +TRACK_COMPLETE_HSV_MIN = np.array([140, 0, 50], np.uint8) +TRACK_COMPLETE_HSV_MAX = np.array([145, 50, 130], np.uint8) + +FULL_RECALL_HSV_MIN = np.array([140, 60, 80], np.uint8) +FULL_RECALL_HSV_MAX = np.array([150, 130, 145], np.uint8) + +PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8) +PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8) + + +class AutoMaskerT1(AutoMasker): + @classmethod + def gray(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + bgr_value_equal_mask = np.max(roi_bgr, axis=2) - np.min(roi_bgr, axis=2) <= 5 + img_bgr = roi_bgr.copy() + img_bgr[~bgr_value_equal_mask] = np.array([0, 0, 0], roi_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_BGR_MIN, GRAY_BGR_MAX) + + @classmethod + def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.gray(roi_bgr) + + @classmethod + def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.gray(roi_bgr) + + @classmethod + def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.gray(roi_bgr) + + @classmethod + def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), WHITE_HSV_MIN, WHITE_HSV_MAX + ) + + @classmethod + def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PST_HSV_MIN, PST_HSV_MAX + ) + + @classmethod + def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PRS_HSV_MIN, PRS_HSV_MAX + ) + + @classmethod + def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), FTR_HSV_MIN, FTR_HSV_MAX + ) + + @classmethod + def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), BYD_HSV_MIN, BYD_HSV_MAX + ) + + @classmethod + def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.gray(roi_bgr) + + @classmethod + def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + TRACK_LOST_HSV_MIN, + TRACK_LOST_HSV_MAX, + ) + + @classmethod + def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + TRACK_COMPLETE_HSV_MIN, + TRACK_COMPLETE_HSV_MAX, + ) + + @classmethod + def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + FULL_RECALL_HSV_MIN, + FULL_RECALL_HSV_MAX, + ) + + @classmethod + def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + PURE_MEMORY_HSV_MIN, + PURE_MEMORY_HSV_MAX, + ) diff --git a/src/arcaea_offline_ocr/masker/auto/t2.py b/src/arcaea_offline_ocr/masker/auto/t2.py new file mode 100644 index 0000000..2c1981a --- /dev/null +++ b/src/arcaea_offline_ocr/masker/auto/t2.py @@ -0,0 +1,128 @@ +import cv2 +import numpy as np + +from .common import AutoMasker + +PFL_HSV_MIN = np.array([0, 0, 248], np.uint8) +PFL_HSV_MAX = np.array([179, 10, 255], np.uint8) + +WHITE_HSV_MIN = np.array([0, 0, 240], np.uint8) +WHITE_HSV_MAX = np.array([179, 10, 255], np.uint8) + + +PST_HSV_MIN = np.array([100, 50, 80], np.uint8) +PST_HSV_MAX = np.array([100, 255, 255], np.uint8) + +PRS_HSV_MIN = np.array([43, 40, 75], np.uint8) +PRS_HSV_MAX = np.array([50, 155, 190], np.uint8) + +FTR_HSV_MIN = np.array([149, 30, 0], np.uint8) +FTR_HSV_MAX = np.array([155, 181, 150], np.uint8) + +BYD_HSV_MIN = np.array([170, 50, 50], np.uint8) +BYD_HSV_MAX = np.array([179, 210, 198], np.uint8) + +MAX_RECALL_HSV_MIN = np.array([125, 0, 0], np.uint8) +MAX_RECALL_HSV_MAX = np.array([130, 100, 150], np.uint8) + +TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8) +TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8) + +TRACK_COMPLETE_HSV_MIN = np.array([140, 0, 50], np.uint8) +TRACK_COMPLETE_HSV_MAX = np.array([145, 50, 130], np.uint8) + +FULL_RECALL_HSV_MIN = np.array([140, 60, 80], np.uint8) +FULL_RECALL_HSV_MAX = np.array([150, 130, 145], np.uint8) + +PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8) +PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8) + + +class AutoMaskerT2(AutoMasker): + @classmethod + def pfl(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PFL_HSV_MIN, PFL_HSV_MAX + ) + + @classmethod + def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.pfl(roi_bgr) + + @classmethod + def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.pfl(roi_bgr) + + @classmethod + def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.pfl(roi_bgr) + + @classmethod + def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), WHITE_HSV_MIN, WHITE_HSV_MAX + ) + + @classmethod + def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PST_HSV_MIN, PST_HSV_MAX + ) + + @classmethod + def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PRS_HSV_MIN, PRS_HSV_MAX + ) + + @classmethod + def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), FTR_HSV_MIN, FTR_HSV_MAX + ) + + @classmethod + def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), BYD_HSV_MIN, BYD_HSV_MAX + ) + + @classmethod + def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + MAX_RECALL_HSV_MIN, + MAX_RECALL_HSV_MAX, + ) + + @classmethod + def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + TRACK_LOST_HSV_MIN, + TRACK_LOST_HSV_MAX, + ) + + @classmethod + def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + TRACK_COMPLETE_HSV_MIN, + TRACK_COMPLETE_HSV_MAX, + ) + + @classmethod + def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + FULL_RECALL_HSV_MIN, + FULL_RECALL_HSV_MAX, + ) + + @classmethod + def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + PURE_MEMORY_HSV_MIN, + PURE_MEMORY_HSV_MAX, + ) diff --git a/src/arcaea_offline_ocr/masker/common.py b/src/arcaea_offline_ocr/masker/common.py index 8b8c867..fabbd55 100644 --- a/src/arcaea_offline_ocr/masker/common.py +++ b/src/arcaea_offline_ocr/masker/common.py @@ -2,54 +2,54 @@ import cv2 class Masker: - @staticmethod - def pure(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def far(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def lost(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def score(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def rating_class_pst(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def rating_class_prs(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def rating_class_ftr(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def rating_class_byd(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def max_recall(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def clear_status_track_lost(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def clear_status_track_complete(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def clear_status_full_recall(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() - @staticmethod - def clear_status_pure_memory(roi_bgr: cv2.Mat) -> cv2.Mat: + @classmethod + def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() From f7cfb841352a8d8b8d3cb9dee27e2045116fb852 Mon Sep 17 00:00:00 2001 From: 283375 Date: Sun, 1 Oct 2023 03:02:06 +0800 Subject: [PATCH 05/26] wip: DeviceOcr --- src/arcaea_offline_ocr/extractor/common.py | 4 + src/arcaea_offline_ocr/ocr_new.py | 101 +++++++++++++++++++++ 2 files changed, 105 insertions(+) create mode 100644 src/arcaea_offline_ocr/ocr_new.py diff --git a/src/arcaea_offline_ocr/extractor/common.py b/src/arcaea_offline_ocr/extractor/common.py index f723aff..01c7965 100644 --- a/src/arcaea_offline_ocr/extractor/common.py +++ b/src/arcaea_offline_ocr/extractor/common.py @@ -28,6 +28,10 @@ class Extractor: def score(self): return crop_xywh(self.img, self.__construct_int_rect(self.sizes.score)) + @property + def jacket(self): + return crop_xywh(self.img, self.__construct_int_rect(self.sizes.jacket)) + @property def rating_class(self): return crop_xywh(self.img, self.__construct_int_rect(self.sizes.rating_class)) diff --git a/src/arcaea_offline_ocr/ocr_new.py b/src/arcaea_offline_ocr/ocr_new.py new file mode 100644 index 0000000..127041c --- /dev/null +++ b/src/arcaea_offline_ocr/ocr_new.py @@ -0,0 +1,101 @@ +import cv2 +import numpy as np +from PIL import Image + +from .crop import crop_xywh +from .extractor import Extractor +from .masker import Masker +from .ocr import ( + FixRects, + ocr_digit_samples_knn, + ocr_digits_by_contour_knn, + preprocess_hog, + resize_fill_square, +) +from .phash_db import ImagePHashDatabase + + +class DeviceOcr: + def __init__( + self, + extractor: Extractor, + masker: Masker, + knn_model: cv2.ml.KNearest, + phash_db: ImagePHashDatabase, + ): + self.extractor = extractor + self.masker = masker + self.knn_model = knn_model + self.phash_db = phash_db + + def pfl(self, roi_gray: cv2.Mat, factor: float = 1.25): + contours, _ = cv2.findContours( + roi_gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE + ) + filtered_contours = [c for c in contours if cv2.contourArea(c) >= 5 * factor] + rects = [cv2.boundingRect(c) for c in filtered_contours] + rects = FixRects.connect_broken(rects, roi_gray.shape[1], roi_gray.shape[0]) + + filtered_rects = [r for r in rects if r[2] >= 5 * factor and r[3] >= 6 * factor] + filtered_rects = FixRects.split_connected(roi_gray, filtered_rects) + filtered_rects = sorted(filtered_rects, key=lambda r: r[0]) + + roi_ocr = roi_gray.copy() + filtered_contours_flattened = {tuple(c.flatten()) for c in filtered_contours} + for contour in contours: + if tuple(contour.flatten()) in filtered_contours_flattened: + continue + roi_ocr = cv2.fillPoly(roi_ocr, [contour], [0]) + digit_rois = [ + resize_fill_square(crop_xywh(roi_ocr, r), 20) + for r in sorted(filtered_rects, key=lambda r: r[0]) + ] + + samples = preprocess_hog(digit_rois) + return ocr_digit_samples_knn(samples, self.knn_model) + + def pure(self): + return self.pfl(self.masker.pure(self.extractor.pure)) + + def far(self): + return self.pfl(self.masker.far(self.extractor.far)) + + def lost(self): + return self.pfl(self.masker.lost(self.extractor.lost)) + + def score(self): + roi = self.masker.score(self.extractor.score) + contours, _ = cv2.findContours(roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) + for contour in contours: + x, y, w, h = cv2.boundingRect(contour) + if h < roi.shape[0] * 0.6: + roi = cv2.fillPoly(roi, [contour], [0]) + return ocr_digits_by_contour_knn(roi, self.knn_model) + + def rating_class(self): + roi = self.extractor.rating_class + results = [ + self.masker.rating_class_pst(roi), + self.masker.rating_class_prs(roi), + self.masker.rating_class_ftr(roi), + self.masker.rating_class_byd(roi), + ] + return max(enumerate(results), key=lambda i: np.count_nonzero(i[1]))[0] + + def max_recall(self): + return ocr_digits_by_contour_knn( + self.masker.max_recall(self.extractor.max_recall), self.knn_model + ) + + def clear_status(self): + roi = self.extractor.clear_status + results = [ + self.masker.clear_status_track_lost(roi), + self.masker.clear_status_track_complete(roi), + self.masker.clear_status_full_recall(roi), + self.masker.clear_status_pure_memory(roi), + ] + return max(enumerate(results), key=lambda i: np.count_nonzero(i[1]))[0] + + def song_id(self): + return self.phash_db.lookup_image(Image.fromarray(self.extractor.jacket))[0] From c6aba3a7e941cb5c6d039022b0d2ba34c7eedb6a Mon Sep 17 00:00:00 2001 From: 283375 Date: Sun, 1 Oct 2023 03:03:39 +0800 Subject: [PATCH 06/26] refactor: module rename --- src/arcaea_offline_ocr/{extractor => roi_extractor}/__init__.py | 0 src/arcaea_offline_ocr/{extractor => roi_extractor}/common.py | 0 .../{extractor => roi_extractor}/sizes/__init__.py | 0 .../{extractor => roi_extractor}/sizes/auto/__init__.py | 0 .../{extractor => roi_extractor}/sizes/auto/common.py | 0 .../{extractor => roi_extractor}/sizes/auto/t1.py | 0 .../{extractor => roi_extractor}/sizes/auto/t2.py | 0 .../{extractor => roi_extractor}/sizes/common.py | 0 .../{extractor => roi_extractor}/sizes/custom.py | 0 9 files changed, 0 insertions(+), 0 deletions(-) rename src/arcaea_offline_ocr/{extractor => roi_extractor}/__init__.py (100%) rename src/arcaea_offline_ocr/{extractor => roi_extractor}/common.py (100%) rename src/arcaea_offline_ocr/{extractor => roi_extractor}/sizes/__init__.py (100%) rename src/arcaea_offline_ocr/{extractor => roi_extractor}/sizes/auto/__init__.py (100%) rename src/arcaea_offline_ocr/{extractor => roi_extractor}/sizes/auto/common.py (100%) rename src/arcaea_offline_ocr/{extractor => roi_extractor}/sizes/auto/t1.py (100%) rename src/arcaea_offline_ocr/{extractor => roi_extractor}/sizes/auto/t2.py (100%) rename src/arcaea_offline_ocr/{extractor => roi_extractor}/sizes/common.py (100%) rename src/arcaea_offline_ocr/{extractor => roi_extractor}/sizes/custom.py (100%) diff --git a/src/arcaea_offline_ocr/extractor/__init__.py b/src/arcaea_offline_ocr/roi_extractor/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/extractor/__init__.py rename to src/arcaea_offline_ocr/roi_extractor/__init__.py diff --git a/src/arcaea_offline_ocr/extractor/common.py b/src/arcaea_offline_ocr/roi_extractor/common.py similarity index 100% rename from src/arcaea_offline_ocr/extractor/common.py rename to src/arcaea_offline_ocr/roi_extractor/common.py diff --git a/src/arcaea_offline_ocr/extractor/sizes/__init__.py b/src/arcaea_offline_ocr/roi_extractor/sizes/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/extractor/sizes/__init__.py rename to src/arcaea_offline_ocr/roi_extractor/sizes/__init__.py diff --git a/src/arcaea_offline_ocr/extractor/sizes/auto/__init__.py b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/extractor/sizes/auto/__init__.py rename to src/arcaea_offline_ocr/roi_extractor/sizes/auto/__init__.py diff --git a/src/arcaea_offline_ocr/extractor/sizes/auto/common.py b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/common.py similarity index 100% rename from src/arcaea_offline_ocr/extractor/sizes/auto/common.py rename to src/arcaea_offline_ocr/roi_extractor/sizes/auto/common.py diff --git a/src/arcaea_offline_ocr/extractor/sizes/auto/t1.py b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t1.py similarity index 100% rename from src/arcaea_offline_ocr/extractor/sizes/auto/t1.py rename to src/arcaea_offline_ocr/roi_extractor/sizes/auto/t1.py diff --git a/src/arcaea_offline_ocr/extractor/sizes/auto/t2.py b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t2.py similarity index 100% rename from src/arcaea_offline_ocr/extractor/sizes/auto/t2.py rename to src/arcaea_offline_ocr/roi_extractor/sizes/auto/t2.py diff --git a/src/arcaea_offline_ocr/extractor/sizes/common.py b/src/arcaea_offline_ocr/roi_extractor/sizes/common.py similarity index 100% rename from src/arcaea_offline_ocr/extractor/sizes/common.py rename to src/arcaea_offline_ocr/roi_extractor/sizes/common.py diff --git a/src/arcaea_offline_ocr/extractor/sizes/custom.py b/src/arcaea_offline_ocr/roi_extractor/sizes/custom.py similarity index 100% rename from src/arcaea_offline_ocr/extractor/sizes/custom.py rename to src/arcaea_offline_ocr/roi_extractor/sizes/custom.py From 897705d23d91ac9cd598842d9b0b36add76413e8 Mon Sep 17 00:00:00 2001 From: 283375 Date: Sun, 1 Oct 2023 03:07:20 +0800 Subject: [PATCH 07/26] refactor: class rename --- src/arcaea_offline_ocr/masker/__init__.py | 4 ++-- src/arcaea_offline_ocr/masker/auto/__init__.py | 6 +++--- src/arcaea_offline_ocr/masker/auto/common.py | 4 ++-- src/arcaea_offline_ocr/masker/auto/t1.py | 4 ++-- src/arcaea_offline_ocr/masker/auto/t2.py | 4 ++-- src/arcaea_offline_ocr/masker/common.py | 2 +- src/arcaea_offline_ocr/roi_extractor/__init__.py | 2 +- src/arcaea_offline_ocr/roi_extractor/common.py | 6 +++--- src/arcaea_offline_ocr/roi_extractor/sizes/auto/__init__.py | 6 +++--- src/arcaea_offline_ocr/roi_extractor/sizes/auto/common.py | 4 ++-- src/arcaea_offline_ocr/roi_extractor/sizes/auto/t1.py | 4 ++-- src/arcaea_offline_ocr/roi_extractor/sizes/auto/t2.py | 4 ++-- src/arcaea_offline_ocr/roi_extractor/sizes/common.py | 2 +- 13 files changed, 26 insertions(+), 26 deletions(-) diff --git a/src/arcaea_offline_ocr/masker/__init__.py b/src/arcaea_offline_ocr/masker/__init__.py index a21c811..fd4f408 100644 --- a/src/arcaea_offline_ocr/masker/__init__.py +++ b/src/arcaea_offline_ocr/masker/__init__.py @@ -1,2 +1,2 @@ -from .auto import AutoMasker, AutoMaskerT1, AutoMaskerT2 -from .common import Masker +from .auto import * +from .common import DeviceRoiMasker diff --git a/src/arcaea_offline_ocr/masker/auto/__init__.py b/src/arcaea_offline_ocr/masker/auto/__init__.py index edc1e5a..cc84af6 100644 --- a/src/arcaea_offline_ocr/masker/auto/__init__.py +++ b/src/arcaea_offline_ocr/masker/auto/__init__.py @@ -1,3 +1,3 @@ -from .common import AutoMasker -from .t1 import AutoMaskerT1 -from .t2 import AutoMaskerT2 +from .common import DeviceAutoRoiMasker +from .t1 import DeviceAutoRoiMaskerT1 +from .t2 import DeviceAutoRoiMaskerT2 diff --git a/src/arcaea_offline_ocr/masker/auto/common.py b/src/arcaea_offline_ocr/masker/auto/common.py index 8ead91f..ed725d2 100644 --- a/src/arcaea_offline_ocr/masker/auto/common.py +++ b/src/arcaea_offline_ocr/masker/auto/common.py @@ -1,5 +1,5 @@ -from ..common import Masker +from ..common import DeviceRoiMasker -class AutoMasker(Masker): +class DeviceAutoRoiMasker(DeviceRoiMasker): ... diff --git a/src/arcaea_offline_ocr/masker/auto/t1.py b/src/arcaea_offline_ocr/masker/auto/t1.py index ee26ecc..84ba4a8 100644 --- a/src/arcaea_offline_ocr/masker/auto/t1.py +++ b/src/arcaea_offline_ocr/masker/auto/t1.py @@ -1,7 +1,7 @@ import cv2 import numpy as np -from .common import AutoMasker +from .common import DeviceAutoRoiMasker GRAY_BGR_MIN = np.array([50] * 3, np.uint8) GRAY_BGR_MAX = np.array([160] * 3, np.uint8) @@ -34,7 +34,7 @@ PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8) PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8) -class AutoMaskerT1(AutoMasker): +class DeviceAutoRoiMaskerT1(DeviceAutoRoiMasker): @classmethod def gray(cls, roi_bgr: cv2.Mat) -> cv2.Mat: bgr_value_equal_mask = np.max(roi_bgr, axis=2) - np.min(roi_bgr, axis=2) <= 5 diff --git a/src/arcaea_offline_ocr/masker/auto/t2.py b/src/arcaea_offline_ocr/masker/auto/t2.py index 2c1981a..430a3fb 100644 --- a/src/arcaea_offline_ocr/masker/auto/t2.py +++ b/src/arcaea_offline_ocr/masker/auto/t2.py @@ -1,7 +1,7 @@ import cv2 import numpy as np -from .common import AutoMasker +from .common import DeviceAutoRoiMasker PFL_HSV_MIN = np.array([0, 0, 248], np.uint8) PFL_HSV_MAX = np.array([179, 10, 255], np.uint8) @@ -38,7 +38,7 @@ PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8) PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8) -class AutoMaskerT2(AutoMasker): +class DeviceAutoRoiMaskerT2(DeviceAutoRoiMasker): @classmethod def pfl(cls, roi_bgr: cv2.Mat) -> cv2.Mat: return cv2.inRange( diff --git a/src/arcaea_offline_ocr/masker/common.py b/src/arcaea_offline_ocr/masker/common.py index fabbd55..39bbaf6 100644 --- a/src/arcaea_offline_ocr/masker/common.py +++ b/src/arcaea_offline_ocr/masker/common.py @@ -1,7 +1,7 @@ import cv2 -class Masker: +class DeviceRoiMasker: @classmethod def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() diff --git a/src/arcaea_offline_ocr/roi_extractor/__init__.py b/src/arcaea_offline_ocr/roi_extractor/__init__.py index c288e7f..050378c 100644 --- a/src/arcaea_offline_ocr/roi_extractor/__init__.py +++ b/src/arcaea_offline_ocr/roi_extractor/__init__.py @@ -1,2 +1,2 @@ -from .common import Extractor +from .common import DeviceRoiExtractor from .sizes import * diff --git a/src/arcaea_offline_ocr/roi_extractor/common.py b/src/arcaea_offline_ocr/roi_extractor/common.py index 01c7965..8bb106e 100644 --- a/src/arcaea_offline_ocr/roi_extractor/common.py +++ b/src/arcaea_offline_ocr/roi_extractor/common.py @@ -1,11 +1,11 @@ import cv2 from ..crop import crop_xywh -from .sizes.common import Sizes +from .sizes.common import DeviceRoiSizes -class Extractor: - def __init__(self, img: cv2.Mat, sizes: Sizes): +class DeviceRoiExtractor: + def __init__(self, img: cv2.Mat, sizes: DeviceRoiSizes): self.img = img self.sizes = sizes diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/__init__.py b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/__init__.py index 681ed11..270a1e0 100644 --- a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/__init__.py +++ b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/__init__.py @@ -1,3 +1,3 @@ -from .common import AutoSizes -from .t1 import AutoSizesT1 -from .t2 import AutoSizesT2 +from .common import DeviceAutoRoiSizes +from .t1 import DeviceAutoRoiSizesT1 +from .t2 import DeviceAutoRoiSizesT2 diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/common.py b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/common.py index 66d42e6..a6e9c11 100644 --- a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/common.py +++ b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/common.py @@ -1,7 +1,7 @@ -from ..common import Sizes +from ..common import DeviceRoiSizes -class AutoSizes(Sizes): +class DeviceAutoRoiSizes(DeviceRoiSizes): def __init__(self, w: int, h: int): self.w = w self.h = h diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t1.py b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t1.py index 415e065..291a29c 100644 --- a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t1.py +++ b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t1.py @@ -1,7 +1,7 @@ -from .common import AutoSizes +from .common import DeviceAutoRoiSizes -class AutoSizesT1(AutoSizes): +class DeviceAutoRoiSizesT1(DeviceAutoRoiSizes): @property def factor(self): return ( diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t2.py b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t2.py index a99107c..f214da0 100644 --- a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t2.py +++ b/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t2.py @@ -1,7 +1,7 @@ -from .common import AutoSizes +from .common import DeviceAutoRoiSizes -class AutoSizesT2(AutoSizes): +class DeviceAutoRoiSizesT2(DeviceAutoRoiSizes): @property def factor(self): return ( diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/common.py b/src/arcaea_offline_ocr/roi_extractor/sizes/common.py index fa85a20..f4ef4a1 100644 --- a/src/arcaea_offline_ocr/roi_extractor/sizes/common.py +++ b/src/arcaea_offline_ocr/roi_extractor/sizes/common.py @@ -3,7 +3,7 @@ from typing import Tuple Rect = Tuple[int, int, int, int] -class Sizes: +class DeviceRoiSizes: pure: Rect far: Rect lost: Rect From 2b01f68a739be5b88f01e4204c76db67cbce6876 Mon Sep 17 00:00:00 2001 From: 283375 Date: Sun, 1 Oct 2023 03:23:43 +0800 Subject: [PATCH 08/26] refactor: replace `device` structure --- .../device/{shared.py => common.py} | 0 .../{ocr_new.py => device/ocr.py} | 14 +- .../roi/definitions}/__init__.py | 0 .../roi/definitions}/auto/__init__.py | 0 .../roi/definitions}/auto/common.py | 0 .../roi/definitions}/auto/t1.py | 0 .../roi/definitions}/auto/t2.py | 0 .../roi/definitions}/common.py | 0 .../roi/definitions}/custom.py | 0 .../roi/extractor}/__init__.py | 1 - .../roi/extractor}/common.py | 4 +- .../{ => device/roi}/masker/__init__.py | 0 .../{ => device/roi}/masker/auto/__init__.py | 0 .../{ => device/roi}/masker/auto/common.py | 0 .../{ => device/roi}/masker/auto/t1.py | 0 .../{ => device/roi}/masker/auto/t2.py | 0 .../{ => device/roi}/masker/common.py | 0 src/arcaea_offline_ocr/device/v1/__init__.py | 0 src/arcaea_offline_ocr/device/v1/crop.py | 53 ---- .../device/v1/definition.py | 37 --- src/arcaea_offline_ocr/device/v1/ocr.py | 86 ------ src/arcaea_offline_ocr/device/v2/__init__.py | 4 - .../device/v2/definition.py | 26 -- src/arcaea_offline_ocr/device/v2/ocr.py | 172 ------------ .../device/v2/preprocess.py | 54 ---- src/arcaea_offline_ocr/device/v2/rois.py | 199 -------------- src/arcaea_offline_ocr/device/v2/shared.py | 9 - src/arcaea_offline_ocr/device/v2/sizes.py | 254 ------------------ 28 files changed, 9 insertions(+), 904 deletions(-) rename src/arcaea_offline_ocr/device/{shared.py => common.py} (100%) rename src/arcaea_offline_ocr/{ocr_new.py => device/ocr.py} (93%) rename src/arcaea_offline_ocr/{roi_extractor/sizes => device/roi/definitions}/__init__.py (100%) rename src/arcaea_offline_ocr/{roi_extractor/sizes => device/roi/definitions}/auto/__init__.py (100%) rename src/arcaea_offline_ocr/{roi_extractor/sizes => device/roi/definitions}/auto/common.py (100%) rename src/arcaea_offline_ocr/{roi_extractor/sizes => device/roi/definitions}/auto/t1.py (100%) rename src/arcaea_offline_ocr/{roi_extractor/sizes => device/roi/definitions}/auto/t2.py (100%) rename src/arcaea_offline_ocr/{roi_extractor/sizes => device/roi/definitions}/common.py (100%) rename src/arcaea_offline_ocr/{roi_extractor/sizes => device/roi/definitions}/custom.py (100%) rename src/arcaea_offline_ocr/{roi_extractor => device/roi/extractor}/__init__.py (65%) rename src/arcaea_offline_ocr/{roi_extractor => device/roi/extractor}/common.py (94%) rename src/arcaea_offline_ocr/{ => device/roi}/masker/__init__.py (100%) rename src/arcaea_offline_ocr/{ => device/roi}/masker/auto/__init__.py (100%) rename src/arcaea_offline_ocr/{ => device/roi}/masker/auto/common.py (100%) rename src/arcaea_offline_ocr/{ => device/roi}/masker/auto/t1.py (100%) rename src/arcaea_offline_ocr/{ => device/roi}/masker/auto/t2.py (100%) rename src/arcaea_offline_ocr/{ => device/roi}/masker/common.py (100%) delete mode 100644 src/arcaea_offline_ocr/device/v1/__init__.py delete mode 100644 src/arcaea_offline_ocr/device/v1/crop.py delete mode 100644 src/arcaea_offline_ocr/device/v1/definition.py delete mode 100644 src/arcaea_offline_ocr/device/v1/ocr.py delete mode 100644 src/arcaea_offline_ocr/device/v2/__init__.py delete mode 100644 src/arcaea_offline_ocr/device/v2/definition.py delete mode 100644 src/arcaea_offline_ocr/device/v2/ocr.py delete mode 100644 src/arcaea_offline_ocr/device/v2/preprocess.py delete mode 100644 src/arcaea_offline_ocr/device/v2/rois.py delete mode 100644 src/arcaea_offline_ocr/device/v2/shared.py delete mode 100644 src/arcaea_offline_ocr/device/v2/sizes.py diff --git a/src/arcaea_offline_ocr/device/shared.py b/src/arcaea_offline_ocr/device/common.py similarity index 100% rename from src/arcaea_offline_ocr/device/shared.py rename to src/arcaea_offline_ocr/device/common.py diff --git a/src/arcaea_offline_ocr/ocr_new.py b/src/arcaea_offline_ocr/device/ocr.py similarity index 93% rename from src/arcaea_offline_ocr/ocr_new.py rename to src/arcaea_offline_ocr/device/ocr.py index 127041c..bf5b9a6 100644 --- a/src/arcaea_offline_ocr/ocr_new.py +++ b/src/arcaea_offline_ocr/device/ocr.py @@ -2,24 +2,24 @@ import cv2 import numpy as np from PIL import Image -from .crop import crop_xywh -from .extractor import Extractor -from .masker import Masker -from .ocr import ( +from ..crop import crop_xywh +from ..ocr import ( FixRects, ocr_digit_samples_knn, ocr_digits_by_contour_knn, preprocess_hog, resize_fill_square, ) -from .phash_db import ImagePHashDatabase +from ..phash_db import ImagePHashDatabase +from .roi.extractor import DeviceRoiExtractor +from .roi.masker import DeviceRoiMasker class DeviceOcr: def __init__( self, - extractor: Extractor, - masker: Masker, + extractor: DeviceRoiExtractor, + masker: DeviceRoiMasker, knn_model: cv2.ml.KNearest, phash_db: ImagePHashDatabase, ): diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/__init__.py b/src/arcaea_offline_ocr/device/roi/definitions/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/__init__.py rename to src/arcaea_offline_ocr/device/roi/definitions/__init__.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/__init__.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/auto/__init__.py rename to src/arcaea_offline_ocr/device/roi/definitions/auto/__init__.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/common.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/common.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/auto/common.py rename to src/arcaea_offline_ocr/device/roi/definitions/auto/common.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t1.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/t1.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/auto/t1.py rename to src/arcaea_offline_ocr/device/roi/definitions/auto/t1.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t2.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/t2.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/auto/t2.py rename to src/arcaea_offline_ocr/device/roi/definitions/auto/t2.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/common.py b/src/arcaea_offline_ocr/device/roi/definitions/common.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/common.py rename to src/arcaea_offline_ocr/device/roi/definitions/common.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/custom.py b/src/arcaea_offline_ocr/device/roi/definitions/custom.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/custom.py rename to src/arcaea_offline_ocr/device/roi/definitions/custom.py diff --git a/src/arcaea_offline_ocr/roi_extractor/__init__.py b/src/arcaea_offline_ocr/device/roi/extractor/__init__.py similarity index 65% rename from src/arcaea_offline_ocr/roi_extractor/__init__.py rename to src/arcaea_offline_ocr/device/roi/extractor/__init__.py index 050378c..66ae350 100644 --- a/src/arcaea_offline_ocr/roi_extractor/__init__.py +++ b/src/arcaea_offline_ocr/device/roi/extractor/__init__.py @@ -1,2 +1 @@ from .common import DeviceRoiExtractor -from .sizes import * diff --git a/src/arcaea_offline_ocr/roi_extractor/common.py b/src/arcaea_offline_ocr/device/roi/extractor/common.py similarity index 94% rename from src/arcaea_offline_ocr/roi_extractor/common.py rename to src/arcaea_offline_ocr/device/roi/extractor/common.py index 8bb106e..d9365e2 100644 --- a/src/arcaea_offline_ocr/roi_extractor/common.py +++ b/src/arcaea_offline_ocr/device/roi/extractor/common.py @@ -1,7 +1,7 @@ import cv2 -from ..crop import crop_xywh -from .sizes.common import DeviceRoiSizes +from ....crop import crop_xywh +from ..definitions.common import DeviceRoiSizes class DeviceRoiExtractor: diff --git a/src/arcaea_offline_ocr/masker/__init__.py b/src/arcaea_offline_ocr/device/roi/masker/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/masker/__init__.py rename to src/arcaea_offline_ocr/device/roi/masker/__init__.py diff --git a/src/arcaea_offline_ocr/masker/auto/__init__.py b/src/arcaea_offline_ocr/device/roi/masker/auto/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/masker/auto/__init__.py rename to src/arcaea_offline_ocr/device/roi/masker/auto/__init__.py diff --git a/src/arcaea_offline_ocr/masker/auto/common.py b/src/arcaea_offline_ocr/device/roi/masker/auto/common.py similarity index 100% rename from src/arcaea_offline_ocr/masker/auto/common.py rename to src/arcaea_offline_ocr/device/roi/masker/auto/common.py diff --git a/src/arcaea_offline_ocr/masker/auto/t1.py b/src/arcaea_offline_ocr/device/roi/masker/auto/t1.py similarity index 100% rename from src/arcaea_offline_ocr/masker/auto/t1.py rename to src/arcaea_offline_ocr/device/roi/masker/auto/t1.py diff --git a/src/arcaea_offline_ocr/masker/auto/t2.py b/src/arcaea_offline_ocr/device/roi/masker/auto/t2.py similarity index 100% rename from src/arcaea_offline_ocr/masker/auto/t2.py rename to src/arcaea_offline_ocr/device/roi/masker/auto/t2.py diff --git a/src/arcaea_offline_ocr/masker/common.py b/src/arcaea_offline_ocr/device/roi/masker/common.py similarity index 100% rename from src/arcaea_offline_ocr/masker/common.py rename to src/arcaea_offline_ocr/device/roi/masker/common.py diff --git a/src/arcaea_offline_ocr/device/v1/__init__.py b/src/arcaea_offline_ocr/device/v1/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/arcaea_offline_ocr/device/v1/crop.py b/src/arcaea_offline_ocr/device/v1/crop.py deleted file mode 100644 index 2f7d6a2..0000000 --- a/src/arcaea_offline_ocr/device/v1/crop.py +++ /dev/null @@ -1,53 +0,0 @@ -from typing import Tuple - -from ...types import Mat -from .definition import DeviceV1 - -__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", -] - - -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: DeviceV1): - return crop_from_device_attr(screenshot, device.pure) - - -def crop_to_far(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.far) - - -def crop_to_lost(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.lost) - - -def crop_to_max_recall(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.max_recall) - - -def crop_to_rating_class(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.rating_class) - - -def crop_to_score(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.score) - - -def crop_to_title(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.title) diff --git a/src/arcaea_offline_ocr/device/v1/definition.py b/src/arcaea_offline_ocr/device/v1/definition.py deleted file mode 100644 index 51ca29c..0000000 --- a/src/arcaea_offline_ocr/device/v1/definition.py +++ /dev/null @@ -1,37 +0,0 @@ -from dataclasses import dataclass -from typing import Any, Dict, Tuple - -__all__ = ["DeviceV1"] - - -@dataclass(kw_only=True) -class DeviceV1: - 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/v1/ocr.py b/src/arcaea_offline_ocr/device/v1/ocr.py deleted file mode 100644 index 2e3a4b1..0000000 --- a/src/arcaea_offline_ocr/device/v1/ocr.py +++ /dev/null @@ -1,86 +0,0 @@ -from typing import List - -import cv2 - -from ...crop import crop_xywh -from ...mask import mask_gray, mask_white -from ...ocr import ocr_digits_by_contour_knn, ocr_rating_class -from ...types import Mat, cv2_ml_KNearest -from ..shared import DeviceOcrResult -from .crop import * -from .definition import DeviceV1 - - -class DeviceV1Ocr: - def __init__(self, device: DeviceV1, knn_model: cv2_ml_KNearest): - self.__device = device - self.__knn_model = knn_model - - @property - def device(self): - return self.__device - - @device.setter - def device(self, value): - self.__device = value - - @property - def knn_model(self): - return self.__knn_model - - @knn_model.setter - def knn_model(self, value): - self.__knn_model = value - - def preprocess_score_roi(self, __roi_gray: Mat) -> List[Mat]: - roi_gray = __roi_gray.copy() - contours, _ = cv2.findContours( - roi_gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE - ) - for contour in contours: - rect = cv2.boundingRect(contour) - if rect[3] > roi_gray.shape[0] * 0.6: - continue - roi_gray = cv2.fillPoly(roi_gray, [contour], 0) - return roi_gray - - def ocr(self, img_bgr: Mat): - rating_class_roi = crop_to_rating_class(img_bgr, self.device) - rating_class = ocr_rating_class(rating_class_roi) - - pfl_mr_roi = [ - crop_to_pure(img_bgr, self.device), - crop_to_far(img_bgr, self.device), - crop_to_lost(img_bgr, self.device), - crop_to_max_recall(img_bgr, self.device), - ] - pfl_mr_roi = [mask_gray(roi) for roi in pfl_mr_roi] - - pure, far, lost = [ - ocr_digits_by_contour_knn(roi, self.knn_model) for roi in pfl_mr_roi[:3] - ] - - max_recall_contours, _ = cv2.findContours( - pfl_mr_roi[3], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE - ) - max_recall_rects = [cv2.boundingRect(c) for c in max_recall_contours] - max_recall_rect = sorted(max_recall_rects, key=lambda r: r[0])[-1] - max_recall_roi = crop_xywh(img_bgr, max_recall_rect) - max_recall = ocr_digits_by_contour_knn(max_recall_roi, self.knn_model) - - score_roi = crop_to_score(img_bgr, self.device) - score_roi = mask_white(score_roi) - score_roi = self.preprocess_score_roi(score_roi) - score = ocr_digits_by_contour_knn(score_roi, self.knn_model) - - return DeviceOcrResult( - song_id=None, - title=None, - rating_class=rating_class, - pure=pure, - far=far, - lost=lost, - score=score, - max_recall=max_recall, - clear_type=None, - ) diff --git a/src/arcaea_offline_ocr/device/v2/__init__.py b/src/arcaea_offline_ocr/device/v2/__init__.py deleted file mode 100644 index 64db9c3..0000000 --- a/src/arcaea_offline_ocr/device/v2/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .definition import DeviceV2 -from .ocr import DeviceV2Ocr -from .rois import DeviceV2AutoRois, DeviceV2Rois -from .shared import MAX_RECALL_CLOSE_KERNEL diff --git a/src/arcaea_offline_ocr/device/v2/definition.py b/src/arcaea_offline_ocr/device/v2/definition.py deleted file mode 100644 index 31dd17a..0000000 --- a/src/arcaea_offline_ocr/device/v2/definition.py +++ /dev/null @@ -1,26 +0,0 @@ -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/ocr.py b/src/arcaea_offline_ocr/device/v2/ocr.py deleted file mode 100644 index bff9c6b..0000000 --- a/src/arcaea_offline_ocr/device/v2/ocr.py +++ /dev/null @@ -1,172 +0,0 @@ -import math -from functools import lru_cache -from typing import Sequence - -import cv2 -import numpy as np -from PIL import Image - -from ...crop import crop_xywh -from ...mask import ( - mask_byd, - mask_ftr, - mask_gray, - mask_max_recall_purple, - mask_pfl_white, - mask_prs, - mask_pst, - mask_white, -) -from ...ocr import ( - FixRects, - ocr_digit_samples_knn, - ocr_digits_by_contour_knn, - preprocess_hog, - resize_fill_square, -) -from ...phash_db import ImagePHashDatabase -from ...sift_db import SIFTDatabase -from ...types import Mat, cv2_ml_KNearest -from ..shared import DeviceOcrResult -from .preprocess import find_digits_preprocess -from .rois import DeviceV2Rois -from .shared import MAX_RECALL_CLOSE_KERNEL -from .sizes import SizesV2 - - -class DeviceV2Ocr: - def __init__(self, knn_model: cv2_ml_KNearest, phash_db: ImagePHashDatabase): - self.__knn_model = knn_model - self.__phash_db = phash_db - - @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, value: cv2_ml_KNearest): - self.__knn_model = value - - @property - def phash_db(self): - if not self.__phash_db: - raise ValueError("`phash_db` unset.") - return self.__phash_db - - @phash_db.setter - def phash_db(self, value: SIFTDatabase): - self.__phash_db = value - - @lru_cache - def _get_digit_widths(self, num_list: Sequence[int], factor: float): - widths = set() - for n in num_list: - lower = math.floor(n * factor) - upper = math.ceil(n * factor) - widths.update(range(lower, upper + 1)) - return widths - - def _base_ocr_pfl(self, roi_masked: Mat, factor: float = 1.0): - contours, _ = cv2.findContours( - roi_masked, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE - ) - filtered_contours = [c for c in contours if cv2.contourArea(c) >= 5 * factor] - rects = [cv2.boundingRect(c) for c in filtered_contours] - rects = FixRects.connect_broken(rects, roi_masked.shape[1], roi_masked.shape[0]) - rect_contour_map = dict(zip(rects, filtered_contours)) - - filtered_rects = [r for r in rects if r[2] >= 5 * factor and r[3] >= 6 * factor] - filtered_rects = FixRects.split_connected(roi_masked, filtered_rects) - filtered_rects = sorted(filtered_rects, key=lambda r: r[0]) - - roi_ocr = roi_masked.copy() - filtered_contours_flattened = {tuple(c.flatten()) for c in filtered_contours} - for contour in contours: - if tuple(contour.flatten()) in filtered_contours_flattened: - continue - roi_ocr = cv2.fillPoly(roi_ocr, [contour], [0]) - digit_rois = [ - resize_fill_square(crop_xywh(roi_ocr, r), 20) - for r in sorted(filtered_rects, key=lambda r: r[0]) - ] - # [cv2.imshow(f"r{i}", r) for i, r in enumerate(digit_rois)] - # cv2.waitKey(0) - samples = preprocess_hog(digit_rois) - return ocr_digit_samples_knn(samples, self.knn_model) - - def ocr_song_id(self, rois: DeviceV2Rois): - jacket = cv2.cvtColor(rois.jacket, cv2.COLOR_BGR2GRAY) - return self.phash_db.lookup_image(Image.fromarray(jacket))[0] - - def ocr_rating_class(self, rois: DeviceV2Rois): - roi = cv2.cvtColor(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 i: np.count_nonzero(i[1]))[0] - - def ocr_score(self, rois: DeviceV2Rois): - roi = cv2.cvtColor(rois.score, cv2.COLOR_BGR2HSV) - roi = mask_white(roi) - contours, _ = cv2.findContours(roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) - for contour in contours: - x, y, w, h = cv2.boundingRect(contour) - if h < roi.shape[0] * 0.6: - roi = cv2.fillPoly(roi, [contour], [0]) - return ocr_digits_by_contour_knn(roi, self.knn_model) - - def mask_pfl(self, pfl_roi: Mat, rois: DeviceV2Rois): - return ( - mask_pfl_white(cv2.cvtColor(pfl_roi, cv2.COLOR_BGR2HSV)) - if isinstance(rois.sizes, SizesV2) - else mask_gray(pfl_roi) - ) - - def ocr_pure(self, rois: DeviceV2Rois): - roi = self.mask_pfl(rois.pure, rois) - return self._base_ocr_pfl(roi, rois.sizes.factor) - - def ocr_far(self, rois: DeviceV2Rois): - roi = self.mask_pfl(rois.far, rois) - return self._base_ocr_pfl(roi, rois.sizes.factor) - - def ocr_lost(self, rois: DeviceV2Rois): - roi = self.mask_pfl(rois.lost, rois) - return self._base_ocr_pfl(roi, rois.sizes.factor) - - def ocr_max_recall(self, rois: DeviceV2Rois): - roi = ( - mask_max_recall_purple( - cv2.cvtColor(rois.max_recall_rating_class, cv2.COLOR_BGR2HSV) - ) - if isinstance(rois.sizes, SizesV2) - else mask_gray(rois.max_recall_rating_class) - ) - roi_closed = cv2.morphologyEx(roi, cv2.MORPH_CLOSE, MAX_RECALL_CLOSE_KERNEL) - contours, _ = cv2.findContours( - roi_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE - ) - rects = sorted( - [cv2.boundingRect(c) for c in contours], key=lambda r: r[0], reverse=True - ) - max_recall_roi = crop_xywh(roi, rects[0]) - return ocr_digits_by_contour_knn(max_recall_roi, self.knn_model) - - def ocr(self, rois: DeviceV2Rois): - song_id = self.ocr_song_id(rois) - rating_class = self.ocr_rating_class(rois) - score = self.ocr_score(rois) - pure = self.ocr_pure(rois) - far = self.ocr_far(rois) - lost = self.ocr_lost(rois) - max_recall = self.ocr_max_recall(rois) - - return DeviceOcrResult( - rating_class=rating_class, - pure=pure, - far=far, - lost=lost, - score=score, - max_recall=max_recall, - song_id=song_id, - ) diff --git a/src/arcaea_offline_ocr/device/v2/preprocess.py b/src/arcaea_offline_ocr/device/v2/preprocess.py deleted file mode 100644 index deef8d5..0000000 --- a/src/arcaea_offline_ocr/device/v2/preprocess.py +++ /dev/null @@ -1,54 +0,0 @@ -import cv2 - -from ...types import Mat -from .shared import * - - -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 diff --git a/src/arcaea_offline_ocr/device/v2/rois.py b/src/arcaea_offline_ocr/device/v2/rois.py deleted file mode 100644 index 100aece..0000000 --- a/src/arcaea_offline_ocr/device/v2/rois.py +++ /dev/null @@ -1,199 +0,0 @@ -from typing import Union - -from ...crop import crop_black_edges, crop_xywh -from ...types import Mat, XYWHRect -from .definition import DeviceV2 -from .sizes import Sizes, SizesV1 - - -def to_int(num: Union[int, float]) -> int: - return round(num) - - -class DeviceV2Rois: - def __init__(self, device: DeviceV2, img: Mat): - self.device = device - self.sizes = SizesV1(self.device.factor) - self.__img = img - - @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 vmid(self): - return self.h / 2 - - @property - def w(self): - return self.img.shape[1] - - @property - def hmid(self): - return self.w / 2 - - @property - def h_without_top_bar(self): - """img_height -= top_bar_height""" - return self.h - self.sizes.TOP_BAR_HEIGHT - - @property - def h_without_top_bar_mid(self): - return self.sizes.TOP_BAR_HEIGHT + self.h_without_top_bar / 2 - - @property - def pfl_top(self): - return self.h_without_top_bar_mid + self.sizes.PFL_TOP_FROM_VMID - - @property - def pfl_left(self): - return self.hmid + self.sizes.PFL_LEFT_FROM_HMID - - @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.hmid - (self.sizes.SCORE_WIDTH / 2), - y=( - self.h_without_top_bar_mid - + self.sizes.SCORE_BOTTOM_FROM_VMID - - 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.hmid - + self.sizes.JACKET_RIGHT_FROM_HOR_MID - - self.sizes.JACKET_WIDTH - - 25 * self.sizes.factor - ) - return self.construct_int_xywh_rect( - x=x, - y=( - self.h_without_top_bar_mid - - self.sizes.SCORE_PANEL[1] / 2 - - self.sizes.MR_RT_HEIGHT - ), - w=self.sizes.MR_RT_WIDTH, - h=self.sizes.MR_RT_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_without_top_bar_mid - + self.sizes.TITLE_BOTTOM_FROM_VMID - - self.sizes.TITLE_FONT_PX, - w=self.hmid + 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 jacket_rect(self): - return self.construct_int_xywh_rect( - x=self.hmid - + self.sizes.JACKET_RIGHT_FROM_HOR_MID - - self.sizes.JACKET_WIDTH, - y=self.h_without_top_bar_mid - self.sizes.SCORE_PANEL[1] / 2, - w=self.sizes.JACKET_WIDTH, - h=self.sizes.JACKET_WIDTH, - ) - - @property - def jacket(self): - return crop_xywh(self.img, self.jacket_rect) - - -class DeviceV2AutoRois(DeviceV2Rois): - @staticmethod - def get_factor(width: int, height: int): - ratio = width / height - return ((width / 16) * 9) / 720 if ratio < (16 / 9) else height / 720 - - def __init__(self, img: Mat): - factor = self.get_factor(img.shape[1], img.shape[0]) - self.sizes = SizesV1(factor) - self.__img = None - self.img = img - - @property - def img(self): - return self.__img - - @img.setter - def img(self, img: Mat): - self.__img = crop_black_edges(img) diff --git a/src/arcaea_offline_ocr/device/v2/shared.py b/src/arcaea_offline_ocr/device/v2/shared.py deleted file mode 100644 index ca511b1..0000000 --- a/src/arcaea_offline_ocr/device/v2/shared.py +++ /dev/null @@ -1,9 +0,0 @@ -from cv2 import 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/device/v2/sizes.py b/src/arcaea_offline_ocr/device/v2/sizes.py deleted file mode 100644 index 3347cb2..0000000 --- a/src/arcaea_offline_ocr/device/v2/sizes.py +++ /dev/null @@ -1,254 +0,0 @@ -from typing import Tuple, Union - - -def apply_factor(num: Union[int, float], factor: float): - return num * factor - - -class Sizes: - def __init__(self, factor: float): - raise NotImplementedError() - - @property - def TOP_BAR_HEIGHT(self): - ... - - @property - def SCORE_PANEL(self) -> Tuple[int, int]: - ... - - @property - def PFL_TOP_FROM_VMID(self): - ... - - @property - def PFL_LEFT_FROM_HMID(self): - ... - - @property - def PFL_WIDTH(self): - ... - - @property - def PFL_FONT_PX(self): - ... - - @property - def PURE_FAR_GAP(self): - ... - - @property - def FAR_LOST_GAP(self): - ... - - @property - def SCORE_BOTTOM_FROM_VMID(self): - ... - - @property - def SCORE_FONT_PX(self): - ... - - @property - def SCORE_WIDTH(self): - ... - - @property - def JACKET_RIGHT_FROM_HOR_MID(self): - ... - - @property - def JACKET_WIDTH(self): - ... - - @property - def MR_RT_RIGHT_FROM_HMID(self): - ... - - @property - def MR_RT_WIDTH(self): - ... - - @property - def MR_RT_HEIGHT(self): - ... - - @property - def TITLE_BOTTOM_FROM_VMID(self): - ... - - @property - def TITLE_FONT_PX(self): - ... - - @property - def TITLE_WIDTH_RIGHT(self): - ... - - -class SizesV1(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_VMID(self): - return self.apply_factor(135) - - @property - def PFL_LEFT_FROM_HMID(self): - return self.apply_factor(5) - - @property - def PFL_WIDTH(self): - return self.apply_factor(76) - - @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_VMID(self): - return self.apply_factor(-50) - - @property - def SCORE_FONT_PX(self): - return self.apply_factor(45) - - @property - def SCORE_WIDTH(self): - return self.apply_factor(280) - - @property - def JACKET_RIGHT_FROM_HOR_MID(self): - return self.apply_factor(-235) - - @property - def JACKET_WIDTH(self): - return self.apply_factor(375) - - @property - def MR_RT_RIGHT_FROM_HMID(self): - return self.apply_factor(-300) - - @property - def MR_RT_WIDTH(self): - return self.apply_factor(275) - - @property - def MR_RT_HEIGHT(self): - return self.apply_factor(75) - - @property - def TITLE_BOTTOM_FROM_VMID(self): - return self.apply_factor(-265) - - @property - def TITLE_FONT_PX(self): - return self.apply_factor(40) - - @property - def TITLE_WIDTH_RIGHT(self): - return self.apply_factor(275) - - -class SizesV2(Sizes): - def __init__(self, factor: float): - self.factor = factor - - def apply_factor(self, num): - return apply_factor(num, self.factor) - - @property - def TOP_BAR_HEIGHT(self): - return self.apply_factor(50) - - @property - def SCORE_PANEL(self) -> Tuple[int, int]: - return tuple(self.apply_factor(num) for num in [447, 233]) - - @property - def PFL_TOP_FROM_VMID(self): - return self.apply_factor(142) - - @property - def PFL_LEFT_FROM_HMID(self): - return self.apply_factor(10) - - @property - def PFL_WIDTH(self): - return self.apply_factor(60) - - @property - def PFL_FONT_PX(self): - return self.apply_factor(16) - - @property - def PURE_FAR_GAP(self): - return self.apply_factor(20) - - @property - def FAR_LOST_GAP(self): - return self.apply_factor(23) - - @property - def SCORE_BOTTOM_FROM_VMID(self): - return self.apply_factor(-50) - - @property - def SCORE_FONT_PX(self): - return self.apply_factor(45) - - @property - def SCORE_WIDTH(self): - return self.apply_factor(280) - - @property - def JACKET_RIGHT_FROM_HOR_MID(self): - return self.apply_factor(-235) - - @property - def JACKET_WIDTH(self): - return self.apply_factor(375) - - @property - def MR_RT_RIGHT_FROM_HMID(self): - return self.apply_factor(-330) - - @property - def MR_RT_WIDTH(self): - return self.apply_factor(330) - - @property - def MR_RT_HEIGHT(self): - return self.apply_factor(75) - - @property - def TITLE_BOTTOM_FROM_VMID(self): - return self.apply_factor(-265) - - @property - def TITLE_FONT_PX(self): - return self.apply_factor(40) - - @property - def TITLE_WIDTH_RIGHT(self): - return self.apply_factor(275) From 0d8e4dea8e8f6014e0a34bcea9f1951b3abc1a81 Mon Sep 17 00:00:00 2001 From: 283375 Date: Sun, 1 Oct 2023 04:00:06 +0800 Subject: [PATCH 09/26] refactor: remove needless module --- src/arcaea_offline_ocr/__init__.py | 1 - src/arcaea_offline_ocr/b30/chieri/v4/ocr.py | 1 - src/arcaea_offline_ocr/mask.py | 119 -------------------- src/arcaea_offline_ocr/ocr.py | 11 -- src/arcaea_offline_ocr/sift_db.py | 110 ------------------ 5 files changed, 242 deletions(-) delete mode 100644 src/arcaea_offline_ocr/mask.py delete mode 100644 src/arcaea_offline_ocr/sift_db.py diff --git a/src/arcaea_offline_ocr/__init__.py b/src/arcaea_offline_ocr/__init__.py index 41057d6..c2e0b50 100644 --- a/src/arcaea_offline_ocr/__init__.py +++ b/src/arcaea_offline_ocr/__init__.py @@ -1,5 +1,4 @@ from .crop import * from .device import * -from .mask import * from .ocr import * from .utils import * diff --git a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py index c033088..07dd434 100644 --- a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py +++ b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py @@ -8,7 +8,6 @@ from PIL import Image from ....crop import crop_xywh from ....ocr import FixRects, ocr_digits_by_contour_knn, preprocess_hog from ....phash_db import ImagePHashDatabase -from ....sift_db import SIFTDatabase from ....types import Mat, cv2_ml_KNearest from ....utils import construct_int_xywh_rect from ...shared import B30OcrResultItem diff --git a/src/arcaea_offline_ocr/mask.py b/src/arcaea_offline_ocr/mask.py deleted file mode 100644 index 085c99d..0000000 --- a/src/arcaea_offline_ocr/mask.py +++ /dev/null @@ -1,119 +0,0 @@ -import cv2 -import numpy as np - -from .types import Mat - -__all__ = [ - "GRAY_MIN_HSV", - "GRAY_MAX_HSV", - "WHITE_MIN_HSV", - "WHITE_MAX_HSV", - "PFL_WHITE_MIN_HSV", - "PFL_WHITE_MAX_HSV", - "PST_MIN_HSV", - "PST_MAX_HSV", - "PRS_MIN_HSV", - "PRS_MAX_HSV", - "FTR_MIN_HSV", - "FTR_MAX_HSV", - "BYD_MIN_HSV", - "BYD_MAX_HSV", - "MAX_RECALL_PURPLE_MIN_HSV", - "MAX_RECALL_PURPLE_MAX_HSV", - "mask_gray", - "mask_white", - "mask_pfl_white", - "mask_pst", - "mask_prs", - "mask_ftr", - "mask_byd", - "mask_rating_class", - "mask_max_recall_purple", -] - -GRAY_MIN_HSV = np.array([0, 0, 70], np.uint8) -GRAY_MAX_HSV = np.array([0, 0, 200], np.uint8) - -GRAY_MIN_BGR = np.array([50] * 3, np.uint8) -GRAY_MAX_BGR = np.array([160] * 3, np.uint8) - -WHITE_MIN_HSV = np.array([0, 0, 240], np.uint8) -WHITE_MAX_HSV = np.array([179, 10, 255], np.uint8) - -PFL_WHITE_MIN_HSV = np.array([0, 0, 248], np.uint8) -PFL_WHITE_MAX_HSV = np.array([179, 10, 255], np.uint8) - -PST_MIN_HSV = np.array([100, 50, 80], np.uint8) -PST_MAX_HSV = np.array([100, 255, 255], np.uint8) - -PRS_MIN_HSV = np.array([43, 40, 75], np.uint8) -PRS_MAX_HSV = np.array([50, 155, 190], np.uint8) - -FTR_MIN_HSV = np.array([149, 30, 0], np.uint8) -FTR_MAX_HSV = np.array([155, 181, 150], np.uint8) - -BYD_MIN_HSV = np.array([170, 50, 50], np.uint8) -BYD_MAX_HSV = np.array([179, 210, 198], np.uint8) - -MAX_RECALL_PURPLE_MIN_HSV = np.array([125, 0, 0], np.uint8) -MAX_RECALL_PURPLE_MAX_HSV = np.array([130, 100, 150], np.uint8) - - -def mask_gray(__img_bgr: Mat): - # bgr_value_equal_mask = all(__img_bgr[:, 1:] == __img_bgr[:, :-1], axis=1) - bgr_value_equal_mask = np.max(__img_bgr, axis=2) - np.min(__img_bgr, axis=2) <= 5 - img_bgr = __img_bgr.copy() - img_bgr[~bgr_value_equal_mask] = np.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 = 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_pfl_white(img_hsv: Mat): - mask = cv2.inRange(img_hsv, PFL_WHITE_MIN_HSV, PFL_WHITE_MAX_HSV) - mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))) - return mask - - -def mask_pst(img_hsv: Mat): - mask = cv2.inRange(img_hsv, PST_MIN_HSV, PST_MAX_HSV) - mask = cv2.dilate(mask, (1, 1)) - return mask - - -def mask_prs(img_hsv: Mat): - 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 = 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 = cv2.inRange(img_hsv, BYD_MIN_HSV, BYD_MAX_HSV) - mask = cv2.dilate(mask, (2, 2)) - return mask - - -def mask_rating_class(img_hsv: Mat): - pst = mask_pst(img_hsv) - prs = mask_prs(img_hsv) - ftr = mask_ftr(img_hsv) - byd = mask_byd(img_hsv) - return cv2.bitwise_or(byd, cv2.bitwise_or(ftr, cv2.bitwise_or(pst, prs))) - - -def mask_max_recall_purple(img_hsv: Mat): - mask = cv2.inRange(img_hsv, MAX_RECALL_PURPLE_MIN_HSV, MAX_RECALL_PURPLE_MAX_HSV) - mask = cv2.dilate(mask, (2, 2)) - return mask diff --git a/src/arcaea_offline_ocr/ocr.py b/src/arcaea_offline_ocr/ocr.py index 1c9c36e..95eb0dc 100644 --- a/src/arcaea_offline_ocr/ocr.py +++ b/src/arcaea_offline_ocr/ocr.py @@ -7,7 +7,6 @@ import numpy as np from numpy.linalg import norm from .crop import crop_xywh -from .mask import mask_byd, mask_ftr, mask_prs, mask_pst from .types import Mat, cv2_ml_KNearest __all__ = [ @@ -199,13 +198,3 @@ def ocr_digits_by_contour_knn( ) -> int: samples = ocr_digits_by_contour_get_samples(__roi_gray, size) return ocr_digit_samples_knn(samples, knn_model, k) - - -def ocr_rating_class(roi_hsv: Mat): - mask_results = [ - mask_pst(roi_hsv), - mask_prs(roi_hsv), - mask_ftr(roi_hsv), - mask_byd(roi_hsv), - ] - return max(enumerate(mask_results), key=lambda e: np.count_nonzero(e[1]))[0] diff --git a/src/arcaea_offline_ocr/sift_db.py b/src/arcaea_offline_ocr/sift_db.py deleted file mode 100644 index d249dae..0000000 --- a/src/arcaea_offline_ocr/sift_db.py +++ /dev/null @@ -1,110 +0,0 @@ -import io -import sqlite3 -from gzip import GzipFile -from typing import Tuple - -import cv2 -import numpy as np - -from .types import Mat - - -class SIFTDatabase: - def __init__(self, db_path: str, load: bool = True): - self.__db_path = db_path - self.__tags = [] - self.__descriptors = [] - self.__size = None - - self.__sift = cv2.SIFT_create() - self.__bf_matcher = cv2.BFMatcher() - - if load: - self.load_db() - - @property - def db_path(self): - return self.__db_path - - @db_path.setter - def db_path(self, value): - self.__db_path = value - - @property - def tags(self): - return self.__tags - - @property - def descriptors(self): - return self.__descriptors - - @property - def size(self): - return self.__size - - @size.setter - def size(self, value: Tuple[int, int]): - self.__size = value - - @property - def sift(self): - return self.__sift - - @property - def bf_matcher(self): - return self.__bf_matcher - - def load_db(self): - conn = sqlite3.connect(self.db_path) - with conn: - cursor = conn.cursor() - - size_str = cursor.execute( - "SELECT value FROM properties WHERE id = 'size'" - ).fetchone()[0] - sizr_str_arr = size_str.split(", ") - self.size = tuple(int(s) for s in sizr_str_arr) - tag__descriptors_bytes = cursor.execute( - "SELECT tag, descriptors FROM sift" - ).fetchall() - - gzipped = int( - cursor.execute( - "SELECT value FROM properties WHERE id = 'gzip'" - ).fetchone()[0] - ) - for tag, descriptor_bytes in tag__descriptors_bytes: - buffer = io.BytesIO(descriptor_bytes) - self.tags.append(tag) - if gzipped == 0: - self.descriptors.append(np.load(buffer)) - else: - gzipped_buffer = GzipFile(None, "rb", fileobj=buffer) - self.descriptors.append(np.load(gzipped_buffer)) - - def lookup_img( - self, - __img: Mat, - *, - sift=None, - bf=None, - ) -> Tuple[str, float]: - sift = sift or self.sift - bf = bf or self.bf_matcher - - img = __img.copy() - if self.size is not None: - img = cv2.resize(img, self.size) - _, descriptors = sift.detectAndCompute(img, None) - - good_results = [] - for des in self.descriptors: - matches = bf.knnMatch(descriptors, des, k=2) - good = sum(m.distance < 0.75 * n.distance for m, n in matches) - good_results.append(good) - best_match_index = max(enumerate(good_results), key=lambda i: i[1])[0] - - return ( - self.tags[best_match_index], - good_results[best_match_index] / len(descriptors), - ) From d5ccbd5a014e6a7282fd9027f37c99b77456bbf6 Mon Sep 17 00:00:00 2001 From: 283375 Date: Tue, 3 Oct 2023 15:05:03 +0800 Subject: [PATCH 10/26] refactor: module structure --- .../device/roi/definitions/auto/__init__.py | 3 - .../device/roi/definitions/auto/common.py | 7 - .../device/roi/definitions/auto/t1.py | 123 --------- .../device/roi/definitions/auto/t2.py | 125 --------- .../device/roi/extractor/__init__.py | 1 - .../device/roi/masker/__init__.py | 2 - .../device/roi/masker/auto/__init__.py | 3 - .../device/roi/masker/auto/common.py | 5 - .../device/roi/masker/auto/t1.py | 123 --------- .../device/roi/masker/auto/t2.py | 128 --------- .../definition}/__init__.py | 0 .../device/rois/definition/auto.py | 255 ++++++++++++++++++ .../definitions => rois/definition}/common.py | 2 +- .../definitions => rois/definition}/custom.py | 0 .../device/rois/extractor/__init__.py | 1 + .../device/{roi => rois}/extractor/common.py | 6 +- .../device/rois/masker/__init__.py | 2 + .../device/rois/masker/auto.py | 254 +++++++++++++++++ .../device/{roi => rois}/masker/common.py | 2 +- 19 files changed, 517 insertions(+), 525 deletions(-) delete mode 100644 src/arcaea_offline_ocr/device/roi/definitions/auto/__init__.py delete mode 100644 src/arcaea_offline_ocr/device/roi/definitions/auto/common.py delete mode 100644 src/arcaea_offline_ocr/device/roi/definitions/auto/t1.py delete mode 100644 src/arcaea_offline_ocr/device/roi/definitions/auto/t2.py delete mode 100644 src/arcaea_offline_ocr/device/roi/extractor/__init__.py delete mode 100644 src/arcaea_offline_ocr/device/roi/masker/__init__.py delete mode 100644 src/arcaea_offline_ocr/device/roi/masker/auto/__init__.py delete mode 100644 src/arcaea_offline_ocr/device/roi/masker/auto/common.py delete mode 100644 src/arcaea_offline_ocr/device/roi/masker/auto/t1.py delete mode 100644 src/arcaea_offline_ocr/device/roi/masker/auto/t2.py rename src/arcaea_offline_ocr/device/{roi/definitions => rois/definition}/__init__.py (100%) create mode 100644 src/arcaea_offline_ocr/device/rois/definition/auto.py rename src/arcaea_offline_ocr/device/{roi/definitions => rois/definition}/common.py (91%) rename src/arcaea_offline_ocr/device/{roi/definitions => rois/definition}/custom.py (100%) create mode 100644 src/arcaea_offline_ocr/device/rois/extractor/__init__.py rename src/arcaea_offline_ocr/device/{roi => rois}/extractor/common.py (90%) create mode 100644 src/arcaea_offline_ocr/device/rois/masker/__init__.py create mode 100644 src/arcaea_offline_ocr/device/rois/masker/auto.py rename src/arcaea_offline_ocr/device/{roi => rois}/masker/common.py (98%) diff --git a/src/arcaea_offline_ocr/device/roi/definitions/auto/__init__.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/__init__.py deleted file mode 100644 index 270a1e0..0000000 --- a/src/arcaea_offline_ocr/device/roi/definitions/auto/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .common import DeviceAutoRoiSizes -from .t1 import DeviceAutoRoiSizesT1 -from .t2 import DeviceAutoRoiSizesT2 diff --git a/src/arcaea_offline_ocr/device/roi/definitions/auto/common.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/common.py deleted file mode 100644 index a6e9c11..0000000 --- a/src/arcaea_offline_ocr/device/roi/definitions/auto/common.py +++ /dev/null @@ -1,7 +0,0 @@ -from ..common import DeviceRoiSizes - - -class DeviceAutoRoiSizes(DeviceRoiSizes): - def __init__(self, w: int, h: int): - self.w = w - self.h = h diff --git a/src/arcaea_offline_ocr/device/roi/definitions/auto/t1.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/t1.py deleted file mode 100644 index 291a29c..0000000 --- a/src/arcaea_offline_ocr/device/roi/definitions/auto/t1.py +++ /dev/null @@ -1,123 +0,0 @@ -from .common import DeviceAutoRoiSizes - - -class DeviceAutoRoiSizesT1(DeviceAutoRoiSizes): - @property - def factor(self): - return ( - ((self.w / 16) * 9) / 720 if (self.w / self.h) < (16 / 9) else self.h / 720 - ) - - @property - def w_mid(self): - return self.w / 2 - - @property - def h_mid(self): - return self.h / 2 - - @property - def top_bar(self): - return (0, 0, self.w, 50 * self.factor) - - @property - def layout_area_h_mid(self): - return self.h / 2 + self.top_bar[3] - - @property - def pfl_left_from_w_mid(self): - return 5 * self.factor - - @property - def pfl_x(self): - return self.w_mid + self.pfl_left_from_w_mid - - @property - def pfl_w(self): - return 76 * self.factor - - @property - def pfl_h(self): - return 26 * self.factor - - @property - def pure(self): - return ( - self.pfl_x, - self.layout_area_h_mid + 110 * self.factor, - self.pfl_w, - self.pfl_h, - ) - - @property - def far(self): - return ( - self.pfl_x, - self.pure[1] + self.pure[3] + 12 * self.factor, - self.pfl_w, - self.pfl_h, - ) - - @property - def lost(self): - return ( - self.pfl_x, - self.far[1] + self.far[3] + 10 * self.factor, - self.pfl_w, - self.pfl_h, - ) - - @property - def score(self): - w = 280 * self.factor - h = 45 * self.factor - return ( - self.w_mid - w / 2, - self.layout_area_h_mid - 75 * self.factor - h, - w, - h, - ) - - @property - def rating_class(self): - return ( - self.w_mid - 610 * self.factor, - self.layout_area_h_mid - 180 * self.factor, - 265 * self.factor, - 35 * self.factor, - ) - - @property - def max_recall(self): - return ( - self.w_mid - 465 * self.factor, - self.layout_area_h_mid - 215 * self.factor, - 150 * self.factor, - 35 * self.factor, - ) - - @property - def jacket(self): - return ( - self.w_mid - 610 * self.factor, - self.layout_area_h_mid - 143 * self.factor, - 375 * self.factor, - 375 * self.factor, - ) - - @property - def clear_status(self): - w = 550 * self.factor - h = 60 * self.factor - return ( - self.w_mid - w / 2, - self.layout_area_h_mid - 155 * self.factor - h, - w, - h, - ) - - @property - def partner_icon(self): - w = 90 * self.factor - h = 75 * self.factor - return (self.w_mid - w / 2, 0, w, h) diff --git a/src/arcaea_offline_ocr/device/roi/definitions/auto/t2.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/t2.py deleted file mode 100644 index f214da0..0000000 --- a/src/arcaea_offline_ocr/device/roi/definitions/auto/t2.py +++ /dev/null @@ -1,125 +0,0 @@ -from .common import DeviceAutoRoiSizes - - -class DeviceAutoRoiSizesT2(DeviceAutoRoiSizes): - @property - def factor(self): - return ( - ((self.w / 16) * 9) / 1080 - if (self.w / self.h) < (16 / 9) - else self.h / 1080 - ) - - @property - def w_mid(self): - return self.w / 2 - - @property - def h_mid(self): - return self.h / 2 - - @property - def top_bar(self): - return (0, 0, self.w, 75 * self.factor) - - @property - def layout_area_h_mid(self): - return self.h / 2 + self.top_bar[3] - - @property - def pfl_mid_from_w_mid(self): - return 60 * self.factor - - @property - def pfl_x(self): - return self.w_mid + 10 * self.factor - - @property - def pfl_w(self): - return 100 * self.factor - - @property - def pfl_h(self): - return 24 * self.factor - - @property - def pure(self): - return ( - self.pfl_x, - self.layout_area_h_mid + 175 * self.factor, - self.pfl_w, - self.pfl_h, - ) - - @property - def far(self): - return ( - self.pfl_x, - self.pure[1] + self.pure[3] + 30 * self.factor, - self.pfl_w, - self.pfl_h, - ) - - @property - def lost(self): - return ( - self.pfl_x, - self.far[1] + self.far[3] + 35 * self.factor, - self.pfl_w, - self.pfl_h, - ) - - @property - def score(self): - w = 420 * self.factor - h = 70 * self.factor - return ( - self.w_mid - w / 2, - self.layout_area_h_mid - 110 * self.factor - h, - w, - h, - ) - - @property - def rating_class(self): - return ( - max(0, self.w_mid - 965 * self.factor), - self.layout_area_h_mid - 330 * self.factor, - 350 * self.factor, - 110 * self.factor, - ) - - @property - def max_recall(self): - return ( - self.w_mid - 625 * self.factor, - self.layout_area_h_mid - 275 * self.factor, - 150 * self.factor, - 50 * self.factor, - ) - - @property - def jacket(self): - return ( - self.w_mid - 915 * self.factor, - self.layout_area_h_mid - 215 * self.factor, - 565 * self.factor, - 565 * self.factor, - ) - - @property - def clear_status(self): - w = 825 * self.factor - h = 90 * self.factor - return ( - self.w_mid - w / 2, - self.layout_area_h_mid - 235 * self.factor - h, - w, - h, - ) - - @property - def partner_icon(self): - w = 135 * self.factor - h = 110 * self.factor - return (self.w_mid - w / 2, 0, w, h) diff --git a/src/arcaea_offline_ocr/device/roi/extractor/__init__.py b/src/arcaea_offline_ocr/device/roi/extractor/__init__.py deleted file mode 100644 index 66ae350..0000000 --- a/src/arcaea_offline_ocr/device/roi/extractor/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .common import DeviceRoiExtractor diff --git a/src/arcaea_offline_ocr/device/roi/masker/__init__.py b/src/arcaea_offline_ocr/device/roi/masker/__init__.py deleted file mode 100644 index fd4f408..0000000 --- a/src/arcaea_offline_ocr/device/roi/masker/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .auto import * -from .common import DeviceRoiMasker diff --git a/src/arcaea_offline_ocr/device/roi/masker/auto/__init__.py b/src/arcaea_offline_ocr/device/roi/masker/auto/__init__.py deleted file mode 100644 index cc84af6..0000000 --- a/src/arcaea_offline_ocr/device/roi/masker/auto/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .common import DeviceAutoRoiMasker -from .t1 import DeviceAutoRoiMaskerT1 -from .t2 import DeviceAutoRoiMaskerT2 diff --git a/src/arcaea_offline_ocr/device/roi/masker/auto/common.py b/src/arcaea_offline_ocr/device/roi/masker/auto/common.py deleted file mode 100644 index ed725d2..0000000 --- a/src/arcaea_offline_ocr/device/roi/masker/auto/common.py +++ /dev/null @@ -1,5 +0,0 @@ -from ..common import DeviceRoiMasker - - -class DeviceAutoRoiMasker(DeviceRoiMasker): - ... diff --git a/src/arcaea_offline_ocr/device/roi/masker/auto/t1.py b/src/arcaea_offline_ocr/device/roi/masker/auto/t1.py deleted file mode 100644 index 84ba4a8..0000000 --- a/src/arcaea_offline_ocr/device/roi/masker/auto/t1.py +++ /dev/null @@ -1,123 +0,0 @@ -import cv2 -import numpy as np - -from .common import DeviceAutoRoiMasker - -GRAY_BGR_MIN = np.array([50] * 3, np.uint8) -GRAY_BGR_MAX = np.array([160] * 3, np.uint8) - -WHITE_HSV_MIN = np.array([0, 0, 240], np.uint8) -WHITE_HSV_MAX = np.array([179, 10, 255], np.uint8) - -PST_HSV_MIN = np.array([100, 50, 80], np.uint8) -PST_HSV_MAX = np.array([100, 255, 255], np.uint8) - -PRS_HSV_MIN = np.array([43, 40, 75], np.uint8) -PRS_HSV_MAX = np.array([50, 155, 190], np.uint8) - -FTR_HSV_MIN = np.array([149, 30, 0], np.uint8) -FTR_HSV_MAX = np.array([155, 181, 150], np.uint8) - -BYD_HSV_MIN = np.array([170, 50, 50], np.uint8) -BYD_HSV_MAX = np.array([179, 210, 198], np.uint8) - -TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8) -TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8) - -TRACK_COMPLETE_HSV_MIN = np.array([140, 0, 50], np.uint8) -TRACK_COMPLETE_HSV_MAX = np.array([145, 50, 130], np.uint8) - -FULL_RECALL_HSV_MIN = np.array([140, 60, 80], np.uint8) -FULL_RECALL_HSV_MAX = np.array([150, 130, 145], np.uint8) - -PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8) -PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8) - - -class DeviceAutoRoiMaskerT1(DeviceAutoRoiMasker): - @classmethod - def gray(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - bgr_value_equal_mask = np.max(roi_bgr, axis=2) - np.min(roi_bgr, axis=2) <= 5 - img_bgr = roi_bgr.copy() - img_bgr[~bgr_value_equal_mask] = np.array([0, 0, 0], roi_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_BGR_MIN, GRAY_BGR_MAX) - - @classmethod - def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cls.gray(roi_bgr) - - @classmethod - def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cls.gray(roi_bgr) - - @classmethod - def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cls.gray(roi_bgr) - - @classmethod - def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), WHITE_HSV_MIN, WHITE_HSV_MAX - ) - - @classmethod - def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PST_HSV_MIN, PST_HSV_MAX - ) - - @classmethod - def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PRS_HSV_MIN, PRS_HSV_MAX - ) - - @classmethod - def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), FTR_HSV_MIN, FTR_HSV_MAX - ) - - @classmethod - def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), BYD_HSV_MIN, BYD_HSV_MAX - ) - - @classmethod - def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cls.gray(roi_bgr) - - @classmethod - def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), - TRACK_LOST_HSV_MIN, - TRACK_LOST_HSV_MAX, - ) - - @classmethod - def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), - TRACK_COMPLETE_HSV_MIN, - TRACK_COMPLETE_HSV_MAX, - ) - - @classmethod - def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), - FULL_RECALL_HSV_MIN, - FULL_RECALL_HSV_MAX, - ) - - @classmethod - def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), - PURE_MEMORY_HSV_MIN, - PURE_MEMORY_HSV_MAX, - ) diff --git a/src/arcaea_offline_ocr/device/roi/masker/auto/t2.py b/src/arcaea_offline_ocr/device/roi/masker/auto/t2.py deleted file mode 100644 index 430a3fb..0000000 --- a/src/arcaea_offline_ocr/device/roi/masker/auto/t2.py +++ /dev/null @@ -1,128 +0,0 @@ -import cv2 -import numpy as np - -from .common import DeviceAutoRoiMasker - -PFL_HSV_MIN = np.array([0, 0, 248], np.uint8) -PFL_HSV_MAX = np.array([179, 10, 255], np.uint8) - -WHITE_HSV_MIN = np.array([0, 0, 240], np.uint8) -WHITE_HSV_MAX = np.array([179, 10, 255], np.uint8) - - -PST_HSV_MIN = np.array([100, 50, 80], np.uint8) -PST_HSV_MAX = np.array([100, 255, 255], np.uint8) - -PRS_HSV_MIN = np.array([43, 40, 75], np.uint8) -PRS_HSV_MAX = np.array([50, 155, 190], np.uint8) - -FTR_HSV_MIN = np.array([149, 30, 0], np.uint8) -FTR_HSV_MAX = np.array([155, 181, 150], np.uint8) - -BYD_HSV_MIN = np.array([170, 50, 50], np.uint8) -BYD_HSV_MAX = np.array([179, 210, 198], np.uint8) - -MAX_RECALL_HSV_MIN = np.array([125, 0, 0], np.uint8) -MAX_RECALL_HSV_MAX = np.array([130, 100, 150], np.uint8) - -TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8) -TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8) - -TRACK_COMPLETE_HSV_MIN = np.array([140, 0, 50], np.uint8) -TRACK_COMPLETE_HSV_MAX = np.array([145, 50, 130], np.uint8) - -FULL_RECALL_HSV_MIN = np.array([140, 60, 80], np.uint8) -FULL_RECALL_HSV_MAX = np.array([150, 130, 145], np.uint8) - -PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8) -PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8) - - -class DeviceAutoRoiMaskerT2(DeviceAutoRoiMasker): - @classmethod - def pfl(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PFL_HSV_MIN, PFL_HSV_MAX - ) - - @classmethod - def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cls.pfl(roi_bgr) - - @classmethod - def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cls.pfl(roi_bgr) - - @classmethod - def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cls.pfl(roi_bgr) - - @classmethod - def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), WHITE_HSV_MIN, WHITE_HSV_MAX - ) - - @classmethod - def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PST_HSV_MIN, PST_HSV_MAX - ) - - @classmethod - def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PRS_HSV_MIN, PRS_HSV_MAX - ) - - @classmethod - def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), FTR_HSV_MIN, FTR_HSV_MAX - ) - - @classmethod - def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), BYD_HSV_MIN, BYD_HSV_MAX - ) - - @classmethod - def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), - MAX_RECALL_HSV_MIN, - MAX_RECALL_HSV_MAX, - ) - - @classmethod - def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), - TRACK_LOST_HSV_MIN, - TRACK_LOST_HSV_MAX, - ) - - @classmethod - def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), - TRACK_COMPLETE_HSV_MIN, - TRACK_COMPLETE_HSV_MAX, - ) - - @classmethod - def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), - FULL_RECALL_HSV_MIN, - FULL_RECALL_HSV_MAX, - ) - - @classmethod - def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat: - return cv2.inRange( - cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), - PURE_MEMORY_HSV_MIN, - PURE_MEMORY_HSV_MAX, - ) diff --git a/src/arcaea_offline_ocr/device/roi/definitions/__init__.py b/src/arcaea_offline_ocr/device/rois/definition/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/device/roi/definitions/__init__.py rename to src/arcaea_offline_ocr/device/rois/definition/__init__.py diff --git a/src/arcaea_offline_ocr/device/rois/definition/auto.py b/src/arcaea_offline_ocr/device/rois/definition/auto.py new file mode 100644 index 0000000..66508c9 --- /dev/null +++ b/src/arcaea_offline_ocr/device/rois/definition/auto.py @@ -0,0 +1,255 @@ +from .common import DeviceRois + +__all__ = ["DeviceRoisAuto", "DeviceRoisAutoT1", "DeviceRoisAutoT2"] + + +class DeviceRoisAuto(DeviceRois): + def __init__(self, w: int, h: int): + self.w = w + self.h = h + + +class DeviceRoisAutoT1(DeviceRoisAuto): + @property + def factor(self): + return ( + ((self.w / 16) * 9) / 720 if (self.w / self.h) < (16 / 9) else self.h / 720 + ) + + @property + def w_mid(self): + return self.w / 2 + + @property + def h_mid(self): + return self.h / 2 + + @property + def top_bar(self): + return (0, 0, self.w, 50 * self.factor) + + @property + def layout_area_h_mid(self): + return self.h / 2 + self.top_bar[3] + + @property + def pfl_left_from_w_mid(self): + return 5 * self.factor + + @property + def pfl_x(self): + return self.w_mid + self.pfl_left_from_w_mid + + @property + def pfl_w(self): + return 76 * self.factor + + @property + def pfl_h(self): + return 26 * self.factor + + @property + def pure(self): + return ( + self.pfl_x, + self.layout_area_h_mid + 110 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def far(self): + return ( + self.pfl_x, + self.pure[1] + self.pure[3] + 12 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def lost(self): + return ( + self.pfl_x, + self.far[1] + self.far[3] + 10 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def score(self): + w = 280 * self.factor + h = 45 * self.factor + return ( + self.w_mid - w / 2, + self.layout_area_h_mid - 75 * self.factor - h, + w, + h, + ) + + @property + def rating_class(self): + return ( + self.w_mid - 610 * self.factor, + self.layout_area_h_mid - 180 * self.factor, + 265 * self.factor, + 35 * self.factor, + ) + + @property + def max_recall(self): + return ( + self.w_mid - 465 * self.factor, + self.layout_area_h_mid - 215 * self.factor, + 150 * self.factor, + 35 * self.factor, + ) + + @property + def jacket(self): + return ( + self.w_mid - 610 * self.factor, + self.layout_area_h_mid - 143 * self.factor, + 375 * self.factor, + 375 * self.factor, + ) + + @property + def clear_status(self): + w = 550 * self.factor + h = 60 * self.factor + return ( + self.w_mid - w / 2, + self.layout_area_h_mid - 155 * self.factor - h, + w, + h, + ) + + @property + def partner_icon(self): + w = 90 * self.factor + h = 75 * self.factor + return (self.w_mid - w / 2, 0, w, h) + + +class DeviceRoisAutoT2(DeviceRoisAuto): + @property + def factor(self): + return ( + ((self.w / 16) * 9) / 1080 + if (self.w / self.h) < (16 / 9) + else self.h / 1080 + ) + + @property + def w_mid(self): + return self.w / 2 + + @property + def h_mid(self): + return self.h / 2 + + @property + def top_bar(self): + return (0, 0, self.w, 75 * self.factor) + + @property + def layout_area_h_mid(self): + return self.h / 2 + self.top_bar[3] + + @property + def pfl_mid_from_w_mid(self): + return 60 * self.factor + + @property + def pfl_x(self): + return self.w_mid + 10 * self.factor + + @property + def pfl_w(self): + return 100 * self.factor + + @property + def pfl_h(self): + return 24 * self.factor + + @property + def pure(self): + return ( + self.pfl_x, + self.layout_area_h_mid + 175 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def far(self): + return ( + self.pfl_x, + self.pure[1] + self.pure[3] + 30 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def lost(self): + return ( + self.pfl_x, + self.far[1] + self.far[3] + 35 * self.factor, + self.pfl_w, + self.pfl_h, + ) + + @property + def score(self): + w = 420 * self.factor + h = 70 * self.factor + return ( + self.w_mid - w / 2, + self.layout_area_h_mid - 110 * self.factor - h, + w, + h, + ) + + @property + def rating_class(self): + return ( + max(0, self.w_mid - 965 * self.factor), + self.layout_area_h_mid - 330 * self.factor, + 350 * self.factor, + 110 * self.factor, + ) + + @property + def max_recall(self): + return ( + self.w_mid - 625 * self.factor, + self.layout_area_h_mid - 275 * self.factor, + 150 * self.factor, + 50 * self.factor, + ) + + @property + def jacket(self): + return ( + self.w_mid - 915 * self.factor, + self.layout_area_h_mid - 215 * self.factor, + 565 * self.factor, + 565 * self.factor, + ) + + @property + def clear_status(self): + w = 825 * self.factor + h = 90 * self.factor + return ( + self.w_mid - w / 2, + self.layout_area_h_mid - 235 * self.factor - h, + w, + h, + ) + + @property + def partner_icon(self): + w = 135 * self.factor + h = 110 * self.factor + return (self.w_mid - w / 2, 0, w, h) diff --git a/src/arcaea_offline_ocr/device/roi/definitions/common.py b/src/arcaea_offline_ocr/device/rois/definition/common.py similarity index 91% rename from src/arcaea_offline_ocr/device/roi/definitions/common.py rename to src/arcaea_offline_ocr/device/rois/definition/common.py index f4ef4a1..96512c4 100644 --- a/src/arcaea_offline_ocr/device/roi/definitions/common.py +++ b/src/arcaea_offline_ocr/device/rois/definition/common.py @@ -3,7 +3,7 @@ from typing import Tuple Rect = Tuple[int, int, int, int] -class DeviceRoiSizes: +class DeviceRois: pure: Rect far: Rect lost: Rect diff --git a/src/arcaea_offline_ocr/device/roi/definitions/custom.py b/src/arcaea_offline_ocr/device/rois/definition/custom.py similarity index 100% rename from src/arcaea_offline_ocr/device/roi/definitions/custom.py rename to src/arcaea_offline_ocr/device/rois/definition/custom.py diff --git a/src/arcaea_offline_ocr/device/rois/extractor/__init__.py b/src/arcaea_offline_ocr/device/rois/extractor/__init__.py new file mode 100644 index 0000000..1b6ae1d --- /dev/null +++ b/src/arcaea_offline_ocr/device/rois/extractor/__init__.py @@ -0,0 +1 @@ +from .common import DeviceRoisExtractor diff --git a/src/arcaea_offline_ocr/device/roi/extractor/common.py b/src/arcaea_offline_ocr/device/rois/extractor/common.py similarity index 90% rename from src/arcaea_offline_ocr/device/roi/extractor/common.py rename to src/arcaea_offline_ocr/device/rois/extractor/common.py index d9365e2..e672e8b 100644 --- a/src/arcaea_offline_ocr/device/roi/extractor/common.py +++ b/src/arcaea_offline_ocr/device/rois/extractor/common.py @@ -1,11 +1,11 @@ import cv2 from ....crop import crop_xywh -from ..definitions.common import DeviceRoiSizes +from ..definition.common import DeviceRois -class DeviceRoiExtractor: - def __init__(self, img: cv2.Mat, sizes: DeviceRoiSizes): +class DeviceRoisExtractor: + def __init__(self, img: cv2.Mat, sizes: DeviceRois): self.img = img self.sizes = sizes diff --git a/src/arcaea_offline_ocr/device/rois/masker/__init__.py b/src/arcaea_offline_ocr/device/rois/masker/__init__.py new file mode 100644 index 0000000..ced796d --- /dev/null +++ b/src/arcaea_offline_ocr/device/rois/masker/__init__.py @@ -0,0 +1,2 @@ +from .auto import * +from .common import DeviceRoisMasker diff --git a/src/arcaea_offline_ocr/device/rois/masker/auto.py b/src/arcaea_offline_ocr/device/rois/masker/auto.py new file mode 100644 index 0000000..77f0696 --- /dev/null +++ b/src/arcaea_offline_ocr/device/rois/masker/auto.py @@ -0,0 +1,254 @@ +import cv2 +import numpy as np + +from .common import DeviceRoisMasker + + +class DeviceRoisMaskerAuto(DeviceRoisMasker): + ... + + +class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto): + GRAY_BGR_MIN = np.array([50] * 3, np.uint8) + GRAY_BGR_MAX = np.array([160] * 3, np.uint8) + + WHITE_HSV_MIN = np.array([0, 0, 240], np.uint8) + WHITE_HSV_MAX = np.array([179, 10, 255], np.uint8) + + PST_HSV_MIN = np.array([100, 50, 80], np.uint8) + PST_HSV_MAX = np.array([100, 255, 255], np.uint8) + + PRS_HSV_MIN = np.array([43, 40, 75], np.uint8) + PRS_HSV_MAX = np.array([50, 155, 190], np.uint8) + + FTR_HSV_MIN = np.array([149, 30, 0], np.uint8) + FTR_HSV_MAX = np.array([155, 181, 150], np.uint8) + + BYD_HSV_MIN = np.array([170, 50, 50], np.uint8) + BYD_HSV_MAX = np.array([179, 210, 198], np.uint8) + + TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8) + TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8) + + TRACK_COMPLETE_HSV_MIN = np.array([140, 0, 50], np.uint8) + TRACK_COMPLETE_HSV_MAX = np.array([145, 50, 130], np.uint8) + + FULL_RECALL_HSV_MIN = np.array([140, 60, 80], np.uint8) + FULL_RECALL_HSV_MAX = np.array([150, 130, 145], np.uint8) + + PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8) + PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8) + + @classmethod + def gray(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + bgr_value_equal_mask = np.max(roi_bgr, axis=2) - np.min(roi_bgr, axis=2) <= 5 + img_bgr = roi_bgr.copy() + img_bgr[~bgr_value_equal_mask] = np.array([0, 0, 0], roi_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, cls.GRAY_BGR_MIN, cls.GRAY_BGR_MAX) + + @classmethod + def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.gray(roi_bgr) + + @classmethod + def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.gray(roi_bgr) + + @classmethod + def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.gray(roi_bgr) + + @classmethod + def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + cls.WHITE_HSV_MIN, + cls.WHITE_HSV_MAX, + ) + + @classmethod + def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PST_HSV_MIN, cls.PST_HSV_MAX + ) + + @classmethod + def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PRS_HSV_MIN, cls.PRS_HSV_MAX + ) + + @classmethod + def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.FTR_HSV_MIN, cls.FTR_HSV_MAX + ) + + @classmethod + def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.BYD_HSV_MIN, cls.BYD_HSV_MAX + ) + + @classmethod + def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.gray(roi_bgr) + + @classmethod + def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + cls.TRACK_LOST_HSV_MIN, + cls.TRACK_LOST_HSV_MAX, + ) + + @classmethod + def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + cls.TRACK_COMPLETE_HSV_MIN, + cls.TRACK_COMPLETE_HSV_MAX, + ) + + @classmethod + def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + cls.FULL_RECALL_HSV_MIN, + cls.FULL_RECALL_HSV_MAX, + ) + + @classmethod + def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + cls.PURE_MEMORY_HSV_MIN, + cls.PURE_MEMORY_HSV_MAX, + ) + + +class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto): + PFL_HSV_MIN = np.array([0, 0, 248], np.uint8) + PFL_HSV_MAX = np.array([179, 10, 255], np.uint8) + + WHITE_HSV_MIN = np.array([0, 0, 240], np.uint8) + WHITE_HSV_MAX = np.array([179, 10, 255], np.uint8) + + PST_HSV_MIN = np.array([100, 50, 80], np.uint8) + PST_HSV_MAX = np.array([100, 255, 255], np.uint8) + + PRS_HSV_MIN = np.array([43, 40, 75], np.uint8) + PRS_HSV_MAX = np.array([50, 155, 190], np.uint8) + + FTR_HSV_MIN = np.array([149, 30, 0], np.uint8) + FTR_HSV_MAX = np.array([155, 181, 150], np.uint8) + + BYD_HSV_MIN = np.array([170, 50, 50], np.uint8) + BYD_HSV_MAX = np.array([179, 210, 198], np.uint8) + + MAX_RECALL_HSV_MIN = np.array([125, 0, 0], np.uint8) + MAX_RECALL_HSV_MAX = np.array([130, 100, 150], np.uint8) + + TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8) + TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8) + + TRACK_COMPLETE_HSV_MIN = np.array([140, 0, 50], np.uint8) + TRACK_COMPLETE_HSV_MAX = np.array([145, 50, 130], np.uint8) + + FULL_RECALL_HSV_MIN = np.array([140, 60, 80], np.uint8) + FULL_RECALL_HSV_MAX = np.array([150, 130, 145], np.uint8) + + PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8) + PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8) + + @classmethod + def pfl(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PFL_HSV_MIN, cls.PFL_HSV_MAX + ) + + @classmethod + def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.pfl(roi_bgr) + + @classmethod + def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.pfl(roi_bgr) + + @classmethod + def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cls.pfl(roi_bgr) + + @classmethod + def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + cls.WHITE_HSV_MIN, + cls.WHITE_HSV_MAX, + ) + + @classmethod + def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PST_HSV_MIN, cls.PST_HSV_MAX + ) + + @classmethod + def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PRS_HSV_MIN, cls.PRS_HSV_MAX + ) + + @classmethod + def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.FTR_HSV_MIN, cls.FTR_HSV_MAX + ) + + @classmethod + def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.BYD_HSV_MIN, cls.BYD_HSV_MAX + ) + + @classmethod + def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + cls.MAX_RECALL_HSV_MIN, + cls.MAX_RECALL_HSV_MAX, + ) + + @classmethod + def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + cls.TRACK_LOST_HSV_MIN, + cls.TRACK_LOST_HSV_MAX, + ) + + @classmethod + def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + cls.TRACK_COMPLETE_HSV_MIN, + cls.TRACK_COMPLETE_HSV_MAX, + ) + + @classmethod + def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + cls.FULL_RECALL_HSV_MIN, + cls.FULL_RECALL_HSV_MAX, + ) + + @classmethod + def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + return cv2.inRange( + cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), + cls.PURE_MEMORY_HSV_MIN, + cls.PURE_MEMORY_HSV_MAX, + ) diff --git a/src/arcaea_offline_ocr/device/roi/masker/common.py b/src/arcaea_offline_ocr/device/rois/masker/common.py similarity index 98% rename from src/arcaea_offline_ocr/device/roi/masker/common.py rename to src/arcaea_offline_ocr/device/rois/masker/common.py index 39bbaf6..f877e2c 100644 --- a/src/arcaea_offline_ocr/device/roi/masker/common.py +++ b/src/arcaea_offline_ocr/device/rois/masker/common.py @@ -1,7 +1,7 @@ import cv2 -class DeviceRoiMasker: +class DeviceRoisMasker: @classmethod def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: raise NotImplementedError() From c009a28f92b6329a9c193af91ec05b697ce73d26 Mon Sep 17 00:00:00 2001 From: 283375 Date: Tue, 3 Oct 2023 16:38:11 +0800 Subject: [PATCH 11/26] refactor: use opencv to calculate image phash --- src/arcaea_offline_ocr/phash_db.py | 38 ++++++++++++++++++++++++------ 1 file changed, 31 insertions(+), 7 deletions(-) diff --git a/src/arcaea_offline_ocr/phash_db.py b/src/arcaea_offline_ocr/phash_db.py index 6bbcd5b..7ffa94e 100644 --- a/src/arcaea_offline_ocr/phash_db.py +++ b/src/arcaea_offline_ocr/phash_db.py @@ -1,8 +1,32 @@ import sqlite3 -import imagehash +import cv2 import numpy as np -from PIL import Image + + +def phash_opencv(img_gray, hash_size=8, highfreq_factor=4): + # type: (cv2.Mat | np.ndarray, int, int) -> np.ndarray + """ + Perceptual Hash computation. + + Implementation follows http://www.hackerfactor.com/blog/index.php?/archives/432-Looks-Like-It.html + + Adapted from `imagehash.phash`, pure opencv implementation + + The result is slightly different from `imagehash.phash`. + """ + if hash_size < 2: + raise ValueError("Hash size must be greater than or equal to 2") + + img_size = hash_size * highfreq_factor + image = cv2.resize(img_gray, (img_size, img_size), interpolation=cv2.INTER_LANCZOS4) + image = np.float32(image) + dct = cv2.dct(image, flags=cv2.DCT_ROWS) + dct = cv2.dct(image) + dctlowfreq = dct[:hash_size, :hash_size] + med = np.median(dctlowfreq) + diff = dctlowfreq > med + return diff def hamming_distance_sql_function(user_input, db_entry) -> int: @@ -46,8 +70,8 @@ class ImagePHashDatabase: self.hashes_head = [h[: self.hashes_slice_size] for h in self.hashes] self.hashes_tail = [h[-self.hashes_slice_size :] for h in self.hashes] - def lookup_hash(self, image_hash: imagehash.ImageHash, *, limit: int = 5): - image_hash = image_hash.hash.flatten() + def lookup_hash(self, image_hash: np.ndarray, *, limit: int = 5): + image_hash = image_hash.flatten() # image_hash_head = image_hash[: self.hashes_slice_size] # image_hash_tail = image_hash[-self.hashes_slice_size :] # head_xor_results = [image_hash_head ^ h for h in self.hashes] @@ -58,8 +82,8 @@ class ImagePHashDatabase: ] return sorted(xor_results, key=lambda r: r[1])[:limit] - def lookup_image(self, pil_image: Image.Image): - image_hash = imagehash.phash( - pil_image, hash_size=self.hash_size, highfreq_factor=self.highfreq_factor + def lookup_image(self, img_gray: cv2.Mat): + image_hash = phash_opencv( + img_gray, hash_size=self.hash_size, highfreq_factor=self.highfreq_factor ) return self.lookup_hash(image_hash)[0] From 1aa71685ce7b05cb72c30fd524b62fe9d71d7622 Mon Sep 17 00:00:00 2001 From: 283375 Date: Mon, 9 Oct 2023 00:16:54 +0800 Subject: [PATCH 12/26] fix: what is this shit --- src/arcaea_offline_ocr/phash_db.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/arcaea_offline_ocr/phash_db.py b/src/arcaea_offline_ocr/phash_db.py index 7ffa94e..4d0d954 100644 --- a/src/arcaea_offline_ocr/phash_db.py +++ b/src/arcaea_offline_ocr/phash_db.py @@ -21,7 +21,6 @@ def phash_opencv(img_gray, hash_size=8, highfreq_factor=4): img_size = hash_size * highfreq_factor image = cv2.resize(img_gray, (img_size, img_size), interpolation=cv2.INTER_LANCZOS4) image = np.float32(image) - dct = cv2.dct(image, flags=cv2.DCT_ROWS) dct = cv2.dct(image) dctlowfreq = dct[:hash_size, :hash_size] med = np.median(dctlowfreq) From d13076c667f666ce395ee78d5203dfd5963f8310 Mon Sep 17 00:00:00 2001 From: 283375 Date: Tue, 10 Oct 2023 01:28:36 +0800 Subject: [PATCH 13/26] feat: split jacket & partner lookup in phash database --- src/arcaea_offline_ocr/phash_db.py | 66 +++++++++++++++++++++--------- 1 file changed, 47 insertions(+), 19 deletions(-) diff --git a/src/arcaea_offline_ocr/phash_db.py b/src/arcaea_offline_ocr/phash_db.py index 4d0d954..384b244 100644 --- a/src/arcaea_offline_ocr/phash_db.py +++ b/src/arcaea_offline_ocr/phash_db.py @@ -1,11 +1,12 @@ import sqlite3 +from typing import List, Union import cv2 import numpy as np def phash_opencv(img_gray, hash_size=8, highfreq_factor=4): - # type: (cv2.Mat | np.ndarray, int, int) -> np.ndarray + # type: (Union[cv2.Mat, np.ndarray], int, int) -> np.ndarray """ Perceptual Hash computation. @@ -53,28 +54,35 @@ class ImagePHashDatabase: ).fetchone()[0] ) - # self.conn.create_function( - # "HAMMING_DISTANCE", - # 2, - # hamming_distance_sql_function, - # deterministic=True, - # ) - - self.ids = [i[0] for i in conn.execute("SELECT id FROM hashes").fetchall()] + self.ids: List[str] = [ + i[0] for i in conn.execute("SELECT id FROM hashes").fetchall() + ] self.hashes_byte = [ i[0] for i in conn.execute("SELECT hash FROM hashes").fetchall() ] self.hashes = [np.frombuffer(hb, bool) for hb in self.hashes_byte] - self.hashes_slice_size = round(len(self.hashes_byte[0]) * 0.25) - self.hashes_head = [h[: self.hashes_slice_size] for h in self.hashes] - self.hashes_tail = [h[-self.hashes_slice_size :] for h in self.hashes] + + self.jacket_ids: List[str] = [] + self.jacket_hashes = [] + self.partner_ids: List[str] = [] + self.partner_hashes = [] + + for id, hash in zip(self.ids, self.hashes): + id_splitted = id.split("||") + if len(id_splitted) > 1 and id_splitted[0] == "partner": + self.partner_ids.append(id) + self.partner_hashes.append(hash) + else: + self.jacket_ids.append(id) + self.jacket_hashes.append(hash) + + def calculate_phash(self, img_gray: cv2.Mat): + return phash_opencv( + img_gray, hash_size=self.hash_size, highfreq_factor=self.highfreq_factor + ) def lookup_hash(self, image_hash: np.ndarray, *, limit: int = 5): image_hash = image_hash.flatten() - # image_hash_head = image_hash[: self.hashes_slice_size] - # image_hash_tail = image_hash[-self.hashes_slice_size :] - # head_xor_results = [image_hash_head ^ h for h in self.hashes] - # tail_xor_results = [image_hash_head ^ h for h in self.hashes] xor_results = [ (id, np.count_nonzero(image_hash ^ h)) for id, h in zip(self.ids, self.hashes) @@ -82,7 +90,27 @@ class ImagePHashDatabase: return sorted(xor_results, key=lambda r: r[1])[:limit] def lookup_image(self, img_gray: cv2.Mat): - image_hash = phash_opencv( - img_gray, hash_size=self.hash_size, highfreq_factor=self.highfreq_factor - ) + image_hash = self.calculate_phash(img_gray) return self.lookup_hash(image_hash)[0] + + def lookup_jackets(self, img_gray: cv2.Mat, *, limit: int = 5): + image_hash = self.calculate_phash(img_gray).flatten() + xor_results = [ + (id, np.count_nonzero(image_hash ^ h)) + for id, h in zip(self.jacket_ids, self.jacket_hashes) + ] + return sorted(xor_results, key=lambda r: r[1])[:limit] + + def lookup_jacket(self, img_gray: cv2.Mat): + return self.lookup_jackets(img_gray)[0] + + def lookup_partners(self, img_gray: cv2.Mat, *, limit: int = 5): + image_hash = self.calculate_phash(img_gray).flatten() + xor_results = [ + (id, np.count_nonzero(image_hash ^ h)) + for id, h in zip(self.partner_ids, self.partner_hashes) + ] + return sorted(xor_results, key=lambda r: r[1])[:limit] + + def lookup_partner(self, img_gray: cv2.Mat): + return self.lookup_partners(img_gray)[0] From 6a19ead8d1f8076fa453bb715c10fcd03ffa1ac3 Mon Sep 17 00:00:00 2001 From: 283375 Date: Tue, 10 Oct 2023 01:30:12 +0800 Subject: [PATCH 14/26] refactor: module structure --- src/arcaea_offline_ocr/b30/chieri/v4/ocr.py | 6 +++--- src/arcaea_offline_ocr/device/ocr.py | 13 +++++++------ src/arcaea_offline_ocr/phash_db.py | 2 +- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py index 07dd434..302b9a5 100644 --- a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py +++ b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py @@ -7,7 +7,7 @@ from PIL import Image from ....crop import crop_xywh from ....ocr import FixRects, ocr_digits_by_contour_knn, preprocess_hog -from ....phash_db import ImagePHashDatabase +from ....phash_db import ImagePhashDatabase from ....types import Mat, cv2_ml_KNearest from ....utils import construct_int_xywh_rect from ...shared import B30OcrResultItem @@ -20,7 +20,7 @@ class ChieriBotV4Ocr: self, score_knn: cv2_ml_KNearest, pfl_knn: cv2_ml_KNearest, - phash_db: ImagePHashDatabase, + phash_db: ImagePhashDatabase, factor: Optional[float] = 1.0, ): self.__score_knn = score_knn @@ -49,7 +49,7 @@ class ChieriBotV4Ocr: return self.__phash_db @phash_db.setter - def phash_db(self, phash_db: ImagePHashDatabase): + def phash_db(self, phash_db: ImagePhashDatabase): self.__phash_db = phash_db @property diff --git a/src/arcaea_offline_ocr/device/ocr.py b/src/arcaea_offline_ocr/device/ocr.py index bf5b9a6..7232bcb 100644 --- a/src/arcaea_offline_ocr/device/ocr.py +++ b/src/arcaea_offline_ocr/device/ocr.py @@ -10,18 +10,19 @@ from ..ocr import ( preprocess_hog, resize_fill_square, ) -from ..phash_db import ImagePHashDatabase -from .roi.extractor import DeviceRoiExtractor -from .roi.masker import DeviceRoiMasker +from ..phash_db import ImagePhashDatabase +from .common import DeviceOcrResult +from .rois.extractor import DeviceRoisExtractor +from .rois.masker import DeviceRoisMasker class DeviceOcr: def __init__( self, - extractor: DeviceRoiExtractor, - masker: DeviceRoiMasker, + extractor: DeviceRoisExtractor, + masker: DeviceRoisMasker, knn_model: cv2.ml.KNearest, - phash_db: ImagePHashDatabase, + phash_db: ImagePhashDatabase, ): self.extractor = extractor self.masker = masker diff --git a/src/arcaea_offline_ocr/phash_db.py b/src/arcaea_offline_ocr/phash_db.py index 384b244..a0a406d 100644 --- a/src/arcaea_offline_ocr/phash_db.py +++ b/src/arcaea_offline_ocr/phash_db.py @@ -35,7 +35,7 @@ def hamming_distance_sql_function(user_input, db_entry) -> int: ) -class ImagePHashDatabase: +class ImagePhashDatabase: def __init__(self, db_path: str): with sqlite3.connect(db_path) as conn: self.hash_size = int( From ede2b4ec516a624cc35b746bfa61e3eec9c291ad Mon Sep 17 00:00:00 2001 From: 283375 Date: Tue, 10 Oct 2023 01:45:02 +0800 Subject: [PATCH 15/26] refactor: DeviceOcr --- src/arcaea_offline_ocr/device/common.py | 8 ++-- src/arcaea_offline_ocr/device/ocr.py | 62 ++++++++++++++++++++++++- src/arcaea_offline_ocr/phash_db.py | 18 +++---- 3 files changed, 74 insertions(+), 14 deletions(-) diff --git a/src/arcaea_offline_ocr/device/common.py b/src/arcaea_offline_ocr/device/common.py index 5a48d37..69c49cb 100644 --- a/src/arcaea_offline_ocr/device/common.py +++ b/src/arcaea_offline_ocr/device/common.py @@ -10,7 +10,9 @@ class DeviceOcrResult: far: int lost: int score: int - max_recall: int + max_recall: Optional[int] = None song_id: Optional[str] = None - title: Optional[str] = None - clear_type: Optional[str] = None + song_id_possibility: Optional[float] = None + clear_status: Optional[str] = None + partner_id: Optional[str] = None + partner_id_possibility: Optional[float] = None diff --git a/src/arcaea_offline_ocr/device/ocr.py b/src/arcaea_offline_ocr/device/ocr.py index 7232bcb..17927b0 100644 --- a/src/arcaea_offline_ocr/device/ocr.py +++ b/src/arcaea_offline_ocr/device/ocr.py @@ -1,6 +1,5 @@ import cv2 import numpy as np -from PIL import Image from ..crop import crop_xywh from ..ocr import ( @@ -98,5 +97,64 @@ class DeviceOcr: ] return max(enumerate(results), key=lambda i: np.count_nonzero(i[1]))[0] + def lookup_song_id(self): + return self.phash_db.lookup_jacket( + cv2.cvtColor(self.extractor.jacket, cv2.COLOR_BGR2GRAY) + ) + def song_id(self): - return self.phash_db.lookup_image(Image.fromarray(self.extractor.jacket))[0] + return self.lookup_song_id()[0] + + @staticmethod + def preprocess_char_icon(img_gray: cv2.Mat): + h, w = img_gray.shape[:2] + img = cv2.copyMakeBorder(img_gray, w - h, 0, 0, 0, cv2.BORDER_REPLICATE) + h, w = img.shape[:2] + img = cv2.fillPoly( + img, + [ + np.array([[0, 0], [round(w / 2), 0], [0, round(h / 2)]], np.int32), + np.array([[w, 0], [round(w / 2), 0], [w, round(h / 2)]], np.int32), + np.array([[0, h], [round(w / 2), h], [0, round(h / 2)]], np.int32), + np.array([[w, h], [round(w / 2), h], [w, round(h / 2)]], np.int32), + ], + (128), + ) + return img + + def lookup_partner_id(self): + return self.phash_db.lookup_partner_icon( + self.preprocess_char_icon( + cv2.cvtColor(self.extractor.partner_icon, cv2.COLOR_BGR2GRAY) + ) + ) + + def partner_id(self): + return self.lookup_partner_id()[0] + + def ocr(self) -> DeviceOcrResult: + rating_class = self.rating_class() + pure = self.pure() + far = self.far() + lost = self.lost() + score = self.score() + max_recall = self.max_recall() + clear_status = self.clear_status() + + hash_len = self.phash_db.hash_size**2 + song_id, song_id_distance = self.lookup_song_id() + partner_id, partner_id_distance = self.lookup_partner_id() + + return DeviceOcrResult( + rating_class=rating_class, + pure=pure, + far=far, + lost=lost, + score=score, + max_recall=max_recall, + song_id=song_id, + song_id_possibility=1 - song_id_distance / hash_len, + clear_status=clear_status, + partner_id=partner_id, + partner_id_possibility=1 - partner_id_distance / hash_len, + ) diff --git a/src/arcaea_offline_ocr/phash_db.py b/src/arcaea_offline_ocr/phash_db.py index a0a406d..8d95b6b 100644 --- a/src/arcaea_offline_ocr/phash_db.py +++ b/src/arcaea_offline_ocr/phash_db.py @@ -64,14 +64,14 @@ class ImagePhashDatabase: self.jacket_ids: List[str] = [] self.jacket_hashes = [] - self.partner_ids: List[str] = [] - self.partner_hashes = [] + self.partner_icon_ids: List[str] = [] + self.partner_icon_hashes = [] for id, hash in zip(self.ids, self.hashes): id_splitted = id.split("||") - if len(id_splitted) > 1 and id_splitted[0] == "partner": - self.partner_ids.append(id) - self.partner_hashes.append(hash) + if len(id_splitted) > 1 and id_splitted[0] == "partner_icon": + self.partner_icon_ids.append(id_splitted[1]) + self.partner_icon_hashes.append(hash) else: self.jacket_ids.append(id) self.jacket_hashes.append(hash) @@ -104,13 +104,13 @@ class ImagePhashDatabase: def lookup_jacket(self, img_gray: cv2.Mat): return self.lookup_jackets(img_gray)[0] - def lookup_partners(self, img_gray: cv2.Mat, *, limit: int = 5): + def lookup_partner_icons(self, img_gray: cv2.Mat, *, limit: int = 5): image_hash = self.calculate_phash(img_gray).flatten() xor_results = [ (id, np.count_nonzero(image_hash ^ h)) - for id, h in zip(self.partner_ids, self.partner_hashes) + for id, h in zip(self.partner_icon_ids, self.partner_icon_hashes) ] return sorted(xor_results, key=lambda r: r[1])[:limit] - def lookup_partner(self, img_gray: cv2.Mat): - return self.lookup_partners(img_gray)[0] + def lookup_partner_icon(self, img_gray: cv2.Mat): + return self.lookup_partner_icons(img_gray)[0] From 5e0642c8328d12456d441d6d53f5af3a1271f7b7 Mon Sep 17 00:00:00 2001 From: 283375 Date: Tue, 10 Oct 2023 01:45:21 +0800 Subject: [PATCH 16/26] impr: max recall masking --- src/arcaea_offline_ocr/device/rois/masker/auto.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/arcaea_offline_ocr/device/rois/masker/auto.py b/src/arcaea_offline_ocr/device/rois/masker/auto.py index 77f0696..164fe2b 100644 --- a/src/arcaea_offline_ocr/device/rois/masker/auto.py +++ b/src/arcaea_offline_ocr/device/rois/masker/auto.py @@ -149,7 +149,7 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto): BYD_HSV_MAX = np.array([179, 210, 198], np.uint8) MAX_RECALL_HSV_MIN = np.array([125, 0, 0], np.uint8) - MAX_RECALL_HSV_MAX = np.array([130, 100, 150], np.uint8) + MAX_RECALL_HSV_MAX = np.array([145, 100, 150], np.uint8) TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8) TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8) From 02599780e38558d83f759e04ea08215847b7c2e4 Mon Sep 17 00:00:00 2001 From: 283375 Date: Thu, 12 Oct 2023 01:36:57 +0800 Subject: [PATCH 17/26] impr: minor improvements --- src/arcaea_offline_ocr/device/common.py | 2 +- src/arcaea_offline_ocr/device/ocr.py | 3 +- src/arcaea_offline_ocr/ocr.py | 70 ++++++++++++------------- 3 files changed, 36 insertions(+), 39 deletions(-) diff --git a/src/arcaea_offline_ocr/device/common.py b/src/arcaea_offline_ocr/device/common.py index 69c49cb..9cdc076 100644 --- a/src/arcaea_offline_ocr/device/common.py +++ b/src/arcaea_offline_ocr/device/common.py @@ -13,6 +13,6 @@ class DeviceOcrResult: max_recall: Optional[int] = None song_id: Optional[str] = None song_id_possibility: Optional[float] = None - clear_status: Optional[str] = None + clear_status: Optional[int] = None partner_id: Optional[str] = None partner_id_possibility: Optional[float] = None diff --git a/src/arcaea_offline_ocr/device/ocr.py b/src/arcaea_offline_ocr/device/ocr.py index 17927b0..1b87a45 100644 --- a/src/arcaea_offline_ocr/device/ocr.py +++ b/src/arcaea_offline_ocr/device/ocr.py @@ -47,8 +47,7 @@ class DeviceOcr: continue roi_ocr = cv2.fillPoly(roi_ocr, [contour], [0]) digit_rois = [ - resize_fill_square(crop_xywh(roi_ocr, r), 20) - for r in sorted(filtered_rects, key=lambda r: r[0]) + resize_fill_square(crop_xywh(roi_ocr, r), 20) for r in filtered_rects ] samples = preprocess_hog(digit_rois) diff --git a/src/arcaea_offline_ocr/ocr.py b/src/arcaea_offline_ocr/ocr.py index 95eb0dc..b15f0d1 100644 --- a/src/arcaea_offline_ocr/ocr.py +++ b/src/arcaea_offline_ocr/ocr.py @@ -1,5 +1,4 @@ import math -from copy import deepcopy from typing import Optional, Sequence, Tuple import cv2 @@ -64,8 +63,7 @@ class FixRects: new_h = new_bottom - new_y new_rects.append((new_x, new_y, new_w, new_h)) - return_rects = deepcopy(rects) - return_rects = [r for r in return_rects if r not in consumed_rects] + return_rects = [r for r in rects if r not in consumed_rects] return_rects.extend(new_rects) return return_rects @@ -80,42 +78,42 @@ class FixRects: new_rects = [] for rect in rects: rx, ry, rw, rh = rect - if rw / rh > rect_wh_ratio: - # consider this is a connected contour - connected_rects.append(rect) + if rw / rh <= rect_wh_ratio: + continue - # find the thinnest part - border_ignore = round(rw * width_range_ratio) - img_cropped = crop_xywh( - img_masked, - (border_ignore, ry, rw - border_ignore, rh), - ) - white_pixels = {} # dict[x, white_pixel_number] - for i in range(img_cropped.shape[1]): - col = img_cropped[:, i] - white_pixels[rx + border_ignore + i] = np.count_nonzero(col > 200) - least_white_pixels = min(v for v in white_pixels.values() if v > 0) - x_values = [ - x - for x, pixel in white_pixels.items() - if pixel == least_white_pixels - ] - # select only middle values - x_mean = np.mean(x_values) - x_std = np.std(x_values) - x_values = [ - x - for x in x_values - if x_mean - x_std * 1.5 <= x <= x_mean + x_std * 1.5 - ] - x_mid = round(np.median(x_values)) + connected_rects.append(rect) - # split the rect - new_rects.extend( - [(rx, ry, x_mid - rx, rh), (x_mid, ry, rx + rw - x_mid, rh)] - ) + # find the thinnest part + border_ignore = round(rw * width_range_ratio) + img_cropped = crop_xywh( + img_masked, + (border_ignore, ry, rw - border_ignore, rh), + ) + white_pixels = {} # dict[x, white_pixel_number] + for i in range(img_cropped.shape[1]): + col = img_cropped[:, i] + white_pixels[rx + border_ignore + i] = np.count_nonzero(col > 200) + + if all(v == 0 for v in white_pixels.values()): + return rects + + least_white_pixels = min(v for v in white_pixels.values() if v > 0) + x_values = [ + x for x, pixel in white_pixels.items() if pixel == least_white_pixels + ] + # select only middle values + x_mean = np.mean(x_values) + x_std = np.std(x_values) + x_values = [ + x for x in x_values if x_mean - x_std * 1.5 <= x <= x_mean + x_std * 1.5 + ] + x_mid = round(np.median(x_values)) + + # split the rect + new_rects.extend( + [(rx, ry, x_mid - rx, rh), (x_mid, ry, rx + rw - x_mid, rh)] + ) - return_rects = deepcopy(rects) return_rects = [r for r in rects if r not in connected_rects] return_rects.extend(new_rects) return return_rects From 2895eb72331a46226774dd5863c202139dea5bcf Mon Sep 17 00:00:00 2001 From: 283375 Date: Thu, 12 Oct 2023 01:37:59 +0800 Subject: [PATCH 18/26] refactor!: hog computing --- src/arcaea_offline_ocr/ocr.py | 24 +++--------------------- 1 file changed, 3 insertions(+), 21 deletions(-) diff --git a/src/arcaea_offline_ocr/ocr.py b/src/arcaea_offline_ocr/ocr.py index b15f0d1..83a30b4 100644 --- a/src/arcaea_offline_ocr/ocr.py +++ b/src/arcaea_offline_ocr/ocr.py @@ -3,7 +3,6 @@ from typing import Optional, Sequence, Tuple import cv2 import numpy as np -from numpy.linalg import norm from .crop import crop_xywh from .types import Mat, cv2_ml_KNearest @@ -142,28 +141,11 @@ def resize_fill_square(img: Mat, target: int = 20): def preprocess_hog(digit_rois): - # https://github.com/opencv/opencv/blob/f834736307c8328340aea48908484052170c9224/samples/python/digits.py + # https://learnopencv.com/handwritten-digits-classification-an-opencv-c-python-tutorial/ samples = [] for digit in digit_rois: - gx = cv2.Sobel(digit, cv2.CV_32F, 1, 0) - gy = cv2.Sobel(digit, cv2.CV_32F, 0, 1) - mag, ang = cv2.cartToPolar(gx, gy) - bin_n = 16 - _bin = np.int32(bin_n * ang / (2 * np.pi)) - bin_cells = _bin[:10, :10], _bin[10:, :10], _bin[:10, 10:], _bin[10:, 10:] - mag_cells = mag[:10, :10], mag[10:, :10], mag[:10, 10:], mag[10:, 10:] - hists = [ - np.bincount(b.ravel(), m.ravel(), bin_n) - for b, m in zip(bin_cells, mag_cells) - ] - hist = np.hstack(hists) - - # transform to Hellinger kernel - eps = 1e-7 - hist /= hist.sum() + eps - hist = np.sqrt(hist) - hist /= norm(hist) + eps - + hog = cv2.HOGDescriptor((20, 20), (10, 10), (5, 5), (10, 10), 9) + hist = hog.compute(digit) samples.append(hist) return np.float32(samples) From 82229b8b5c99036c3aff6a71e04883713598f051 Mon Sep 17 00:00:00 2001 From: 283375 Date: Thu, 12 Oct 2023 01:50:27 +0800 Subject: [PATCH 19/26] chore: remove custom cv2 type annotations (#8) --- src/arcaea_offline_ocr/b30/chieri/v4/ocr.py | 57 +++++--------------- src/arcaea_offline_ocr/b30/chieri/v4/rois.py | 6 ++- src/arcaea_offline_ocr/crop.py | 11 ++-- src/arcaea_offline_ocr/ocr.py | 13 +++-- src/arcaea_offline_ocr/types.py | 23 +------- src/arcaea_offline_ocr/utils.py | 6 +-- 6 files changed, 32 insertions(+), 84 deletions(-) diff --git a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py index 302b9a5..a68a9c1 100644 --- a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py +++ b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py @@ -8,7 +8,6 @@ from PIL import Image from ....crop import crop_xywh from ....ocr import FixRects, ocr_digits_by_contour_knn, preprocess_hog from ....phash_db import ImagePhashDatabase -from ....types import Mat, cv2_ml_KNearest from ....utils import construct_int_xywh_rect from ...shared import B30OcrResultItem from .colors import * @@ -18,8 +17,8 @@ from .rois import ChieriBotV4Rois class ChieriBotV4Ocr: def __init__( self, - score_knn: cv2_ml_KNearest, - pfl_knn: cv2_ml_KNearest, + score_knn: cv2.ml.KNearest, + pfl_knn: cv2.ml.KNearest, phash_db: ImagePhashDatabase, factor: Optional[float] = 1.0, ): @@ -33,7 +32,7 @@ class ChieriBotV4Ocr: return self.__score_knn @score_knn.setter - def score_knn(self, knn_digits_model: Mat): + def score_knn(self, knn_digits_model: cv2.ml.KNearest): self.__score_knn = knn_digits_model @property @@ -41,7 +40,7 @@ class ChieriBotV4Ocr: return self.__pfl_knn @pfl_knn.setter - def pfl_knn(self, knn_digits_model: Mat): + def pfl_knn(self, knn_digits_model: cv2.ml.KNearest): self.__pfl_knn = knn_digits_model @property @@ -64,10 +63,10 @@ class ChieriBotV4Ocr: def factor(self, factor: float): self.__rois.factor = factor - def set_factor(self, img: Mat): + def set_factor(self, img: cv2.Mat): self.factor = img.shape[0] / 4400 - def ocr_component_rating_class(self, component_bgr: Mat) -> int: + def ocr_component_rating_class(self, component_bgr: cv2.Mat) -> int: rating_class_rect = construct_int_xywh_rect( self.rois.component_rois.rating_class_rect ) @@ -84,15 +83,7 @@ class ChieriBotV4Ocr: else: return max(enumerate(rating_class_results), key=lambda i: i[1])[0] + 1 - # def ocr_component_title(self, component_bgr: Mat) -> str: - # # sourcery skip: inline-immediately-returned-variable - # title_rect = construct_int_xywh_rect(self.rois.component_rois.title_rect) - # title_roi = crop_xywh(component_bgr, title_rect) - # ocr_result = self.sift_db.ocr(title_roi, cls=False) - # title = ocr_result[0][-1][1][0] if ocr_result and ocr_result[0] else "" - # return title - - def ocr_component_song_id(self, component_bgr: Mat): + def ocr_component_song_id(self, component_bgr: cv2.Mat): jacket_rect = construct_int_xywh_rect( self.rois.component_rois.jacket_rect, floor ) @@ -101,20 +92,7 @@ class ChieriBotV4Ocr: ) return self.phash_db.lookup_image(Image.fromarray(jacket_roi))[0] - # def ocr_component_score_paddle(self, component_bgr: Mat) -> int: - # # sourcery skip: inline-immediately-returned-variable - # score_rect = construct_int_xywh_rect(self.rois.component_rois.score_rect) - # score_roi = cv2.cvtColor( - # crop_xywh(component_bgr, score_rect), cv2.COLOR_BGR2GRAY - # ) - # _, score_roi = cv2.threshold( - # score_roi, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU - # ) - # score_str = self.sift_db.ocr(score_roi, cls=False)[0][-1][1][0] - # score = int(score_str.replace("'", "").replace(" ", "")) - # return score - - def ocr_component_score_knn(self, component_bgr: Mat) -> int: + def ocr_component_score_knn(self, component_bgr: cv2.Mat) -> int: # sourcery skip: inline-immediately-returned-variable score_rect = construct_int_xywh_rect(self.rois.component_rois.score_rect) score_roi = cv2.cvtColor( @@ -136,7 +114,7 @@ class ChieriBotV4Ocr: score_roi = cv2.fillPoly(score_roi, [contour], 0) return ocr_digits_by_contour_knn(score_roi, self.score_knn) - def find_pfl_rects(self, component_pfl_processed: Mat) -> List[List[int]]: + def find_pfl_rects(self, component_pfl_processed: cv2.Mat) -> List[List[int]]: # sourcery skip: inline-immediately-returned-variable pfl_roi_find = cv2.morphologyEx( component_pfl_processed, @@ -162,7 +140,7 @@ class ChieriBotV4Ocr: ] return pfl_rects_adjusted - def preprocess_component_pfl(self, component_bgr: Mat) -> Mat: + def preprocess_component_pfl(self, component_bgr: cv2.Mat) -> cv2.Mat: pfl_rect = construct_int_xywh_rect(self.rois.component_rois.pfl_rect) pfl_roi = crop_xywh(component_bgr, pfl_rect) pfl_roi_hsv = cv2.cvtColor(pfl_roi, cv2.COLOR_BGR2HSV) @@ -202,7 +180,7 @@ class ChieriBotV4Ocr: return result_eroded if len(self.find_pfl_rects(result_eroded)) == 3 else result def ocr_component_pfl( - self, component_bgr: Mat + self, component_bgr: cv2.Mat ) -> Tuple[Optional[int], Optional[int], Optional[int]]: try: pfl_roi = self.preprocess_component_pfl(component_bgr) @@ -233,16 +211,7 @@ class ChieriBotV4Ocr: except Exception: return (None, None, None) - # def ocr_component_date(self, component_bgr: Mat): - # date_rect = construct_int_xywh_rect(self.rois.component_rois.date_rect) - # date_roi = cv2.cvtColor(crop_xywh(component_bgr, date_rect), cv2.COLOR_BGR2GRAY) - # _, date_roi = cv2.threshold( - # date_roi, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU - # ) - # date_str = self.sift_db.ocr(date_roi, cls=False)[0][-1][1][0] - # return date_str - - def ocr_component(self, component_bgr: Mat) -> B30OcrResultItem: + def ocr_component(self, component_bgr: cv2.Mat) -> B30OcrResultItem: component_blur = cv2.GaussianBlur(component_bgr, (5, 5), 0) rating_class = self.ocr_component_rating_class(component_blur) song_id = self.ocr_component_song_id(component_bgr) @@ -261,7 +230,7 @@ class ChieriBotV4Ocr: date=None, ) - def ocr(self, img_bgr: Mat) -> List[B30OcrResultItem]: + def ocr(self, img_bgr: cv2.Mat) -> List[B30OcrResultItem]: self.set_factor(img_bgr) return [ self.ocr_component(component_bgr) diff --git a/src/arcaea_offline_ocr/b30/chieri/v4/rois.py b/src/arcaea_offline_ocr/b30/chieri/v4/rois.py index 9926b8a..6bc5c22 100644 --- a/src/arcaea_offline_ocr/b30/chieri/v4/rois.py +++ b/src/arcaea_offline_ocr/b30/chieri/v4/rois.py @@ -1,7 +1,9 @@ from typing import List, Optional +import cv2 + from ....crop import crop_xywh -from ....types import Mat, XYWHRect +from ....types import XYWHRect from ....utils import apply_factor, construct_int_xywh_rect @@ -108,7 +110,7 @@ class ChieriBotV4Rois: def b33_vertical_gap(self): return apply_factor(121, self.factor) - def components(self, img_bgr: Mat) -> List[Mat]: + def components(self, img_bgr: cv2.Mat) -> List[cv2.Mat]: first_rect = XYWHRect(x=self.left, y=self.top, w=self.width, h=self.height) results = [] diff --git a/src/arcaea_offline_ocr/crop.py b/src/arcaea_offline_ocr/crop.py index a65b6ea..5499acf 100644 --- a/src/arcaea_offline_ocr/crop.py +++ b/src/arcaea_offline_ocr/crop.py @@ -1,26 +1,25 @@ from math import floor from typing import Tuple +import cv2 import numpy as np -from .types import Mat - __all__ = ["crop_xywh", "crop_black_edges", "crop_black_edges_grayscale"] -def crop_xywh(mat: Mat, rect: Tuple[int, int, int, int]): +def crop_xywh(mat: cv2.Mat, rect: Tuple[int, int, int, int]): x, y, w, h = rect return mat[y : y + h, x : x + w] -def is_black_edge(list_of_pixels: Mat, black_pixel: Mat, ratio: float = 0.6): +def is_black_edge(list_of_pixels: cv2.Mat, black_pixel: cv2.Mat, ratio: float = 0.6): pixels = list_of_pixels.reshape([-1, 3]) return np.count_nonzero(np.all(pixels < black_pixel, axis=1)) > floor( len(pixels) * ratio ) -def crop_black_edges(img_bgr: Mat, black_threshold: int = 50): +def crop_black_edges(img_bgr: cv2.Mat, black_threshold: int = 50): cropped = img_bgr.copy() black_pixel = np.array([black_threshold] * 3, img_bgr.dtype) height, width = img_bgr.shape[:2] @@ -66,7 +65,7 @@ def is_black_edge_grayscale( def crop_black_edges_grayscale( - img_gray: Mat, black_threshold: int = 50 + img_gray: cv2.Mat, black_threshold: int = 50 ) -> Tuple[int, int, int, int]: """Returns cropped rect""" height, width = img_gray.shape[:2] diff --git a/src/arcaea_offline_ocr/ocr.py b/src/arcaea_offline_ocr/ocr.py index 83a30b4..cfa96e3 100644 --- a/src/arcaea_offline_ocr/ocr.py +++ b/src/arcaea_offline_ocr/ocr.py @@ -5,7 +5,6 @@ import cv2 import numpy as np from .crop import crop_xywh -from .types import Mat, cv2_ml_KNearest __all__ = [ "FixRects", @@ -68,7 +67,7 @@ class FixRects: @staticmethod def split_connected( - img_masked: Mat, + img_masked: cv2.Mat, rects: Sequence[Tuple[int, int, int, int]], rect_wh_ratio: float = 1.05, width_range_ratio: float = 0.1, @@ -118,7 +117,7 @@ class FixRects: return return_rects -def resize_fill_square(img: Mat, target: int = 20): +def resize_fill_square(img: cv2.Mat, target: int = 20): h, w = img.shape[:2] if h > w: new_h = target @@ -150,14 +149,14 @@ def preprocess_hog(digit_rois): return np.float32(samples) -def ocr_digit_samples_knn(__samples, knn_model: cv2_ml_KNearest, k: int = 4): +def ocr_digit_samples_knn(__samples, knn_model: cv2.ml.KNearest, k: int = 4): _, results, _, _ = knn_model.findNearest(__samples, k) result_list = [int(r) for r in results.ravel()] result_str = "".join(str(r) for r in result_list if r > -1) return int(result_str) if result_str else 0 -def ocr_digits_by_contour_get_samples(__roi_gray: Mat, size: int): +def ocr_digits_by_contour_get_samples(__roi_gray: cv2.Mat, size: int): roi = __roi_gray.copy() contours, _ = cv2.findContours(roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) rects = [cv2.boundingRect(c) for c in contours] @@ -170,8 +169,8 @@ def ocr_digits_by_contour_get_samples(__roi_gray: Mat, size: int): def ocr_digits_by_contour_knn( - __roi_gray: Mat, - knn_model: cv2_ml_KNearest, + __roi_gray: cv2.Mat, + knn_model: cv2.ml.KNearest, *, k=4, size: int = 20, diff --git a/src/arcaea_offline_ocr/types.py b/src/arcaea_offline_ocr/types.py index 3a3dc92..dc0bd76 100644 --- a/src/arcaea_offline_ocr/types.py +++ b/src/arcaea_offline_ocr/types.py @@ -1,10 +1,5 @@ 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]] +from typing import NamedTuple, Tuple, Union class XYWHRect(NamedTuple): @@ -24,19 +19,3 @@ class XYWHRect(NamedTuple): 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""" - ... diff --git a/src/arcaea_offline_ocr/utils.py b/src/arcaea_offline_ocr/utils.py index e55ea0f..d7466d2 100644 --- a/src/arcaea_offline_ocr/utils.py +++ b/src/arcaea_offline_ocr/utils.py @@ -1,17 +1,17 @@ import io from collections.abc import Iterable -from typing import Callable, Tuple, TypeVar, Union, overload +from typing import Callable, TypeVar, Union, overload import cv2 import numpy as np from PIL import Image, ImageCms -from .types import Mat, XYWHRect +from .types import XYWHRect __all__ = ["imread_unicode"] -def imread_unicode(filepath: str, flags: int = cv2.IMREAD_UNCHANGED) -> Mat: +def imread_unicode(filepath: str, flags: int = cv2.IMREAD_UNCHANGED): # https://stackoverflow.com/a/57872297/16484891 # CC BY-SA 4.0 return cv2.imdecode(np.fromfile(filepath, dtype=np.uint8), flags) From 4fd31b1e9b13e0bc1697258fc0e48500f4d7aa8e Mon Sep 17 00:00:00 2001 From: 283375 Date: Thu, 12 Oct 2023 17:53:51 +0800 Subject: [PATCH 20/26] impr: module `__init__` entries --- src/arcaea_offline_ocr/device/__init__.py | 2 ++ src/arcaea_offline_ocr/device/rois/__init__.py | 3 +++ src/arcaea_offline_ocr/device/rois/definition/__init__.py | 1 + src/arcaea_offline_ocr/device/rois/extractor/common.py | 4 ++-- 4 files changed, 8 insertions(+), 2 deletions(-) create mode 100644 src/arcaea_offline_ocr/device/rois/__init__.py diff --git a/src/arcaea_offline_ocr/device/__init__.py b/src/arcaea_offline_ocr/device/__init__.py index e69de29..bd6cd44 100644 --- a/src/arcaea_offline_ocr/device/__init__.py +++ b/src/arcaea_offline_ocr/device/__init__.py @@ -0,0 +1,2 @@ +from .common import DeviceOcrResult +from .ocr import DeviceOcr diff --git a/src/arcaea_offline_ocr/device/rois/__init__.py b/src/arcaea_offline_ocr/device/rois/__init__.py new file mode 100644 index 0000000..e73f32a --- /dev/null +++ b/src/arcaea_offline_ocr/device/rois/__init__.py @@ -0,0 +1,3 @@ +from .definition import * +from .extractor import * +from .masker import * diff --git a/src/arcaea_offline_ocr/device/rois/definition/__init__.py b/src/arcaea_offline_ocr/device/rois/definition/__init__.py index 4083875..37f6ef1 100644 --- a/src/arcaea_offline_ocr/device/rois/definition/__init__.py +++ b/src/arcaea_offline_ocr/device/rois/definition/__init__.py @@ -1 +1,2 @@ from .auto import * +from .common import DeviceRois diff --git a/src/arcaea_offline_ocr/device/rois/extractor/common.py b/src/arcaea_offline_ocr/device/rois/extractor/common.py index e672e8b..a90a0a4 100644 --- a/src/arcaea_offline_ocr/device/rois/extractor/common.py +++ b/src/arcaea_offline_ocr/device/rois/extractor/common.py @@ -5,9 +5,9 @@ from ..definition.common import DeviceRois class DeviceRoisExtractor: - def __init__(self, img: cv2.Mat, sizes: DeviceRois): + def __init__(self, img: cv2.Mat, rois: DeviceRois): self.img = img - self.sizes = sizes + self.sizes = rois def __construct_int_rect(self, rect): return tuple(round(r) for r in rect) From 3400df2d522e9823977452520e53fa9cdfa5a487 Mon Sep 17 00:00:00 2001 From: 283375 Date: Thu, 12 Oct 2023 18:08:54 +0800 Subject: [PATCH 21/26] refactor!: remove `utils.convert_to_srgb` --- src/arcaea_offline_ocr/utils.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/src/arcaea_offline_ocr/utils.py b/src/arcaea_offline_ocr/utils.py index d7466d2..9fa9390 100644 --- a/src/arcaea_offline_ocr/utils.py +++ b/src/arcaea_offline_ocr/utils.py @@ -1,10 +1,8 @@ -import io from collections.abc import Iterable from typing import Callable, TypeVar, Union, overload import cv2 import numpy as np -from PIL import Image, ImageCms from .types import XYWHRect @@ -46,25 +44,3 @@ def apply_factor(item, factor: float): return item * factor elif isinstance(item, Iterable): return item.__class__([i * factor for i in item]) - - -def convert_to_srgb(pil_img: Image.Image): - """ - Convert PIL image to sRGB color space (if possible) - and save the converted file. - - https://stackoverflow.com/a/65667797/16484891 - - CC BY-SA 4.0 - """ - icc = pil_img.info.get("icc_profile", "") - icc_conv = "" - - if icc: - io_handle = io.BytesIO(icc) # virtual file - src_profile = ImageCms.ImageCmsProfile(io_handle) - dst_profile = ImageCms.createProfile("sRGB") - img_conv = ImageCms.profileToProfile(pil_img, src_profile, dst_profile) - icc_conv = img_conv.info.get("icc_profile", "") - - return img_conv if icc != icc_conv else pil_img From 42bcd7b430a5a12a16fe56e1b7ca8ba491e9bdf7 Mon Sep 17 00:00:00 2001 From: 283375 Date: Fri, 13 Oct 2023 18:32:33 +0800 Subject: [PATCH 22/26] refactor: B30 ocr --- src/arcaea_offline_ocr/b30/chieri/v4/ocr.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py index a68a9c1..98e1292 100644 --- a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py +++ b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py @@ -3,10 +3,14 @@ from typing import List, Optional, Tuple import cv2 import numpy as np -from PIL import Image from ....crop import crop_xywh -from ....ocr import FixRects, ocr_digits_by_contour_knn, preprocess_hog +from ....ocr import ( + FixRects, + ocr_digits_by_contour_knn, + preprocess_hog, + resize_fill_square, +) from ....phash_db import ImagePhashDatabase from ....utils import construct_int_xywh_rect from ...shared import B30OcrResultItem @@ -90,7 +94,7 @@ class ChieriBotV4Ocr: jacket_roi = cv2.cvtColor( crop_xywh(component_bgr, jacket_rect), cv2.COLOR_BGR2GRAY ) - return self.phash_db.lookup_image(Image.fromarray(jacket_roi))[0] + return self.phash_db.lookup_jacket(jacket_roi)[0] def ocr_component_score_knn(self, component_bgr: cv2.Mat) -> int: # sourcery skip: inline-immediately-returned-variable @@ -200,7 +204,7 @@ class ChieriBotV4Ocr: digits = [] for digit_rect in digit_rects: digit = crop_xywh(roi, digit_rect) - digit = cv2.resize(digit, (20, 20)) + digit = resize_fill_square(digit, 20) digits.append(digit) samples = preprocess_hog(digits) From 122a546174156a4488981ef8bda90531bb9e5990 Mon Sep 17 00:00:00 2001 From: 283375 Date: Sun, 22 Oct 2023 01:55:20 +0800 Subject: [PATCH 23/26] refactor: `CropBlackEdges` --- src/arcaea_offline_ocr/crop.py | 129 ++++++++++++--------------------- 1 file changed, 45 insertions(+), 84 deletions(-) diff --git a/src/arcaea_offline_ocr/crop.py b/src/arcaea_offline_ocr/crop.py index 5499acf..aa9d1f2 100644 --- a/src/arcaea_offline_ocr/crop.py +++ b/src/arcaea_offline_ocr/crop.py @@ -1,10 +1,10 @@ -from math import floor +import math from typing import Tuple import cv2 import numpy as np -__all__ = ["crop_xywh", "crop_black_edges", "crop_black_edges_grayscale"] +__all__ = ["crop_xywh", "CropBlackEdges"] def crop_xywh(mat: cv2.Mat, rect: Tuple[int, int, int, int]): @@ -12,92 +12,53 @@ def crop_xywh(mat: cv2.Mat, rect: Tuple[int, int, int, int]): return mat[y : y + h, x : x + w] -def is_black_edge(list_of_pixels: cv2.Mat, black_pixel: cv2.Mat, ratio: float = 0.6): - pixels = list_of_pixels.reshape([-1, 3]) - return np.count_nonzero(np.all(pixels < black_pixel, axis=1)) > floor( - len(pixels) * ratio - ) +class CropBlackEdges: + @staticmethod + def is_black_edge(__img_gray_slice: cv2.Mat, black_pixel: int, ratio: float = 0.6): + pixels_compared = __img_gray_slice < black_pixel + return np.count_nonzero(pixels_compared) > math.floor( + __img_gray_slice.size * ratio + ) + @classmethod + def get_crop_rect(cls, img_gray: cv2.Mat, black_threshold: int = 25): + height, width = img_gray.shape[:2] + left = 0 + right = width + top = 0 + bottom = height -def crop_black_edges(img_bgr: cv2.Mat, black_threshold: int = 50): - cropped = img_bgr.copy() - black_pixel = np.array([black_threshold] * 3, img_bgr.dtype) - height, width = img_bgr.shape[:2] - left = 0 - right = width - top = 0 - bottom = height + for i in range(width): + column = img_gray[:, i] + if not cls.is_black_edge(column, black_threshold): + break + left += 1 - for i in range(width): - column = cropped[:, i] - if not is_black_edge(column, black_pixel): - break - left += 1 + for i in sorted(range(width), reverse=True): + column = img_gray[:, i] + if i <= left + 1 or not cls.is_black_edge(column, black_threshold): + break + right -= 1 - 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 = img_gray[i] + if not cls.is_black_edge(row, black_threshold): + break + top += 1 - for i in range(height): - row = cropped[i] - if not is_black_edge(row, black_pixel): - break - top += 1 + for i in sorted(range(height), reverse=True): + row = img_gray[i] + if i <= top + 1 or not cls.is_black_edge(row, black_threshold): + break + bottom -= 1 - 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 + assert right > left, "cropped width < 0" + assert bottom > top, "cropped height < 0" + return (left, top, right - left, bottom - top) - return cropped[top:bottom, left:right] - - -def is_black_edge_grayscale( - gray_value_list: np.ndarray, black_threshold: int = 50, ratio: float = 0.6 -) -> bool: - return ( - np.count_nonzero(gray_value_list < black_threshold) - > len(gray_value_list) * ratio - ) - - -def crop_black_edges_grayscale( - img_gray: cv2.Mat, black_threshold: int = 50 -) -> Tuple[int, int, int, int]: - """Returns cropped rect""" - height, width = img_gray.shape[:2] - left = 0 - right = width - top = 0 - bottom = height - - for i in range(width): - column = img_gray[:, i] - if not is_black_edge_grayscale(column, black_threshold): - break - left += 1 - - for i in sorted(range(width), reverse=True): - column = img_gray[:, i] - if i <= left + 1 or not is_black_edge_grayscale(column, black_threshold): - break - right -= 1 - - for i in range(height): - row = img_gray[i] - if not is_black_edge_grayscale(row, black_threshold): - break - top += 1 - - for i in sorted(range(height), reverse=True): - row = img_gray[i] - if i <= top + 1 or not is_black_edge_grayscale(row, black_threshold): - break - bottom -= 1 - - assert right > left, "cropped width > 0" - assert bottom > top, "cropped height > 0" - return (left, top, right - left, bottom - top) + @classmethod + def crop( + cls, img: cv2.Mat, convert_flag: cv2.COLOR_BGR2GRAY, black_threshold: int = 25 + ) -> cv2.Mat: + rect = cls.get_crop_rect(cv2.cvtColor(img, convert_flag), black_threshold) + return crop_xywh(img, rect) From cf46bf7a5978147203aec6cd32ca9152a99daccb Mon Sep 17 00:00:00 2001 From: 283375 Date: Mon, 23 Oct 2023 13:28:59 +0800 Subject: [PATCH 24/26] chore: update dependencies --- pyproject.toml | 2 +- requirements.txt | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 45e97f0..aaea2bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ authors = [{ name = "283375", email = "log_283375@163.com" }] description = "Extract your Arcaea play result from screenshot." readme = "README.md" requires-python = ">=3.8" -dependencies = ["attrs==23.1.0", "numpy==1.25.2", "opencv-python==4.8.0.76"] +dependencies = ["attrs==23.1.0", "numpy==1.26.1", "opencv-python==4.8.1.78"] classifiers = [ "Development Status :: 3 - Alpha", "Programming Language :: Python :: 3", diff --git a/requirements.txt b/requirements.txt index b065292..b1b83a0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,3 @@ attrs==23.1.0 -numpy==1.25.2 -opencv-python==4.8.0.76 +numpy==1.26.1 +opencv-python==4.8.1.78 From f2f854040f6b99841f709f868267e0908808ceb5 Mon Sep 17 00:00:00 2001 From: 283375 Date: Mon, 23 Oct 2023 13:29:17 +0800 Subject: [PATCH 25/26] fix: `cv2.Mat` typing annotations --- src/arcaea_offline_ocr/b30/chieri/v4/ocr.py | 19 ++++--- src/arcaea_offline_ocr/b30/chieri/v4/rois.py | 6 +- src/arcaea_offline_ocr/crop.py | 12 ++-- src/arcaea_offline_ocr/device/ocr.py | 5 +- .../device/rois/extractor/common.py | 5 +- .../device/rois/masker/auto.py | 57 ++++++++++--------- .../device/rois/masker/common.py | 28 ++++----- src/arcaea_offline_ocr/ocr.py | 9 +-- src/arcaea_offline_ocr/phash_db.py | 16 +++--- src/arcaea_offline_ocr/types.py | 4 ++ 10 files changed, 85 insertions(+), 76 deletions(-) diff --git a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py index 98e1292..845145a 100644 --- a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py +++ b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py @@ -12,6 +12,7 @@ from ....ocr import ( resize_fill_square, ) from ....phash_db import ImagePhashDatabase +from ....types import Mat from ....utils import construct_int_xywh_rect from ...shared import B30OcrResultItem from .colors import * @@ -67,10 +68,10 @@ class ChieriBotV4Ocr: def factor(self, factor: float): self.__rois.factor = factor - def set_factor(self, img: cv2.Mat): + def set_factor(self, img: Mat): self.factor = img.shape[0] / 4400 - def ocr_component_rating_class(self, component_bgr: cv2.Mat) -> int: + def ocr_component_rating_class(self, component_bgr: Mat) -> int: rating_class_rect = construct_int_xywh_rect( self.rois.component_rois.rating_class_rect ) @@ -87,7 +88,7 @@ class ChieriBotV4Ocr: else: return max(enumerate(rating_class_results), key=lambda i: i[1])[0] + 1 - def ocr_component_song_id(self, component_bgr: cv2.Mat): + def ocr_component_song_id(self, component_bgr: Mat): jacket_rect = construct_int_xywh_rect( self.rois.component_rois.jacket_rect, floor ) @@ -96,7 +97,7 @@ class ChieriBotV4Ocr: ) return self.phash_db.lookup_jacket(jacket_roi)[0] - def ocr_component_score_knn(self, component_bgr: cv2.Mat) -> int: + def ocr_component_score_knn(self, component_bgr: Mat) -> int: # sourcery skip: inline-immediately-returned-variable score_rect = construct_int_xywh_rect(self.rois.component_rois.score_rect) score_roi = cv2.cvtColor( @@ -118,7 +119,7 @@ class ChieriBotV4Ocr: score_roi = cv2.fillPoly(score_roi, [contour], 0) return ocr_digits_by_contour_knn(score_roi, self.score_knn) - def find_pfl_rects(self, component_pfl_processed: cv2.Mat) -> List[List[int]]: + def find_pfl_rects(self, component_pfl_processed: Mat) -> List[List[int]]: # sourcery skip: inline-immediately-returned-variable pfl_roi_find = cv2.morphologyEx( component_pfl_processed, @@ -144,7 +145,7 @@ class ChieriBotV4Ocr: ] return pfl_rects_adjusted - def preprocess_component_pfl(self, component_bgr: cv2.Mat) -> cv2.Mat: + def preprocess_component_pfl(self, component_bgr: Mat) -> Mat: pfl_rect = construct_int_xywh_rect(self.rois.component_rois.pfl_rect) pfl_roi = crop_xywh(component_bgr, pfl_rect) pfl_roi_hsv = cv2.cvtColor(pfl_roi, cv2.COLOR_BGR2HSV) @@ -184,7 +185,7 @@ class ChieriBotV4Ocr: return result_eroded if len(self.find_pfl_rects(result_eroded)) == 3 else result def ocr_component_pfl( - self, component_bgr: cv2.Mat + self, component_bgr: Mat ) -> Tuple[Optional[int], Optional[int], Optional[int]]: try: pfl_roi = self.preprocess_component_pfl(component_bgr) @@ -215,7 +216,7 @@ class ChieriBotV4Ocr: except Exception: return (None, None, None) - def ocr_component(self, component_bgr: cv2.Mat) -> B30OcrResultItem: + def ocr_component(self, component_bgr: Mat) -> B30OcrResultItem: component_blur = cv2.GaussianBlur(component_bgr, (5, 5), 0) rating_class = self.ocr_component_rating_class(component_blur) song_id = self.ocr_component_song_id(component_bgr) @@ -234,7 +235,7 @@ class ChieriBotV4Ocr: date=None, ) - def ocr(self, img_bgr: cv2.Mat) -> List[B30OcrResultItem]: + def ocr(self, img_bgr: Mat) -> List[B30OcrResultItem]: self.set_factor(img_bgr) return [ self.ocr_component(component_bgr) diff --git a/src/arcaea_offline_ocr/b30/chieri/v4/rois.py b/src/arcaea_offline_ocr/b30/chieri/v4/rois.py index 6bc5c22..9926b8a 100644 --- a/src/arcaea_offline_ocr/b30/chieri/v4/rois.py +++ b/src/arcaea_offline_ocr/b30/chieri/v4/rois.py @@ -1,9 +1,7 @@ from typing import List, Optional -import cv2 - from ....crop import crop_xywh -from ....types import XYWHRect +from ....types import Mat, XYWHRect from ....utils import apply_factor, construct_int_xywh_rect @@ -110,7 +108,7 @@ class ChieriBotV4Rois: def b33_vertical_gap(self): return apply_factor(121, self.factor) - def components(self, img_bgr: cv2.Mat) -> List[cv2.Mat]: + def components(self, img_bgr: Mat) -> List[Mat]: first_rect = XYWHRect(x=self.left, y=self.top, w=self.width, h=self.height) results = [] diff --git a/src/arcaea_offline_ocr/crop.py b/src/arcaea_offline_ocr/crop.py index aa9d1f2..12c531d 100644 --- a/src/arcaea_offline_ocr/crop.py +++ b/src/arcaea_offline_ocr/crop.py @@ -4,24 +4,26 @@ from typing import Tuple import cv2 import numpy as np +from .types import Mat + __all__ = ["crop_xywh", "CropBlackEdges"] -def crop_xywh(mat: cv2.Mat, rect: Tuple[int, int, int, int]): +def crop_xywh(mat: Mat, rect: Tuple[int, int, int, int]): x, y, w, h = rect return mat[y : y + h, x : x + w] class CropBlackEdges: @staticmethod - def is_black_edge(__img_gray_slice: cv2.Mat, black_pixel: int, ratio: float = 0.6): + def is_black_edge(__img_gray_slice: Mat, black_pixel: int, ratio: float = 0.6): pixels_compared = __img_gray_slice < black_pixel return np.count_nonzero(pixels_compared) > math.floor( __img_gray_slice.size * ratio ) @classmethod - def get_crop_rect(cls, img_gray: cv2.Mat, black_threshold: int = 25): + def get_crop_rect(cls, img_gray: Mat, black_threshold: int = 25): height, width = img_gray.shape[:2] left = 0 right = width @@ -58,7 +60,7 @@ class CropBlackEdges: @classmethod def crop( - cls, img: cv2.Mat, convert_flag: cv2.COLOR_BGR2GRAY, black_threshold: int = 25 - ) -> cv2.Mat: + cls, img: Mat, convert_flag: cv2.COLOR_BGR2GRAY, black_threshold: int = 25 + ) -> Mat: rect = cls.get_crop_rect(cv2.cvtColor(img, convert_flag), black_threshold) return crop_xywh(img, rect) diff --git a/src/arcaea_offline_ocr/device/ocr.py b/src/arcaea_offline_ocr/device/ocr.py index 1b87a45..d0cbf88 100644 --- a/src/arcaea_offline_ocr/device/ocr.py +++ b/src/arcaea_offline_ocr/device/ocr.py @@ -10,6 +10,7 @@ from ..ocr import ( resize_fill_square, ) from ..phash_db import ImagePhashDatabase +from ..types import Mat from .common import DeviceOcrResult from .rois.extractor import DeviceRoisExtractor from .rois.masker import DeviceRoisMasker @@ -28,7 +29,7 @@ class DeviceOcr: self.knn_model = knn_model self.phash_db = phash_db - def pfl(self, roi_gray: cv2.Mat, factor: float = 1.25): + def pfl(self, roi_gray: Mat, factor: float = 1.25): contours, _ = cv2.findContours( roi_gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE ) @@ -105,7 +106,7 @@ class DeviceOcr: return self.lookup_song_id()[0] @staticmethod - def preprocess_char_icon(img_gray: cv2.Mat): + def preprocess_char_icon(img_gray: Mat): h, w = img_gray.shape[:2] img = cv2.copyMakeBorder(img_gray, w - h, 0, 0, 0, cv2.BORDER_REPLICATE) h, w = img.shape[:2] diff --git a/src/arcaea_offline_ocr/device/rois/extractor/common.py b/src/arcaea_offline_ocr/device/rois/extractor/common.py index a90a0a4..671ae2c 100644 --- a/src/arcaea_offline_ocr/device/rois/extractor/common.py +++ b/src/arcaea_offline_ocr/device/rois/extractor/common.py @@ -1,11 +1,10 @@ -import cv2 - from ....crop import crop_xywh +from ....types import Mat from ..definition.common import DeviceRois class DeviceRoisExtractor: - def __init__(self, img: cv2.Mat, rois: DeviceRois): + def __init__(self, img: Mat, rois: DeviceRois): self.img = img self.sizes = rois diff --git a/src/arcaea_offline_ocr/device/rois/masker/auto.py b/src/arcaea_offline_ocr/device/rois/masker/auto.py index 164fe2b..ec92548 100644 --- a/src/arcaea_offline_ocr/device/rois/masker/auto.py +++ b/src/arcaea_offline_ocr/device/rois/masker/auto.py @@ -1,6 +1,7 @@ import cv2 import numpy as np +from ....types import Mat from .common import DeviceRoisMasker @@ -40,7 +41,7 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto): PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8) @classmethod - def gray(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def gray(cls, roi_bgr: Mat) -> Mat: bgr_value_equal_mask = np.max(roi_bgr, axis=2) - np.min(roi_bgr, axis=2) <= 5 img_bgr = roi_bgr.copy() img_bgr[~bgr_value_equal_mask] = np.array([0, 0, 0], roi_bgr.dtype) @@ -49,19 +50,19 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto): return cv2.inRange(img_bgr, cls.GRAY_BGR_MIN, cls.GRAY_BGR_MAX) @classmethod - def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def pure(cls, roi_bgr: Mat) -> Mat: return cls.gray(roi_bgr) @classmethod - def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def far(cls, roi_bgr: Mat) -> Mat: return cls.gray(roi_bgr) @classmethod - def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def lost(cls, roi_bgr: Mat) -> Mat: return cls.gray(roi_bgr) @classmethod - def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def score(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.WHITE_HSV_MIN, @@ -69,35 +70,35 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto): ) @classmethod - def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_pst(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PST_HSV_MIN, cls.PST_HSV_MAX ) @classmethod - def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_prs(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PRS_HSV_MIN, cls.PRS_HSV_MAX ) @classmethod - def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_ftr(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.FTR_HSV_MIN, cls.FTR_HSV_MAX ) @classmethod - def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_byd(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.BYD_HSV_MIN, cls.BYD_HSV_MAX ) @classmethod - def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def max_recall(cls, roi_bgr: Mat) -> Mat: return cls.gray(roi_bgr) @classmethod - def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_track_lost(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.TRACK_LOST_HSV_MIN, @@ -105,7 +106,7 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto): ) @classmethod - def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_track_complete(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.TRACK_COMPLETE_HSV_MIN, @@ -113,7 +114,7 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto): ) @classmethod - def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_full_recall(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.FULL_RECALL_HSV_MIN, @@ -121,7 +122,7 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto): ) @classmethod - def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_pure_memory(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PURE_MEMORY_HSV_MIN, @@ -164,25 +165,25 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto): PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8) @classmethod - def pfl(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def pfl(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PFL_HSV_MIN, cls.PFL_HSV_MAX ) @classmethod - def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def pure(cls, roi_bgr: Mat) -> Mat: return cls.pfl(roi_bgr) @classmethod - def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def far(cls, roi_bgr: Mat) -> Mat: return cls.pfl(roi_bgr) @classmethod - def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def lost(cls, roi_bgr: Mat) -> Mat: return cls.pfl(roi_bgr) @classmethod - def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def score(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.WHITE_HSV_MIN, @@ -190,31 +191,31 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto): ) @classmethod - def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_pst(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PST_HSV_MIN, cls.PST_HSV_MAX ) @classmethod - def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_prs(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PRS_HSV_MIN, cls.PRS_HSV_MAX ) @classmethod - def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_ftr(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.FTR_HSV_MIN, cls.FTR_HSV_MAX ) @classmethod - def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_byd(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.BYD_HSV_MIN, cls.BYD_HSV_MAX ) @classmethod - def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def max_recall(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.MAX_RECALL_HSV_MIN, @@ -222,7 +223,7 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto): ) @classmethod - def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_track_lost(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.TRACK_LOST_HSV_MIN, @@ -230,7 +231,7 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto): ) @classmethod - def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_track_complete(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.TRACK_COMPLETE_HSV_MIN, @@ -238,7 +239,7 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto): ) @classmethod - def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_full_recall(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.FULL_RECALL_HSV_MIN, @@ -246,7 +247,7 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto): ) @classmethod - def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_pure_memory(cls, roi_bgr: Mat) -> Mat: return cv2.inRange( cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PURE_MEMORY_HSV_MIN, diff --git a/src/arcaea_offline_ocr/device/rois/masker/common.py b/src/arcaea_offline_ocr/device/rois/masker/common.py index f877e2c..cb1223f 100644 --- a/src/arcaea_offline_ocr/device/rois/masker/common.py +++ b/src/arcaea_offline_ocr/device/rois/masker/common.py @@ -1,55 +1,55 @@ -import cv2 +from ....types import Mat class DeviceRoisMasker: @classmethod - def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def pure(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def far(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def lost(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def score(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_pst(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_prs(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_ftr(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def rating_class_byd(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def max_recall(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_track_lost(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_track_complete(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_full_recall(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() @classmethod - def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat: + def clear_status_pure_memory(cls, roi_bgr: Mat) -> Mat: raise NotImplementedError() diff --git a/src/arcaea_offline_ocr/ocr.py b/src/arcaea_offline_ocr/ocr.py index cfa96e3..44ca73a 100644 --- a/src/arcaea_offline_ocr/ocr.py +++ b/src/arcaea_offline_ocr/ocr.py @@ -5,6 +5,7 @@ import cv2 import numpy as np from .crop import crop_xywh +from .types import Mat __all__ = [ "FixRects", @@ -67,7 +68,7 @@ class FixRects: @staticmethod def split_connected( - img_masked: cv2.Mat, + img_masked: Mat, rects: Sequence[Tuple[int, int, int, int]], rect_wh_ratio: float = 1.05, width_range_ratio: float = 0.1, @@ -117,7 +118,7 @@ class FixRects: return return_rects -def resize_fill_square(img: cv2.Mat, target: int = 20): +def resize_fill_square(img: Mat, target: int = 20): h, w = img.shape[:2] if h > w: new_h = target @@ -156,7 +157,7 @@ def ocr_digit_samples_knn(__samples, knn_model: cv2.ml.KNearest, k: int = 4): return int(result_str) if result_str else 0 -def ocr_digits_by_contour_get_samples(__roi_gray: cv2.Mat, size: int): +def ocr_digits_by_contour_get_samples(__roi_gray: Mat, size: int): roi = __roi_gray.copy() contours, _ = cv2.findContours(roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) rects = [cv2.boundingRect(c) for c in contours] @@ -169,7 +170,7 @@ def ocr_digits_by_contour_get_samples(__roi_gray: cv2.Mat, size: int): def ocr_digits_by_contour_knn( - __roi_gray: cv2.Mat, + __roi_gray: Mat, knn_model: cv2.ml.KNearest, *, k=4, diff --git a/src/arcaea_offline_ocr/phash_db.py b/src/arcaea_offline_ocr/phash_db.py index 8d95b6b..dba7c04 100644 --- a/src/arcaea_offline_ocr/phash_db.py +++ b/src/arcaea_offline_ocr/phash_db.py @@ -4,9 +4,11 @@ from typing import List, Union import cv2 import numpy as np +from .types import Mat + def phash_opencv(img_gray, hash_size=8, highfreq_factor=4): - # type: (Union[cv2.Mat, np.ndarray], int, int) -> np.ndarray + # type: (Union[Mat, np.ndarray], int, int) -> np.ndarray """ Perceptual Hash computation. @@ -76,7 +78,7 @@ class ImagePhashDatabase: self.jacket_ids.append(id) self.jacket_hashes.append(hash) - def calculate_phash(self, img_gray: cv2.Mat): + def calculate_phash(self, img_gray: Mat): return phash_opencv( img_gray, hash_size=self.hash_size, highfreq_factor=self.highfreq_factor ) @@ -89,11 +91,11 @@ class ImagePhashDatabase: ] return sorted(xor_results, key=lambda r: r[1])[:limit] - def lookup_image(self, img_gray: cv2.Mat): + def lookup_image(self, img_gray: Mat): image_hash = self.calculate_phash(img_gray) return self.lookup_hash(image_hash)[0] - def lookup_jackets(self, img_gray: cv2.Mat, *, limit: int = 5): + def lookup_jackets(self, img_gray: Mat, *, limit: int = 5): image_hash = self.calculate_phash(img_gray).flatten() xor_results = [ (id, np.count_nonzero(image_hash ^ h)) @@ -101,10 +103,10 @@ class ImagePhashDatabase: ] return sorted(xor_results, key=lambda r: r[1])[:limit] - def lookup_jacket(self, img_gray: cv2.Mat): + def lookup_jacket(self, img_gray: Mat): return self.lookup_jackets(img_gray)[0] - def lookup_partner_icons(self, img_gray: cv2.Mat, *, limit: int = 5): + def lookup_partner_icons(self, img_gray: Mat, *, limit: int = 5): image_hash = self.calculate_phash(img_gray).flatten() xor_results = [ (id, np.count_nonzero(image_hash ^ h)) @@ -112,5 +114,5 @@ class ImagePhashDatabase: ] return sorted(xor_results, key=lambda r: r[1])[:limit] - def lookup_partner_icon(self, img_gray: cv2.Mat): + def lookup_partner_icon(self, img_gray: Mat): return self.lookup_partner_icons(img_gray)[0] diff --git a/src/arcaea_offline_ocr/types.py b/src/arcaea_offline_ocr/types.py index dc0bd76..7f1bc5b 100644 --- a/src/arcaea_offline_ocr/types.py +++ b/src/arcaea_offline_ocr/types.py @@ -1,6 +1,10 @@ from collections.abc import Iterable from typing import NamedTuple, Tuple, Union +import numpy as np + +Mat = np.ndarray + class XYWHRect(NamedTuple): x: int From f9c867e180a02bacc1112682d8fa14ffbb0cc905 Mon Sep 17 00:00:00 2001 From: 283375 Date: Wed, 25 Oct 2023 18:08:29 +0800 Subject: [PATCH 26/26] chore: v0.0.95 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index aaea2bb..1390b88 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "arcaea-offline-ocr" -version = "0.1.0" +version = "0.0.95" authors = [{ name = "283375", email = "log_283375@163.com" }] description = "Extract your Arcaea play result from screenshot." readme = "README.md"