283375/arcaea-offline-ocr
1
0
Fork 0
mirror of https://github.com/283375/arcaea-offline-ocr.git synced 2025-04-18 21:10:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

impr: split common preprocess functions

Browse Source
This commit is contained in:
283375 2023-06-10 17:27:43 +08:00
parent c7e870f42e
commit 4f26fd1b59
1 changed files with 25 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
src/arcaea_offline_ocr/recognize.py
View File

@ -1,7 +1,7 @@
from dataclasses import dataclass from dataclasses import dataclass
from typing import Optional from typing import Callable, Optional
from cv2 import COLOR_BGR2HSV, GaussianBlur, cvtColor, imread from cv2 import COLOR_BGR2HSV, GaussianBlur, Mat, cvtColor, imread
from .crop import * from .crop import *
from .device import Device from .device import Device
@ -9,6 +9,18 @@ from .mask import *
from .ocr import * from .ocr import *
def process_digit_ocr_img(img_hsv, mask=Callable[[Mat], Mat]):
img_hsv = mask(img_hsv)
img_hsv = GaussianBlur(img_hsv, (3, 3), 0)
return img_hsv
def process_tesseract_ocr_img(img_hsv, mask=Callable[[Mat], Mat]):
img_hsv = mask(img_hsv)
img_hsv = GaussianBlur(img_hsv, (1, 1), 0)
return img_hsv
@dataclass(kw_only=True) @dataclass(kw_only=True)
class RecognizeResult: class RecognizeResult:
pure: Optional[int] pure: Optional[int]
@ -25,36 +37,29 @@ def recognize(img_filename: str, device: Device):
img_hsv = cvtColor(img, COLOR_BGR2HSV) img_hsv = cvtColor(img, COLOR_BGR2HSV)
pure_roi = crop_to_pure(img_hsv, device) pure_roi = crop_to_pure(img_hsv, device)
pure_roi = mask_gray(pure_roi) pure = ocr_pure(process_digit_ocr_img(pure_roi, mask=mask_gray))
pure_roi = GaussianBlur(pure_roi, (3, 3), 0)
pure = ocr_pure(pure_roi)
far_roi = crop_to_far(img_hsv, device) far_roi = crop_to_far(img_hsv, device)
far_roi = mask_gray(far_roi) far = ocr_far_lost(process_digit_ocr_img(far_roi, mask=mask_gray))
far_roi = GaussianBlur(far_roi, (3, 3), 0)
far = ocr_far_lost(far_roi)
lost_roi = crop_to_lost(img_hsv, device) lost_roi = crop_to_lost(img_hsv, device)
lost_roi = mask_gray(lost_roi) lost = ocr_far_lost(process_digit_ocr_img(lost_roi, mask=mask_gray))
lost_roi = GaussianBlur(lost_roi, (3, 3), 0)
lost = ocr_far_lost(lost_roi)
score_roi = crop_to_score(img_hsv, device) score_roi = crop_to_score(img_hsv, device)
score_roi = mask_white(score_roi) score = ocr_score(process_digit_ocr_img(score_roi, mask=mask_white))
score_roi = GaussianBlur(score_roi, (3, 3), 0)
score = ocr_score(score_roi)
max_recall_roi = crop_to_max_recall(img_hsv, device) max_recall_roi = crop_to_max_recall(img_hsv, device)
max_recall_roi = mask_gray(max_recall_roi) max_recall = ocr_max_recall(
max_recall = ocr_max_recall(max_recall_roi) process_tesseract_ocr_img(max_recall_roi, mask=mask_gray)
)
rating_class_roi = crop_to_rating_class(img_hsv, device) rating_class_roi = crop_to_rating_class(img_hsv, device)
rating_class_roi = mask_rating_class(rating_class_roi) rating_class = ocr_rating_class(
rating_class = ocr_rating_class(rating_class_roi) process_tesseract_ocr_img(rating_class_roi, mask=mask_rating_class)
)
title_roi = crop_to_title(img_hsv, device) title_roi = crop_to_title(img_hsv, device)
title_roi = mask_white(title_roi) title = ocr_title(process_tesseract_ocr_img(title_roi, mask=mask_white))
title = ocr_title(title_roi)
return RecognizeResult( return RecognizeResult(
pure=pure, pure=pure,