refactor!: device versions

src/arcaea_offline_ocr/device/shared.py

@@ -0,0 +1,14 @@
+import attrs
+
+
+@attrs.define
+class DeviceOcrResult:
+    song_id: None
+    title: None
+    rating_class: int
+    pure: int
+    far: int
+    lost: int
+    score: int
+    max_recall: int
+    clear_type: None

src/arcaea_offline_ocr/device/v1/crop.py

@@ -1,10 +1,7 @@
-from math import floor
-from typing import Any, Tuple
-
-import numpy as np
+from typing import Tuple
 
 from ...types import Mat
-from .definition import Device
+from .definition import DeviceV1
 
 __all__ = [
     "crop_img",
@@ -16,7 +13,6 @@ __all__ = [
     "crop_to_rating_class",
     "crop_to_score",
     "crop_to_title",
-    "crop_black_edges",
 ]
 
 
@@ -29,38 +25,29 @@ def crop_from_device_attr(img: Mat, rect: Tuple[int, int, int, int]):
     return crop_img(img, top=y, left=x, bottom=y + h, right=x + w)
 
 
-def crop_to_pure(screenshot: Mat, device: Device):
+def crop_to_pure(screenshot: Mat, device: DeviceV1):
     return crop_from_device_attr(screenshot, device.pure)
 
 
-def crop_to_far(screenshot: Mat, device: Device):
+def crop_to_far(screenshot: Mat, device: DeviceV1):
     return crop_from_device_attr(screenshot, device.far)
 
 
-def crop_to_lost(screenshot: Mat, device: Device):
+def crop_to_lost(screenshot: Mat, device: DeviceV1):
     return crop_from_device_attr(screenshot, device.lost)
 
 
-def crop_to_max_recall(screenshot: Mat, device: Device):
+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: Device):
+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: Device):
+def crop_to_score(screenshot: Mat, device: DeviceV1):
     return crop_from_device_attr(screenshot, device.score)
 
 
-def crop_to_title(screenshot: Mat, device: Device):
+def crop_to_title(screenshot: Mat, device: DeviceV1):
     return crop_from_device_attr(screenshot, device.title)
-
-
-def is_black_edge(list_of_pixels: Mat, black_pixel=None):
-    if black_pixel is None:
-        black_pixel = np.array([0, 0, 0], list_of_pixels.dtype)
-    pixels = list_of_pixels.reshape([-1, 3])
-    return np.count_nonzero(all(pixels < black_pixel, axis=1)) > floor(
-        len(pixels) * 0.6
-    )

src/arcaea_offline_ocr/device/v1/definition.py

@@ -1,11 +1,11 @@
 from dataclasses import dataclass
 from typing import Any, Dict, Tuple
 
-__all__ = ["Device"]
+__all__ = ["DeviceV1"]
 
 
 @dataclass(kw_only=True)
-class Device:
+class DeviceV1:
     version: int
     uuid: str
     name: str

src/arcaea_offline_ocr/device/v1/ocr.py

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

src/arcaea_offline_ocr/device/v2/ocr.py

@@ -1,31 +1,18 @@
-from typing import Optional
-
-import attrs
 import cv2
 import numpy as np
 
 from ...mask import mask_byd, mask_ftr, mask_gray, mask_prs, mask_pst, mask_white
 from ...ocr import ocr_digits_knn_model
 from ...types import Mat, cv2_ml_KNearest
+from ..shared import DeviceOcrResult
 from .find import find_digits
 from .rois import DeviceV2Rois
 
 
-@attrs.define
-class DeviceV2OcrResult:
-    pure: int
-    far: int
-    lost: int
-    score: int
-    rating_class: int
-    max_recall: int
-    title: Optional[str]
-
-
 class DeviceV2Ocr:
-    def __init__(self):
-        self.__rois = None
-        self.__knn_model = None
+    def __init__(self, rois: DeviceV2Rois, knn_model: cv2_ml_KNearest):
+        self.__rois = rois
+        self.__knn_model = knn_model
 
     @property
     def rois(self):