refactor: CropBlackEdges

src/arcaea_offline_ocr/crop.py

@@ -1,10 +1,10 @@
-from math import floor
+import math
 from typing import Tuple
 
 import cv2
 import numpy as np
 
-__all__ = ["crop_xywh", "crop_black_edges", "crop_black_edges_grayscale"]
+__all__ = ["crop_xywh", "CropBlackEdges"]
 
 
 def crop_xywh(mat: cv2.Mat, rect: Tuple[int, int, int, int]):
@@ -12,92 +12,53 @@ def crop_xywh(mat: cv2.Mat, rect: Tuple[int, int, int, int]):
     return mat[y : y + h, x : x + w]
 
 
-def is_black_edge(list_of_pixels: cv2.Mat, black_pixel: cv2.Mat, ratio: float = 0.6):
-    pixels = list_of_pixels.reshape([-1, 3])
-    return np.count_nonzero(np.all(pixels < black_pixel, axis=1)) > floor(
-        len(pixels) * ratio
-    )
+class CropBlackEdges:
+    @staticmethod
+    def is_black_edge(__img_gray_slice: cv2.Mat, black_pixel: int, ratio: float = 0.6):
+        pixels_compared = __img_gray_slice < black_pixel
+        return np.count_nonzero(pixels_compared) > math.floor(
+            __img_gray_slice.size * ratio
+        )
 
+    @classmethod
+    def get_crop_rect(cls, img_gray: cv2.Mat, black_threshold: int = 25):
+        height, width = img_gray.shape[:2]
+        left = 0
+        right = width
+        top = 0
+        bottom = height
 
-def crop_black_edges(img_bgr: cv2.Mat, black_threshold: int = 50):
-    cropped = img_bgr.copy()
-    black_pixel = np.array([black_threshold] * 3, img_bgr.dtype)
-    height, width = img_bgr.shape[:2]
-    left = 0
-    right = width
-    top = 0
-    bottom = height
+        for i in range(width):
+            column = img_gray[:, i]
+            if not cls.is_black_edge(column, black_threshold):
+                break
+            left += 1
 
-    for i in range(width):
-        column = cropped[:, i]
-        if not is_black_edge(column, black_pixel):
-            break
-        left += 1
+        for i in sorted(range(width), reverse=True):
+            column = img_gray[:, i]
+            if i <= left + 1 or not cls.is_black_edge(column, black_threshold):
+                break
+            right -= 1
 
-    for i in sorted(range(width), reverse=True):
-        column = cropped[:, i]
-        if i <= left + 1 or not is_black_edge(column, black_pixel):
-            break
-        right -= 1
+        for i in range(height):
+            row = img_gray[i]
+            if not cls.is_black_edge(row, black_threshold):
+                break
+            top += 1
 
-    for i in range(height):
-        row = cropped[i]
-        if not is_black_edge(row, black_pixel):
-            break
-        top += 1
+        for i in sorted(range(height), reverse=True):
+            row = img_gray[i]
+            if i <= top + 1 or not cls.is_black_edge(row, black_threshold):
+                break
+            bottom -= 1
 
-    for i in sorted(range(height), reverse=True):
-        row = cropped[i]
-        if i <= top + 1 or not is_black_edge(row, black_pixel):
-            break
-        bottom -= 1
+        assert right > left, "cropped width < 0"
+        assert bottom > top, "cropped height < 0"
+        return (left, top, right - left, bottom - top)
 
-    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: cv2.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)
+    @classmethod
+    def crop(
+        cls, img: cv2.Mat, convert_flag: cv2.COLOR_BGR2GRAY, black_threshold: int = 25
+    ) -> cv2.Mat:
+        rect = cls.get_crop_rect(cv2.cvtColor(img, convert_flag), black_threshold)
+        return crop_xywh(img, rect)