mirror of
https://github.com/283375/arcaea-offline-ocr.git
synced 2025-07-02 12:56:28 +00:00
87 lines
2.7 KiB
Python
87 lines
2.7 KiB
Python
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,
|
|
)
|