wip: SizesV2

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

@@ -3,9 +3,11 @@ 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 ....phash_db import ImagePHashDatabase
 from ....sift_db import SIFTDatabase
 from ....types import Mat, cv2_ml_KNearest
 from ....utils import construct_int_xywh_rect
@@ -19,12 +21,12 @@ class ChieriBotV4Ocr:
         self,
         score_knn: cv2_ml_KNearest,
         pfl_knn: cv2_ml_KNearest,
-        sift_db: SIFTDatabase,
+        phash_db: ImagePHashDatabase,
         factor: Optional[float] = 1.0,
     ):
         self.__score_knn = score_knn
         self.__pfl_knn = pfl_knn
-        self.__sift_db = sift_db
+        self.__phash_db = phash_db
         self.__rois = ChieriBotV4Rois(factor)
 
     @property
@@ -44,12 +46,12 @@ class ChieriBotV4Ocr:
         self.__pfl_knn = knn_digits_model
 
     @property
-    def sift_db(self):
-        return self.__sift_db
+    def phash_db(self):
+        return self.__phash_db
 
-    @sift_db.setter
-    def sift_db(self, sift_db: SIFTDatabase):
-        self.__sift_db = sift_db
+    @phash_db.setter
+    def phash_db(self, phash_db: ImagePHashDatabase):
+        self.__phash_db = phash_db
 
     @property
     def rois(self):
@@ -98,7 +100,7 @@ class ChieriBotV4Ocr:
         jacket_roi = cv2.cvtColor(
             crop_xywh(component_bgr, jacket_rect), cv2.COLOR_BGR2GRAY
         )
-        return self.sift_db.lookup_img(jacket_roi)[0]
+        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

src/arcaea_offline_ocr/device/v2/ocr.py

@@ -4,9 +4,19 @@ 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_prs, mask_pst, mask_white
+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,
@@ -14,18 +24,20 @@ from ...ocr import (
     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, sift_db: SIFTDatabase):
+    def __init__(self, knn_model: cv2_ml_KNearest, phash_db: ImagePHashDatabase):
         self.__knn_model = knn_model
-        self.__sift_db = sift_db
+        self.__phash_db = phash_db
 
     @property
     def knn_model(self):
@@ -38,14 +50,14 @@ class DeviceV2Ocr:
         self.__knn_model = value
 
     @property
-    def sift_db(self):
-        if not self.__sift_db:
-            raise ValueError("`sift_db` unset.")
-        return self.__sift_db
+    def phash_db(self):
+        if not self.__phash_db:
+            raise ValueError("`phash_db` unset.")
+        return self.__phash_db
 
-    @sift_db.setter
-    def sift_db(self, value: SIFTDatabase):
-        self.__sift_db = value
+    @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):
@@ -86,7 +98,7 @@ class DeviceV2Ocr:
 
     def ocr_song_id(self, rois: DeviceV2Rois):
         jacket = cv2.cvtColor(rois.jacket, cv2.COLOR_BGR2GRAY)
-        return self.sift_db.lookup_img(jacket)[0]
+        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)
@@ -103,20 +115,33 @@ class DeviceV2Ocr:
                 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 = mask_gray(rois.pure)
+        roi = self.mask_pfl(rois.pure, rois)
         return self._base_ocr_pfl(roi, rois.sizes.factor)
 
     def ocr_far(self, rois: DeviceV2Rois):
-        roi = mask_gray(rois.far)
+        roi = self.mask_pfl(rois.far, rois)
         return self._base_ocr_pfl(roi, rois.sizes.factor)
 
     def ocr_lost(self, rois: DeviceV2Rois):
-        roi = mask_gray(rois.lost)
+        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_gray(rois.max_recall_rating_class)
+        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

src/arcaea_offline_ocr/device/v2/rois.py

@@ -1,183 +1,15 @@
-from typing import Tuple, Union
+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)
 
 
-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 DeviceV2Rois:
     def __init__(self, device: DeviceV2, img: Mat):
         self.device = device

src/arcaea_offline_ocr/device/v2/sizes.py

@@ -0,0 +1,254 @@
+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)

src/arcaea_offline_ocr/mask.py

@@ -8,6 +8,8 @@ __all__ = [
     "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",
@@ -16,13 +18,17 @@ __all__ = [
     "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)
@@ -34,6 +40,9 @@ 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)
 
@@ -46,6 +55,9 @@ 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)
@@ -63,6 +75,12 @@ def mask_white(img_hsv: Mat):
     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))
@@ -93,3 +111,9 @@ def mask_rating_class(img_hsv: Mat):
     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

src/arcaea_offline_ocr/phash_db.py

@@ -0,0 +1,65 @@
+import sqlite3
+
+import imagehash
+import numpy as np
+from PIL import Image
+
+
+def hamming_distance_sql_function(user_input, db_entry) -> int:
+    return np.count_nonzero(
+        np.frombuffer(user_input, bool) ^ np.frombuffer(db_entry, bool)
+    )
+
+
+class ImagePHashDatabase:
+    def __init__(self, db_path: str):
+        with sqlite3.connect(db_path) as conn:
+            self.hash_size = int(
+                conn.execute(
+                    "SELECT value FROM properties WHERE key = 'hash_size'"
+                ).fetchone()[0]
+            )
+            self.highfreq_factor = int(
+                conn.execute(
+                    "SELECT value FROM properties WHERE key = 'highfreq_factor'"
+                ).fetchone()[0]
+            )
+            self.built_timestamp = int(
+                conn.execute(
+                    "SELECT value FROM properties WHERE key = 'built_timestamp'"
+                ).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.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]
+
+    def lookup_hash(self, image_hash: imagehash.ImageHash, *, limit: int = 5):
+        image_hash = image_hash.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)
+        ]
+        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
+        )
+        return self.lookup_hash(image_hash)[0]