refactor!: device versions

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,
+        )