From 95987f6a2b233ebff33f49cda5d45cb20b39f902 Mon Sep 17 00:00:00 2001 From: 283375 Date: Wed, 2 Aug 2023 01:41:07 +0800 Subject: [PATCH] feat: new gray masking algorithm --- src/arcaea_offline_ocr/mask.py | 47 ++++++++++++++++++++++------------ 1 file changed, 31 insertions(+), 16 deletions(-) diff --git a/src/arcaea_offline_ocr/mask.py b/src/arcaea_offline_ocr/mask.py index 9704082..1300b4c 100644 --- a/src/arcaea_offline_ocr/mask.py +++ b/src/arcaea_offline_ocr/mask.py @@ -1,5 +1,5 @@ -from cv2 import BORDER_CONSTANT, BORDER_ISOLATED, bitwise_or, dilate, inRange -from numpy import array, uint8 +import cv2 +from numpy import array, max, min, uint8 from .types import Mat @@ -26,7 +26,10 @@ __all__ = [ ] GRAY_MIN_HSV = array([0, 0, 70], uint8) -GRAY_MAX_HSV = array([0, 70, 200], uint8) +GRAY_MAX_HSV = array([0, 0, 200], uint8) + +GRAY_MIN_BGR = array([50] * 3, uint8) +GRAY_MAX_BGR = array([160] * 3, uint8) WHITE_MIN_HSV = array([0, 0, 240], uint8) WHITE_MAX_HSV = array([179, 10, 255], uint8) @@ -44,39 +47,51 @@ BYD_MIN_HSV = array([170, 50, 50], uint8) BYD_MAX_HSV = array([179, 210, 198], uint8) +def mask_gray_new(__img_bgr: Mat): + # bgr_value_equal_mask = all(__img_bgr[:, 1:] == __img_bgr[:, :-1], axis=1) + bgr_value_equal_mask = max(__img_bgr, axis=2) - min(__img_bgr, axis=2) <= 5 + img_bgr = __img_bgr.copy() + img_bgr[~bgr_value_equal_mask] = array([0, 0, 0], __img_bgr.dtype) + img_bgr = cv2.erode(img_bgr, cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))) + img_bgr = cv2.dilate(img_bgr, cv2.getStructuringElement(cv2.MORPH_RECT, (1, 1))) + return cv2.inRange(img_bgr, GRAY_MIN_BGR, GRAY_MAX_BGR) + + def mask_gray(img_hsv: Mat): - mask = inRange(img_hsv, GRAY_MIN_HSV, GRAY_MAX_HSV) - mask = dilate(mask, (2, 2)) + mask = cv2.inRange(img_hsv, GRAY_MIN_HSV, GRAY_MAX_HSV) + mask = cv2.dilate(mask, (2, 2)) return mask def mask_white(img_hsv: Mat): - mask = inRange(img_hsv, WHITE_MIN_HSV, WHITE_MAX_HSV) - mask = dilate(mask, (5, 5), borderType=BORDER_CONSTANT | BORDER_ISOLATED) + mask = cv2.inRange(img_hsv, WHITE_MIN_HSV, WHITE_MAX_HSV) + mask = cv2.dilate( + mask, (5, 5), borderType=cv2.BORDER_CONSTANT | cv2.BORDER_ISOLATED + ) return mask def mask_pst(img_hsv: Mat): - mask = inRange(img_hsv, PST_MIN_HSV, PST_MAX_HSV) - mask = dilate(mask, (1, 1)) + mask = cv2.inRange(img_hsv, PST_MIN_HSV, PST_MAX_HSV) + mask = cv2.dilate(mask, (1, 1)) return mask def mask_prs(img_hsv: Mat): - mask = inRange(img_hsv, PRS_MIN_HSV, PRS_MAX_HSV) - mask = dilate(mask, (1, 1)) + mask = cv2.inRange(img_hsv, PRS_MIN_HSV, PRS_MAX_HSV) + mask = cv2.dilate(mask, (1, 1)) return mask def mask_ftr(img_hsv: Mat): - mask = inRange(img_hsv, FTR_MIN_HSV, FTR_MAX_HSV) - mask = dilate(mask, (1, 1)) + mask = cv2.inRange(img_hsv, FTR_MIN_HSV, FTR_MAX_HSV) + mask = cv2.dilate(mask, (1, 1)) return mask def mask_byd(img_hsv: Mat): - mask = inRange(img_hsv, BYD_MIN_HSV, BYD_MAX_HSV) - mask = dilate(mask, (2, 2)) + mask = cv2.inRange(img_hsv, BYD_MIN_HSV, BYD_MAX_HSV) + mask = cv2.dilate(mask, (2, 2)) return mask @@ -85,4 +100,4 @@ def mask_rating_class(img_hsv: Mat): prs = mask_prs(img_hsv) ftr = mask_ftr(img_hsv) byd = mask_byd(img_hsv) - return bitwise_or(byd, bitwise_or(ftr, bitwise_or(pst, prs))) + return cv2.bitwise_or(byd, cv2.bitwise_or(ftr, cv2.bitwise_or(pst, prs)))