diff --git a/src/arcaea_offline_ocr/crop.py b/src/arcaea_offline_ocr/crop.py index efc816b..59b663f 100644 --- a/src/arcaea_offline_ocr/crop.py +++ b/src/arcaea_offline_ocr/crop.py @@ -5,7 +5,7 @@ import numpy as np from .types import Mat -__all__ = ["crop_xywh", "crop_black_edges"] +__all__ = ["crop_xywh", "crop_black_edges", "crop_black_edges_grayscale"] def crop_xywh(mat: Mat, rect: Tuple[int, int, int, int]): @@ -56,3 +56,51 @@ def crop_black_edges(screenshot: Mat): bottom -= 1 return cropped[top:bottom, left:right] + + +def is_black_edge_grayscale( + gray_value_list: np.ndarray, black_threshold: int = 50, ratio: float = 0.6 +) -> bool: + return ( + np.count_nonzero(gray_value_list < black_threshold) + > len(gray_value_list) * ratio + ) + + +def crop_black_edges_grayscale( + img_gray: Mat, black_threshold: int = 50 +) -> Tuple[int, int, int, int]: + """Returns cropped rect""" + height, width = img_gray.shape[:2] + left = 0 + right = width + top = 0 + bottom = height + + for i in range(width): + column = img_gray[:, i] + if not is_black_edge_grayscale(column, black_threshold): + break + left += 1 + + for i in sorted(range(width), reverse=True): + column = img_gray[:, i] + if i <= left + 1 or not is_black_edge_grayscale(column, black_threshold): + break + right -= 1 + + for i in range(height): + row = img_gray[i] + if not is_black_edge_grayscale(row, black_threshold): + break + top += 1 + + for i in sorted(range(height), reverse=True): + row = img_gray[i] + if i <= top + 1 or not is_black_edge_grayscale(row, black_threshold): + break + bottom -= 1 + + assert right > left, "cropped width > 0" + assert bottom > top, "cropped height > 0" + return (left, top, right - left, bottom - top)