Merge branch 'structure-refactor'

pyproject.toml

@@ -4,12 +4,12 @@ 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"
 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",

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

src/arcaea_offline_ocr/b30/chieri/v4/ocr.py

@@ -3,11 +3,16 @@ 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 ....types import Mat
 from ....utils import construct_int_xywh_rect
 from ...shared import B30OcrResultItem
 from .colors import *
@@ -63,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
         )
@@ -83,16 +88,16 @@ 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
         )
         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:
+    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(
@@ -114,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,
@@ -140,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)
@@ -180,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)
@@ -200,7 +205,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)
 
@@ -211,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)
@@ -230,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)

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 = []
 

src/arcaea_offline_ocr/crop.py

@@ -1,103 +1,66 @@
-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"]
+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]
 
 
-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: 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: 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: 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)

src/arcaea_offline_ocr/device/__init__.py

@@ -0,0 +1,2 @@
+from .common import DeviceOcrResult
+from .ocr import DeviceOcr

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]

src/arcaea_offline_ocr/device/rois/__init__.py

@@ -0,0 +1,3 @@
+from .definition import *
+from .extractor import *
+from .masker import *

src/arcaea_offline_ocr/device/rois/definition/__init__.py

@@ -1 +1,2 @@
 from .auto import *
+from .common import DeviceRois

src/arcaea_offline_ocr/device/rois/extractor/common.py

@@ -1,13 +1,12 @@
-import cv2
-
 from ....crop import crop_xywh
+from ....types import Mat
 from ..definition.common import DeviceRois
 
 
 class DeviceRoisExtractor:
-    def __init__(self, img: cv2.Mat, sizes: DeviceRois):
+    def __init__(self, img: 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)

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,

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

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,

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]

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

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