chore!: add recognize parts

This commit is contained in:
283375 2023-06-10 18:32:23 +08:00
parent 4f26fd1b59
commit ea31aaf00c
2 changed files with 46 additions and 20 deletions

View File

@ -88,7 +88,11 @@ def filter_digit_results(
if possible_results: if possible_results:
# Sort the results based on "max_val" in descending order and select the top result # Sort the results based on "max_val" in descending order and select the top result
result = sorted(possible_results, key=lambda d: d["max_val"], reverse=True)[0] result = sorted(
possible_results,
key=lambda d: d["max_val"],
reverse=True,
)[0]
result_digit = result["digit"] # Get the digit value from the result result_digit = result["digit"] # Get the digit value from the result
result.pop("digit", None) # Remove the digit key from the result result.pop("digit", None) # Remove the digit key from the result

View File

@ -9,16 +9,42 @@ from .mask import *
from .ocr import * from .ocr import *
def process_digit_ocr_img(img_hsv, mask=Callable[[Mat], Mat]): def process_digits_ocr_img(img_hsv_cropped: Mat, mask=Callable[[Mat], Mat]):
img_hsv = mask(img_hsv) img_hsv_cropped = mask(img_hsv_cropped)
img_hsv = GaussianBlur(img_hsv, (3, 3), 0) img_hsv_cropped = GaussianBlur(img_hsv_cropped, (3, 3), 0)
return img_hsv return img_hsv_cropped
def process_tesseract_ocr_img(img_hsv, mask=Callable[[Mat], Mat]): def process_tesseract_ocr_img(img_hsv_cropped: Mat, mask=Callable[[Mat], Mat]):
img_hsv = mask(img_hsv) img_hsv_cropped = mask(img_hsv_cropped)
img_hsv = GaussianBlur(img_hsv, (1, 1), 0) img_hsv_cropped = GaussianBlur(img_hsv_cropped, (1, 1), 0)
return img_hsv return img_hsv_cropped
def recognize_pure(img_hsv_cropped: Mat):
return ocr_pure(process_digits_ocr_img(img_hsv_cropped, mask=mask_gray))
def recognize_far_lost(img_hsv_cropped: Mat):
return ocr_far_lost(process_digits_ocr_img(img_hsv_cropped, mask=mask_gray))
def recognize_score(img_hsv_cropped: Mat):
return ocr_score(process_digits_ocr_img(img_hsv_cropped, mask=mask_white))
def recognize_max_recall(img_hsv_cropped: Mat):
return ocr_max_recall(process_tesseract_ocr_img(img_hsv_cropped, mask=mask_gray))
def recognize_rating_class(img_hsv_cropped: Mat):
return ocr_rating_class(
process_tesseract_ocr_img(img_hsv_cropped, mask=mask_rating_class)
)
def recognize_title(img_hsv_cropped: Mat):
return ocr_title(process_tesseract_ocr_img(img_hsv_cropped, mask=mask_white))
@dataclass(kw_only=True) @dataclass(kw_only=True)
@ -37,29 +63,25 @@ 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 = ocr_pure(process_digit_ocr_img(pure_roi, mask=mask_gray)) pure = recognize_pure(pure_roi)
far_roi = crop_to_far(img_hsv, device) far_roi = crop_to_far(img_hsv, device)
far = ocr_far_lost(process_digit_ocr_img(far_roi, mask=mask_gray)) far = recognize_far_lost(far_roi)
lost_roi = crop_to_lost(img_hsv, device) lost_roi = crop_to_lost(img_hsv, device)
lost = ocr_far_lost(process_digit_ocr_img(lost_roi, mask=mask_gray)) lost = recognize_far_lost(lost_roi)
score_roi = crop_to_score(img_hsv, device) score_roi = crop_to_score(img_hsv, device)
score = ocr_score(process_digit_ocr_img(score_roi, mask=mask_white)) score = recognize_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 = ocr_max_recall( max_recall = recognize_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 = ocr_rating_class( rating_class = recognize_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 = ocr_title(process_tesseract_ocr_img(title_roi, mask=mask_white)) title = recognize_title(title_roi)
return RecognizeResult( return RecognizeResult(
pure=pure, pure=pure,