diff --git a/src/arcaea_offline_ocr/crop.py b/src/arcaea_offline_ocr/crop.py index 7af0b34..2a8e5ad 100644 --- a/src/arcaea_offline_ocr/crop.py +++ b/src/arcaea_offline_ocr/crop.py @@ -1,7 +1,7 @@ from math import floor from typing import Tuple -from numpy import all, array, count_nonzero +import numpy as np from .types import Mat @@ -15,14 +15,16 @@ def crop_xywh(mat: Mat, rect: Tuple[int, int, int, int]): def is_black_edge(list_of_pixels: Mat, black_pixel=None): if black_pixel is None: - black_pixel = array([0, 0, 0], list_of_pixels.dtype) + black_pixel = np.array([0, 0, 0], list_of_pixels.dtype) pixels = list_of_pixels.reshape([-1, 3]) - return count_nonzero(all(pixels < black_pixel, axis=1)) > floor(len(pixels) * 0.6) + return np.count_nonzero(all(pixels < black_pixel, axis=1)) > floor( + len(pixels) * 0.6 + ) def crop_black_edges(screenshot: Mat): cropped = screenshot.copy() - black_pixel = array([50, 50, 50], screenshot.dtype) + black_pixel = np.array([50, 50, 50], screenshot.dtype) height, width = screenshot.shape[:2] left = 0 right = width diff --git a/src/arcaea_offline_ocr/device/v1/crop.py b/src/arcaea_offline_ocr/device/v1/crop.py index 5977d15..15922f2 100644 --- a/src/arcaea_offline_ocr/device/v1/crop.py +++ b/src/arcaea_offline_ocr/device/v1/crop.py @@ -1,7 +1,7 @@ from math import floor from typing import Any, Tuple -from numpy import all, array, count_nonzero +import numpy as np from ...types import Mat from .definition import Device @@ -59,6 +59,8 @@ def crop_to_title(screenshot: Mat, device: Device): def is_black_edge(list_of_pixels: Mat, black_pixel=None): if black_pixel is None: - black_pixel = array([0, 0, 0], list_of_pixels.dtype) + black_pixel = np.array([0, 0, 0], list_of_pixels.dtype) pixels = list_of_pixels.reshape([-1, 3]) - return count_nonzero(all(pixels < black_pixel, axis=1)) > floor(len(pixels) * 0.6) + return np.count_nonzero(all(pixels < black_pixel, axis=1)) > floor( + len(pixels) * 0.6 + ) diff --git a/src/arcaea_offline_ocr/device/v2/shared.py b/src/arcaea_offline_ocr/device/v2/shared.py index bc68755..ca511b1 100644 --- a/src/arcaea_offline_ocr/device/v2/shared.py +++ b/src/arcaea_offline_ocr/device/v2/shared.py @@ -1,4 +1,4 @@ -from cv2 import MORPH_CROSS, MORPH_ELLIPSE, MORPH_RECT, getStructuringElement +from cv2 import MORPH_RECT, getStructuringElement PFL_DENOISE_KERNEL = getStructuringElement(MORPH_RECT, [2, 2]) PFL_ERODE_KERNEL = getStructuringElement(MORPH_RECT, [3, 3]) diff --git a/src/arcaea_offline_ocr/mask.py b/src/arcaea_offline_ocr/mask.py index 23c9454..08de893 100644 --- a/src/arcaea_offline_ocr/mask.py +++ b/src/arcaea_offline_ocr/mask.py @@ -1,5 +1,5 @@ import cv2 -from numpy import array, max, min, uint8 +import numpy as np from .types import Mat @@ -25,33 +25,33 @@ __all__ = [ "mask_rating_class", ] -GRAY_MIN_HSV = array([0, 0, 70], uint8) -GRAY_MAX_HSV = array([0, 0, 200], uint8) +GRAY_MIN_HSV = np.array([0, 0, 70], np.uint8) +GRAY_MAX_HSV = np.array([0, 0, 200], np.uint8) -GRAY_MIN_BGR = array([50] * 3, uint8) -GRAY_MAX_BGR = array([160] * 3, uint8) +GRAY_MIN_BGR = np.array([50] * 3, np.uint8) +GRAY_MAX_BGR = np.array([160] * 3, np.uint8) -WHITE_MIN_HSV = array([0, 0, 240], uint8) -WHITE_MAX_HSV = array([179, 10, 255], uint8) +WHITE_MIN_HSV = np.array([0, 0, 240], np.uint8) +WHITE_MAX_HSV = np.array([179, 10, 255], np.uint8) -PST_MIN_HSV = array([100, 50, 80], uint8) -PST_MAX_HSV = array([100, 255, 255], uint8) +PST_MIN_HSV = np.array([100, 50, 80], np.uint8) +PST_MAX_HSV = np.array([100, 255, 255], np.uint8) -PRS_MIN_HSV = array([43, 40, 75], uint8) -PRS_MAX_HSV = array([50, 155, 190], uint8) +PRS_MIN_HSV = np.array([43, 40, 75], np.uint8) +PRS_MAX_HSV = np.array([50, 155, 190], np.uint8) -FTR_MIN_HSV = array([149, 30, 0], uint8) -FTR_MAX_HSV = array([155, 181, 150], uint8) +FTR_MIN_HSV = np.array([149, 30, 0], np.uint8) +FTR_MAX_HSV = np.array([155, 181, 150], np.uint8) -BYD_MIN_HSV = array([170, 50, 50], uint8) -BYD_MAX_HSV = array([179, 210, 198], uint8) +BYD_MIN_HSV = np.array([170, 50, 50], np.uint8) +BYD_MAX_HSV = np.array([179, 210, 198], np.uint8) def mask_gray(__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[~bgr_value_equal_mask] = np.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) diff --git a/src/arcaea_offline_ocr/recognize.py b/src/arcaea_offline_ocr/recognize.py index 8c666e3..80f338c 100644 --- a/src/arcaea_offline_ocr/recognize.py +++ b/src/arcaea_offline_ocr/recognize.py @@ -1,7 +1,7 @@ from dataclasses import dataclass from typing import Callable, Optional -from cv2 import COLOR_BGR2HSV, GaussianBlur, cvtColor, imread +import cv2 from .crop import * @@ -29,13 +29,13 @@ __all__ = [ def process_digits_ocr_img(img_hsv_cropped: Mat, mask=Callable[[Mat], Mat]): img_hsv_cropped = mask(img_hsv_cropped) - img_hsv_cropped = GaussianBlur(img_hsv_cropped, (3, 3), 0) + img_hsv_cropped = cv2.GaussianBlur(img_hsv_cropped, (3, 3), 0) return img_hsv_cropped def process_tesseract_ocr_img(img_hsv_cropped: Mat, mask=Callable[[Mat], Mat]): img_hsv_cropped = mask(img_hsv_cropped) - img_hsv_cropped = GaussianBlur(img_hsv_cropped, (1, 1), 0) + img_hsv_cropped = cv2.GaussianBlur(img_hsv_cropped, (1, 1), 0) return img_hsv_cropped @@ -78,7 +78,7 @@ class RecognizeResult: def recognize(img_filename: str, device: Device): img = imread_unicode(img_filename) - img_hsv = cvtColor(img, COLOR_BGR2HSV) + img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) pure_roi = crop_to_pure(img_hsv, device) pure = recognize_pure(pure_roi) diff --git a/src/arcaea_offline_ocr/template.py b/src/arcaea_offline_ocr/template.py index a2ab2a2..037f646 100644 --- a/src/arcaea_offline_ocr/template.py +++ b/src/arcaea_offline_ocr/template.py @@ -2,30 +2,8 @@ import pickle from base64 import b64decode from typing import Any, Dict, List, Literal, Tuple, TypedDict, Union -from cv2 import ( - CHAIN_APPROX_SIMPLE, - COLOR_BGR2GRAY, - COLOR_GRAY2BGR, - FONT_HERSHEY_SIMPLEX, - IMREAD_GRAYSCALE, - RETR_EXTERNAL, - THRESH_BINARY_INV, - TM_CCOEFF_NORMED, - boundingRect, - cvtColor, - destroyAllWindows, - findContours, - imdecode, - imread, - imshow, - matchTemplate, - minMaxLoc, - putText, - rectangle, - threshold, - waitKey, -) -from numpy import ndarray +import cv2 +import numpy as np from ._builtin_templates import ( DEFAULT_ITALIC, @@ -60,7 +38,7 @@ class DigitTemplate: return ( isinstance(item, (list, tuple)) and len(item) == 11 - and all(isinstance(i, ndarray) for i in item) + and all(isinstance(i, np.ndarray) for i in item) ) def __init__(self, regular, italic, regular_eroded, italic_eroded): @@ -109,8 +87,8 @@ class MatchTemplateMultipleResult(TypedDict): def matchTemplateMultiple( src: Mat, template: Mat, threshold: float = 0.1 ) -> List[MatchTemplateMultipleResult]: - template_result = matchTemplate(src, template, TM_CCOEFF_NORMED) - min_val, max_val, min_loc, max_loc = minMaxLoc(template_result) + template_result = cv2.matchTemplate(src, template, cv2.TM_CCOEFF_NORMED) + min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(template_result) template_h, template_w = template.shape[:2] results = [] @@ -123,7 +101,7 @@ def matchTemplateMultiple( # CC BY-SA 4.0 prev_min_val, prev_max_val, prev_min_loc, prev_max_loc = None, None, None, None while max_val > threshold: - min_val, max_val, min_loc, max_loc = minMaxLoc(template_result) + min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(template_result) # Prevent infinite loop. If those 4 values are the same as previous ones, break the loop. if ( diff --git a/src/arcaea_offline_ocr/utils.py b/src/arcaea_offline_ocr/utils.py index aa771bf..33eec22 100644 --- a/src/arcaea_offline_ocr/utils.py +++ b/src/arcaea_offline_ocr/utils.py @@ -1,9 +1,8 @@ from collections.abc import Iterable from typing import Callable, Tuple, TypeVar, Union, overload -from cv2 import IMREAD_UNCHANGED, imdecode -from numpy import fromfile as np_fromfile -from numpy import uint8 +import cv2 +import numpy as np from .types import Mat, XYWHRect @@ -13,7 +12,7 @@ __all__ = ["imread_unicode"] def imread_unicode(filepath: str) -> Mat: # https://stackoverflow.com/a/57872297/16484891 # CC BY-SA 4.0 - return imdecode(np_fromfile(filepath, dtype=uint8), IMREAD_UNCHANGED) + return cv2.imdecode(np.fromfile(filepath, dtype=np.uint8), cv2.IMREAD_UNCHANGED) def construct_int_xywh_rect(