diff --git a/src/arcaea_offline_ocr/device/shared.py b/src/arcaea_offline_ocr/device/common.py similarity index 100% rename from src/arcaea_offline_ocr/device/shared.py rename to src/arcaea_offline_ocr/device/common.py diff --git a/src/arcaea_offline_ocr/ocr_new.py b/src/arcaea_offline_ocr/device/ocr.py similarity index 93% rename from src/arcaea_offline_ocr/ocr_new.py rename to src/arcaea_offline_ocr/device/ocr.py index 127041c..bf5b9a6 100644 --- a/src/arcaea_offline_ocr/ocr_new.py +++ b/src/arcaea_offline_ocr/device/ocr.py @@ -2,24 +2,24 @@ import cv2 import numpy as np from PIL import Image -from .crop import crop_xywh -from .extractor import Extractor -from .masker import Masker -from .ocr import ( +from ..crop import crop_xywh +from ..ocr import ( FixRects, ocr_digit_samples_knn, ocr_digits_by_contour_knn, preprocess_hog, resize_fill_square, ) -from .phash_db import ImagePHashDatabase +from ..phash_db import ImagePHashDatabase +from .roi.extractor import DeviceRoiExtractor +from .roi.masker import DeviceRoiMasker class DeviceOcr: def __init__( self, - extractor: Extractor, - masker: Masker, + extractor: DeviceRoiExtractor, + masker: DeviceRoiMasker, knn_model: cv2.ml.KNearest, phash_db: ImagePHashDatabase, ): diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/__init__.py b/src/arcaea_offline_ocr/device/roi/definitions/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/__init__.py rename to src/arcaea_offline_ocr/device/roi/definitions/__init__.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/__init__.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/auto/__init__.py rename to src/arcaea_offline_ocr/device/roi/definitions/auto/__init__.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/common.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/common.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/auto/common.py rename to src/arcaea_offline_ocr/device/roi/definitions/auto/common.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t1.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/t1.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/auto/t1.py rename to src/arcaea_offline_ocr/device/roi/definitions/auto/t1.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/auto/t2.py b/src/arcaea_offline_ocr/device/roi/definitions/auto/t2.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/auto/t2.py rename to src/arcaea_offline_ocr/device/roi/definitions/auto/t2.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/common.py b/src/arcaea_offline_ocr/device/roi/definitions/common.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/common.py rename to src/arcaea_offline_ocr/device/roi/definitions/common.py diff --git a/src/arcaea_offline_ocr/roi_extractor/sizes/custom.py b/src/arcaea_offline_ocr/device/roi/definitions/custom.py similarity index 100% rename from src/arcaea_offline_ocr/roi_extractor/sizes/custom.py rename to src/arcaea_offline_ocr/device/roi/definitions/custom.py diff --git a/src/arcaea_offline_ocr/roi_extractor/__init__.py b/src/arcaea_offline_ocr/device/roi/extractor/__init__.py similarity index 65% rename from src/arcaea_offline_ocr/roi_extractor/__init__.py rename to src/arcaea_offline_ocr/device/roi/extractor/__init__.py index 050378c..66ae350 100644 --- a/src/arcaea_offline_ocr/roi_extractor/__init__.py +++ b/src/arcaea_offline_ocr/device/roi/extractor/__init__.py @@ -1,2 +1 @@ from .common import DeviceRoiExtractor -from .sizes import * diff --git a/src/arcaea_offline_ocr/roi_extractor/common.py b/src/arcaea_offline_ocr/device/roi/extractor/common.py similarity index 94% rename from src/arcaea_offline_ocr/roi_extractor/common.py rename to src/arcaea_offline_ocr/device/roi/extractor/common.py index 8bb106e..d9365e2 100644 --- a/src/arcaea_offline_ocr/roi_extractor/common.py +++ b/src/arcaea_offline_ocr/device/roi/extractor/common.py @@ -1,7 +1,7 @@ import cv2 -from ..crop import crop_xywh -from .sizes.common import DeviceRoiSizes +from ....crop import crop_xywh +from ..definitions.common import DeviceRoiSizes class DeviceRoiExtractor: diff --git a/src/arcaea_offline_ocr/masker/__init__.py b/src/arcaea_offline_ocr/device/roi/masker/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/masker/__init__.py rename to src/arcaea_offline_ocr/device/roi/masker/__init__.py diff --git a/src/arcaea_offline_ocr/masker/auto/__init__.py b/src/arcaea_offline_ocr/device/roi/masker/auto/__init__.py similarity index 100% rename from src/arcaea_offline_ocr/masker/auto/__init__.py rename to src/arcaea_offline_ocr/device/roi/masker/auto/__init__.py diff --git a/src/arcaea_offline_ocr/masker/auto/common.py b/src/arcaea_offline_ocr/device/roi/masker/auto/common.py similarity index 100% rename from src/arcaea_offline_ocr/masker/auto/common.py rename to src/arcaea_offline_ocr/device/roi/masker/auto/common.py diff --git a/src/arcaea_offline_ocr/masker/auto/t1.py b/src/arcaea_offline_ocr/device/roi/masker/auto/t1.py similarity index 100% rename from src/arcaea_offline_ocr/masker/auto/t1.py rename to src/arcaea_offline_ocr/device/roi/masker/auto/t1.py diff --git a/src/arcaea_offline_ocr/masker/auto/t2.py b/src/arcaea_offline_ocr/device/roi/masker/auto/t2.py similarity index 100% rename from src/arcaea_offline_ocr/masker/auto/t2.py rename to src/arcaea_offline_ocr/device/roi/masker/auto/t2.py diff --git a/src/arcaea_offline_ocr/masker/common.py b/src/arcaea_offline_ocr/device/roi/masker/common.py similarity index 100% rename from src/arcaea_offline_ocr/masker/common.py rename to src/arcaea_offline_ocr/device/roi/masker/common.py diff --git a/src/arcaea_offline_ocr/device/v1/__init__.py b/src/arcaea_offline_ocr/device/v1/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/arcaea_offline_ocr/device/v1/crop.py b/src/arcaea_offline_ocr/device/v1/crop.py deleted file mode 100644 index 2f7d6a2..0000000 --- a/src/arcaea_offline_ocr/device/v1/crop.py +++ /dev/null @@ -1,53 +0,0 @@ -from typing import Tuple - -from ...types import Mat -from .definition import DeviceV1 - -__all__ = [ - "crop_img", - "crop_from_device_attr", - "crop_to_pure", - "crop_to_far", - "crop_to_lost", - "crop_to_max_recall", - "crop_to_rating_class", - "crop_to_score", - "crop_to_title", -] - - -def crop_img(img: Mat, *, top: int, left: int, bottom: int, right: int): - return img[top:bottom, left:right] - - -def crop_from_device_attr(img: Mat, rect: Tuple[int, int, int, int]): - x, y, w, h = rect - return crop_img(img, top=y, left=x, bottom=y + h, right=x + w) - - -def crop_to_pure(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.pure) - - -def crop_to_far(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.far) - - -def crop_to_lost(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.lost) - - -def crop_to_max_recall(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.max_recall) - - -def crop_to_rating_class(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.rating_class) - - -def crop_to_score(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.score) - - -def crop_to_title(screenshot: Mat, device: DeviceV1): - return crop_from_device_attr(screenshot, device.title) diff --git a/src/arcaea_offline_ocr/device/v1/definition.py b/src/arcaea_offline_ocr/device/v1/definition.py deleted file mode 100644 index 51ca29c..0000000 --- a/src/arcaea_offline_ocr/device/v1/definition.py +++ /dev/null @@ -1,37 +0,0 @@ -from dataclasses import dataclass -from typing import Any, Dict, Tuple - -__all__ = ["DeviceV1"] - - -@dataclass(kw_only=True) -class DeviceV1: - version: int - uuid: str - name: str - pure: Tuple[int, int, int, int] - far: Tuple[int, int, int, int] - lost: Tuple[int, int, int, int] - max_recall: Tuple[int, int, int, int] - rating_class: Tuple[int, int, int, int] - score: Tuple[int, int, int, int] - title: Tuple[int, int, int, int] - - @classmethod - def from_json_object(cls, json_dict: Dict[str, Any]): - if json_dict["version"] == 1: - return cls( - version=1, - uuid=json_dict["uuid"], - name=json_dict["name"], - pure=json_dict["pure"], - far=json_dict["far"], - lost=json_dict["lost"], - max_recall=json_dict["max_recall"], - rating_class=json_dict["rating_class"], - score=json_dict["score"], - title=json_dict["title"], - ) - - def repr_info(self): - return f"Device(version={self.version}, uuid={repr(self.uuid)}, name={repr(self.name)})" diff --git a/src/arcaea_offline_ocr/device/v1/ocr.py b/src/arcaea_offline_ocr/device/v1/ocr.py deleted file mode 100644 index 2e3a4b1..0000000 --- a/src/arcaea_offline_ocr/device/v1/ocr.py +++ /dev/null @@ -1,86 +0,0 @@ -from typing import List - -import cv2 - -from ...crop import crop_xywh -from ...mask import mask_gray, mask_white -from ...ocr import ocr_digits_by_contour_knn, ocr_rating_class -from ...types import Mat, cv2_ml_KNearest -from ..shared import DeviceOcrResult -from .crop import * -from .definition import DeviceV1 - - -class DeviceV1Ocr: - def __init__(self, device: DeviceV1, knn_model: cv2_ml_KNearest): - self.__device = device - self.__knn_model = knn_model - - @property - def device(self): - return self.__device - - @device.setter - def device(self, value): - self.__device = value - - @property - def knn_model(self): - return self.__knn_model - - @knn_model.setter - def knn_model(self, value): - self.__knn_model = value - - def preprocess_score_roi(self, __roi_gray: Mat) -> List[Mat]: - roi_gray = __roi_gray.copy() - contours, _ = cv2.findContours( - roi_gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE - ) - for contour in contours: - rect = cv2.boundingRect(contour) - if rect[3] > roi_gray.shape[0] * 0.6: - continue - roi_gray = cv2.fillPoly(roi_gray, [contour], 0) - return roi_gray - - def ocr(self, img_bgr: Mat): - rating_class_roi = crop_to_rating_class(img_bgr, self.device) - rating_class = ocr_rating_class(rating_class_roi) - - pfl_mr_roi = [ - crop_to_pure(img_bgr, self.device), - crop_to_far(img_bgr, self.device), - crop_to_lost(img_bgr, self.device), - crop_to_max_recall(img_bgr, self.device), - ] - pfl_mr_roi = [mask_gray(roi) for roi in pfl_mr_roi] - - pure, far, lost = [ - ocr_digits_by_contour_knn(roi, self.knn_model) for roi in pfl_mr_roi[:3] - ] - - max_recall_contours, _ = cv2.findContours( - pfl_mr_roi[3], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE - ) - max_recall_rects = [cv2.boundingRect(c) for c in max_recall_contours] - max_recall_rect = sorted(max_recall_rects, key=lambda r: r[0])[-1] - max_recall_roi = crop_xywh(img_bgr, max_recall_rect) - max_recall = ocr_digits_by_contour_knn(max_recall_roi, self.knn_model) - - score_roi = crop_to_score(img_bgr, self.device) - score_roi = mask_white(score_roi) - score_roi = self.preprocess_score_roi(score_roi) - score = ocr_digits_by_contour_knn(score_roi, self.knn_model) - - return DeviceOcrResult( - song_id=None, - title=None, - rating_class=rating_class, - pure=pure, - far=far, - lost=lost, - score=score, - max_recall=max_recall, - clear_type=None, - ) diff --git a/src/arcaea_offline_ocr/device/v2/__init__.py b/src/arcaea_offline_ocr/device/v2/__init__.py deleted file mode 100644 index 64db9c3..0000000 --- a/src/arcaea_offline_ocr/device/v2/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .definition import DeviceV2 -from .ocr import DeviceV2Ocr -from .rois import DeviceV2AutoRois, DeviceV2Rois -from .shared import MAX_RECALL_CLOSE_KERNEL diff --git a/src/arcaea_offline_ocr/device/v2/definition.py b/src/arcaea_offline_ocr/device/v2/definition.py deleted file mode 100644 index 31dd17a..0000000 --- a/src/arcaea_offline_ocr/device/v2/definition.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing import Iterable - -from attrs import define, field - -from ...types import XYWHRect - - -def iterable_to_xywh_rect(__iter: Iterable) -> XYWHRect: - return XYWHRect(*__iter) - - -@define(kw_only=True) -class DeviceV2: - version = field(type=int) - uuid = field(type=str) - name = field(type=str) - crop_black_edges = field(type=bool) - factor = field(type=float) - pure = field(converter=iterable_to_xywh_rect, default=[0, 0, 0, 0]) - far = field(converter=iterable_to_xywh_rect, default=[0, 0, 0, 0]) - lost = field(converter=iterable_to_xywh_rect, default=[0, 0, 0, 0]) - score = field(converter=iterable_to_xywh_rect, default=[0, 0, 0, 0]) - max_recall_rating_class = field( - converter=iterable_to_xywh_rect, default=[0, 0, 0, 0] - ) - title = field(converter=iterable_to_xywh_rect, default=[0, 0, 0, 0]) diff --git a/src/arcaea_offline_ocr/device/v2/ocr.py b/src/arcaea_offline_ocr/device/v2/ocr.py deleted file mode 100644 index bff9c6b..0000000 --- a/src/arcaea_offline_ocr/device/v2/ocr.py +++ /dev/null @@ -1,172 +0,0 @@ -import math -from functools import lru_cache -from typing import Sequence - -import cv2 -import numpy as np -from PIL import Image - -from ...crop import crop_xywh -from ...mask import ( - mask_byd, - mask_ftr, - mask_gray, - mask_max_recall_purple, - mask_pfl_white, - mask_prs, - mask_pst, - mask_white, -) -from ...ocr import ( - FixRects, - ocr_digit_samples_knn, - ocr_digits_by_contour_knn, - preprocess_hog, - resize_fill_square, -) -from ...phash_db import ImagePHashDatabase -from ...sift_db import SIFTDatabase -from ...types import Mat, cv2_ml_KNearest -from ..shared import DeviceOcrResult -from .preprocess import find_digits_preprocess -from .rois import DeviceV2Rois -from .shared import MAX_RECALL_CLOSE_KERNEL -from .sizes import SizesV2 - - -class DeviceV2Ocr: - def __init__(self, knn_model: cv2_ml_KNearest, phash_db: ImagePHashDatabase): - self.__knn_model = knn_model - self.__phash_db = phash_db - - @property - def knn_model(self): - if not self.__knn_model: - raise ValueError("`knn_model` unset.") - return self.__knn_model - - @knn_model.setter - def knn_model(self, value: cv2_ml_KNearest): - self.__knn_model = value - - @property - def phash_db(self): - if not self.__phash_db: - raise ValueError("`phash_db` unset.") - return self.__phash_db - - @phash_db.setter - def phash_db(self, value: SIFTDatabase): - self.__phash_db = value - - @lru_cache - def _get_digit_widths(self, num_list: Sequence[int], factor: float): - widths = set() - for n in num_list: - lower = math.floor(n * factor) - upper = math.ceil(n * factor) - widths.update(range(lower, upper + 1)) - return widths - - def _base_ocr_pfl(self, roi_masked: Mat, factor: float = 1.0): - contours, _ = cv2.findContours( - roi_masked, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE - ) - filtered_contours = [c for c in contours if cv2.contourArea(c) >= 5 * factor] - rects = [cv2.boundingRect(c) for c in filtered_contours] - rects = FixRects.connect_broken(rects, roi_masked.shape[1], roi_masked.shape[0]) - rect_contour_map = dict(zip(rects, filtered_contours)) - - filtered_rects = [r for r in rects if r[2] >= 5 * factor and r[3] >= 6 * factor] - filtered_rects = FixRects.split_connected(roi_masked, filtered_rects) - filtered_rects = sorted(filtered_rects, key=lambda r: r[0]) - - roi_ocr = roi_masked.copy() - filtered_contours_flattened = {tuple(c.flatten()) for c in filtered_contours} - for contour in contours: - if tuple(contour.flatten()) in filtered_contours_flattened: - continue - roi_ocr = cv2.fillPoly(roi_ocr, [contour], [0]) - digit_rois = [ - resize_fill_square(crop_xywh(roi_ocr, r), 20) - for r in sorted(filtered_rects, key=lambda r: r[0]) - ] - # [cv2.imshow(f"r{i}", r) for i, r in enumerate(digit_rois)] - # cv2.waitKey(0) - samples = preprocess_hog(digit_rois) - return ocr_digit_samples_knn(samples, self.knn_model) - - def ocr_song_id(self, rois: DeviceV2Rois): - jacket = cv2.cvtColor(rois.jacket, cv2.COLOR_BGR2GRAY) - return self.phash_db.lookup_image(Image.fromarray(jacket))[0] - - def ocr_rating_class(self, rois: DeviceV2Rois): - roi = cv2.cvtColor(rois.max_recall_rating_class, cv2.COLOR_BGR2HSV) - results = [mask_pst(roi), mask_prs(roi), mask_ftr(roi), mask_byd(roi)] - return max(enumerate(results), key=lambda i: np.count_nonzero(i[1]))[0] - - def ocr_score(self, rois: DeviceV2Rois): - roi = cv2.cvtColor(rois.score, cv2.COLOR_BGR2HSV) - roi = mask_white(roi) - contours, _ = cv2.findContours(roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) - for contour in contours: - x, y, w, h = cv2.boundingRect(contour) - if h < roi.shape[0] * 0.6: - roi = cv2.fillPoly(roi, [contour], [0]) - return ocr_digits_by_contour_knn(roi, self.knn_model) - - def mask_pfl(self, pfl_roi: Mat, rois: DeviceV2Rois): - return ( - mask_pfl_white(cv2.cvtColor(pfl_roi, cv2.COLOR_BGR2HSV)) - if isinstance(rois.sizes, SizesV2) - else mask_gray(pfl_roi) - ) - - def ocr_pure(self, rois: DeviceV2Rois): - roi = self.mask_pfl(rois.pure, rois) - return self._base_ocr_pfl(roi, rois.sizes.factor) - - def ocr_far(self, rois: DeviceV2Rois): - roi = self.mask_pfl(rois.far, rois) - return self._base_ocr_pfl(roi, rois.sizes.factor) - - def ocr_lost(self, rois: DeviceV2Rois): - roi = self.mask_pfl(rois.lost, rois) - return self._base_ocr_pfl(roi, rois.sizes.factor) - - def ocr_max_recall(self, rois: DeviceV2Rois): - roi = ( - mask_max_recall_purple( - cv2.cvtColor(rois.max_recall_rating_class, cv2.COLOR_BGR2HSV) - ) - if isinstance(rois.sizes, SizesV2) - else mask_gray(rois.max_recall_rating_class) - ) - roi_closed = cv2.morphologyEx(roi, cv2.MORPH_CLOSE, MAX_RECALL_CLOSE_KERNEL) - contours, _ = cv2.findContours( - roi_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE - ) - rects = sorted( - [cv2.boundingRect(c) for c in contours], key=lambda r: r[0], reverse=True - ) - max_recall_roi = crop_xywh(roi, rects[0]) - return ocr_digits_by_contour_knn(max_recall_roi, self.knn_model) - - def ocr(self, rois: DeviceV2Rois): - song_id = self.ocr_song_id(rois) - rating_class = self.ocr_rating_class(rois) - score = self.ocr_score(rois) - pure = self.ocr_pure(rois) - far = self.ocr_far(rois) - lost = self.ocr_lost(rois) - max_recall = self.ocr_max_recall(rois) - - return DeviceOcrResult( - rating_class=rating_class, - pure=pure, - far=far, - lost=lost, - score=score, - max_recall=max_recall, - song_id=song_id, - ) diff --git a/src/arcaea_offline_ocr/device/v2/preprocess.py b/src/arcaea_offline_ocr/device/v2/preprocess.py deleted file mode 100644 index deef8d5..0000000 --- a/src/arcaea_offline_ocr/device/v2/preprocess.py +++ /dev/null @@ -1,54 +0,0 @@ -import cv2 - -from ...types import Mat -from .shared import * - - -def find_digits_preprocess(__img_masked: Mat) -> Mat: - img = __img_masked.copy() - img_denoised = cv2.morphologyEx(img, cv2.MORPH_OPEN, PFL_DENOISE_KERNEL) - # img_denoised = cv2.bitwise_and(img, img_denoised) - - denoise_contours, _ = cv2.findContours( - img_denoised, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE - ) - # cv2.drawContours(img_denoised, contours, -1, [128], 2) - - # fill all contour.area < max(contour.area) * ratio with black pixels - # for denoise purposes - - # define threshold contour area - # we assume the smallest digit "1", is 80% height of the image, - # and at least 1.5 pixel wide, considering cv2.contourArea always - # returns a smaller value than the actual contour area. - max_contour_area = __img_masked.shape[0] * 0.8 * 1.5 - filtered_contours = list( - filter(lambda c: cv2.contourArea(c) >= max_contour_area, denoise_contours) - ) - - filtered_contours_flattened = {tuple(c.flatten()) for c in filtered_contours} - - for contour in denoise_contours: - if tuple(contour.flatten()) not in filtered_contours_flattened: - img_denoised = cv2.fillPoly(img_denoised, [contour], [0]) - - # old algorithm, finding the largest contour area - ## contour_area_tuples = [(contour, cv2.contourArea(contour)) for contour in contours] - ## contour_area_tuples = sorted( - ## contour_area_tuples, key=lambda item: item[1], reverse=True - ## ) - ## max_contour_area = contour_area_tuples[0][1] - ## print(max_contour_area, [item[1] for item in contour_area_tuples]) - ## contours_filter_end_index = len(contours) - ## for i, item in enumerate(contour_area_tuples): - ## contour, area = item - ## if area < max_contour_area * 0.15: - ## contours_filter_end_index = i - ## break - ## contours = [item[0] for item in contour_area_tuples] - ## for contour in contours[-contours_filter_end_index - 1:]: - ## img = cv2.fillPoly(img, [contour], [0]) - ## img_denoised = cv2.fillPoly(img_denoised, [contour], [0]) - ## contours = contours[:contours_filter_end_index] - - return img_denoised diff --git a/src/arcaea_offline_ocr/device/v2/rois.py b/src/arcaea_offline_ocr/device/v2/rois.py deleted file mode 100644 index 100aece..0000000 --- a/src/arcaea_offline_ocr/device/v2/rois.py +++ /dev/null @@ -1,199 +0,0 @@ -from typing import Union - -from ...crop import crop_black_edges, crop_xywh -from ...types import Mat, XYWHRect -from .definition import DeviceV2 -from .sizes import Sizes, SizesV1 - - -def to_int(num: Union[int, float]) -> int: - return round(num) - - -class DeviceV2Rois: - def __init__(self, device: DeviceV2, img: Mat): - self.device = device - self.sizes = SizesV1(self.device.factor) - self.__img = img - - @staticmethod - def construct_int_xywh_rect(x, y, w, h) -> XYWHRect: - return XYWHRect(*[to_int(item) for item in [x, y, w, h]]) - - @property - def img(self): - return self.__img - - @img.setter - def img(self, img: Mat): - self.__img = ( - crop_black_edges(img) if self.device.crop_black_edges else img.copy() - ) - - @property - def h(self): - return self.img.shape[0] - - @property - def vmid(self): - return self.h / 2 - - @property - def w(self): - return self.img.shape[1] - - @property - def hmid(self): - return self.w / 2 - - @property - def h_without_top_bar(self): - """img_height -= top_bar_height""" - return self.h - self.sizes.TOP_BAR_HEIGHT - - @property - def h_without_top_bar_mid(self): - return self.sizes.TOP_BAR_HEIGHT + self.h_without_top_bar / 2 - - @property - def pfl_top(self): - return self.h_without_top_bar_mid + self.sizes.PFL_TOP_FROM_VMID - - @property - def pfl_left(self): - return self.hmid + self.sizes.PFL_LEFT_FROM_HMID - - @property - def pure_rect(self): - return self.construct_int_xywh_rect( - x=self.pfl_left, - y=self.pfl_top, - w=self.sizes.PFL_WIDTH, - h=self.sizes.PFL_FONT_PX, - ) - - @property - def pure(self): - return crop_xywh(self.img, self.pure_rect) - - @property - def far_rect(self): - return self.construct_int_xywh_rect( - x=self.pfl_left, - y=self.pfl_top + self.sizes.PFL_FONT_PX + self.sizes.PURE_FAR_GAP, - w=self.sizes.PFL_WIDTH, - h=self.sizes.PFL_FONT_PX, - ) - - @property - def far(self): - return crop_xywh(self.img, self.far_rect) - - @property - def lost_rect(self): - return self.construct_int_xywh_rect( - x=self.pfl_left, - y=( - self.pfl_top - + self.sizes.PFL_FONT_PX * 2 - + self.sizes.PURE_FAR_GAP - + self.sizes.FAR_LOST_GAP - ), - w=self.sizes.PFL_WIDTH, - h=self.sizes.PFL_FONT_PX, - ) - - @property - def lost(self): - return crop_xywh(self.img, self.lost_rect) - - @property - def score_rect(self): - return self.construct_int_xywh_rect( - x=self.hmid - (self.sizes.SCORE_WIDTH / 2), - y=( - self.h_without_top_bar_mid - + self.sizes.SCORE_BOTTOM_FROM_VMID - - self.sizes.SCORE_FONT_PX - ), - w=self.sizes.SCORE_WIDTH, - h=self.sizes.SCORE_FONT_PX, - ) - - @property - def score(self): - return crop_xywh(self.img, self.score_rect) - - @property - def max_recall_rating_class_rect(self): - x = ( - self.hmid - + self.sizes.JACKET_RIGHT_FROM_HOR_MID - - self.sizes.JACKET_WIDTH - - 25 * self.sizes.factor - ) - return self.construct_int_xywh_rect( - x=x, - y=( - self.h_without_top_bar_mid - - self.sizes.SCORE_PANEL[1] / 2 - - self.sizes.MR_RT_HEIGHT - ), - w=self.sizes.MR_RT_WIDTH, - h=self.sizes.MR_RT_HEIGHT, - ) - - @property - def max_recall_rating_class(self): - return crop_xywh(self.img, self.max_recall_rating_class_rect) - - @property - def title_rect(self): - return self.construct_int_xywh_rect( - x=0, - y=self.h_without_top_bar_mid - + self.sizes.TITLE_BOTTOM_FROM_VMID - - self.sizes.TITLE_FONT_PX, - w=self.hmid + self.sizes.TITLE_WIDTH_RIGHT, - h=self.sizes.TITLE_FONT_PX, - ) - - @property - def title(self): - return crop_xywh(self.img, self.title_rect) - - @property - def jacket_rect(self): - return self.construct_int_xywh_rect( - x=self.hmid - + self.sizes.JACKET_RIGHT_FROM_HOR_MID - - self.sizes.JACKET_WIDTH, - y=self.h_without_top_bar_mid - self.sizes.SCORE_PANEL[1] / 2, - w=self.sizes.JACKET_WIDTH, - h=self.sizes.JACKET_WIDTH, - ) - - @property - def jacket(self): - return crop_xywh(self.img, self.jacket_rect) - - -class DeviceV2AutoRois(DeviceV2Rois): - @staticmethod - def get_factor(width: int, height: int): - ratio = width / height - return ((width / 16) * 9) / 720 if ratio < (16 / 9) else height / 720 - - def __init__(self, img: Mat): - factor = self.get_factor(img.shape[1], img.shape[0]) - self.sizes = SizesV1(factor) - self.__img = None - self.img = img - - @property - def img(self): - return self.__img - - @img.setter - def img(self, img: Mat): - self.__img = crop_black_edges(img) diff --git a/src/arcaea_offline_ocr/device/v2/shared.py b/src/arcaea_offline_ocr/device/v2/shared.py deleted file mode 100644 index ca511b1..0000000 --- a/src/arcaea_offline_ocr/device/v2/shared.py +++ /dev/null @@ -1,9 +0,0 @@ -from cv2 import MORPH_RECT, getStructuringElement - -PFL_DENOISE_KERNEL = getStructuringElement(MORPH_RECT, [2, 2]) -PFL_ERODE_KERNEL = getStructuringElement(MORPH_RECT, [3, 3]) -PFL_CLOSE_HORIZONTAL_KERNEL = getStructuringElement(MORPH_RECT, [10, 1]) - -MAX_RECALL_DENOISE_KERNEL = getStructuringElement(MORPH_RECT, [3, 3]) -MAX_RECALL_ERODE_KERNEL = getStructuringElement(MORPH_RECT, [2, 2]) -MAX_RECALL_CLOSE_KERNEL = getStructuringElement(MORPH_RECT, [20, 1]) diff --git a/src/arcaea_offline_ocr/device/v2/sizes.py b/src/arcaea_offline_ocr/device/v2/sizes.py deleted file mode 100644 index 3347cb2..0000000 --- a/src/arcaea_offline_ocr/device/v2/sizes.py +++ /dev/null @@ -1,254 +0,0 @@ -from typing import Tuple, Union - - -def apply_factor(num: Union[int, float], factor: float): - return num * factor - - -class Sizes: - def __init__(self, factor: float): - raise NotImplementedError() - - @property - def TOP_BAR_HEIGHT(self): - ... - - @property - def SCORE_PANEL(self) -> Tuple[int, int]: - ... - - @property - def PFL_TOP_FROM_VMID(self): - ... - - @property - def PFL_LEFT_FROM_HMID(self): - ... - - @property - def PFL_WIDTH(self): - ... - - @property - def PFL_FONT_PX(self): - ... - - @property - def PURE_FAR_GAP(self): - ... - - @property - def FAR_LOST_GAP(self): - ... - - @property - def SCORE_BOTTOM_FROM_VMID(self): - ... - - @property - def SCORE_FONT_PX(self): - ... - - @property - def SCORE_WIDTH(self): - ... - - @property - def JACKET_RIGHT_FROM_HOR_MID(self): - ... - - @property - def JACKET_WIDTH(self): - ... - - @property - def MR_RT_RIGHT_FROM_HMID(self): - ... - - @property - def MR_RT_WIDTH(self): - ... - - @property - def MR_RT_HEIGHT(self): - ... - - @property - def TITLE_BOTTOM_FROM_VMID(self): - ... - - @property - def TITLE_FONT_PX(self): - ... - - @property - def TITLE_WIDTH_RIGHT(self): - ... - - -class SizesV1(Sizes): - def __init__(self, factor: float): - self.factor = factor - - def apply_factor(self, num): - return apply_factor(num, self.factor) - - @property - def TOP_BAR_HEIGHT(self): - return self.apply_factor(50) - - @property - def SCORE_PANEL(self) -> Tuple[int, int]: - return tuple(self.apply_factor(num) for num in [485, 239]) - - @property - def PFL_TOP_FROM_VMID(self): - return self.apply_factor(135) - - @property - def PFL_LEFT_FROM_HMID(self): - return self.apply_factor(5) - - @property - def PFL_WIDTH(self): - return self.apply_factor(76) - - @property - def PFL_FONT_PX(self): - return self.apply_factor(26) - - @property - def PURE_FAR_GAP(self): - return self.apply_factor(12) - - @property - def FAR_LOST_GAP(self): - return self.apply_factor(10) - - @property - def SCORE_BOTTOM_FROM_VMID(self): - return self.apply_factor(-50) - - @property - def SCORE_FONT_PX(self): - return self.apply_factor(45) - - @property - def SCORE_WIDTH(self): - return self.apply_factor(280) - - @property - def JACKET_RIGHT_FROM_HOR_MID(self): - return self.apply_factor(-235) - - @property - def JACKET_WIDTH(self): - return self.apply_factor(375) - - @property - def MR_RT_RIGHT_FROM_HMID(self): - return self.apply_factor(-300) - - @property - def MR_RT_WIDTH(self): - return self.apply_factor(275) - - @property - def MR_RT_HEIGHT(self): - return self.apply_factor(75) - - @property - def TITLE_BOTTOM_FROM_VMID(self): - return self.apply_factor(-265) - - @property - def TITLE_FONT_PX(self): - return self.apply_factor(40) - - @property - def TITLE_WIDTH_RIGHT(self): - return self.apply_factor(275) - - -class SizesV2(Sizes): - def __init__(self, factor: float): - self.factor = factor - - def apply_factor(self, num): - return apply_factor(num, self.factor) - - @property - def TOP_BAR_HEIGHT(self): - return self.apply_factor(50) - - @property - def SCORE_PANEL(self) -> Tuple[int, int]: - return tuple(self.apply_factor(num) for num in [447, 233]) - - @property - def PFL_TOP_FROM_VMID(self): - return self.apply_factor(142) - - @property - def PFL_LEFT_FROM_HMID(self): - return self.apply_factor(10) - - @property - def PFL_WIDTH(self): - return self.apply_factor(60) - - @property - def PFL_FONT_PX(self): - return self.apply_factor(16) - - @property - def PURE_FAR_GAP(self): - return self.apply_factor(20) - - @property - def FAR_LOST_GAP(self): - return self.apply_factor(23) - - @property - def SCORE_BOTTOM_FROM_VMID(self): - return self.apply_factor(-50) - - @property - def SCORE_FONT_PX(self): - return self.apply_factor(45) - - @property - def SCORE_WIDTH(self): - return self.apply_factor(280) - - @property - def JACKET_RIGHT_FROM_HOR_MID(self): - return self.apply_factor(-235) - - @property - def JACKET_WIDTH(self): - return self.apply_factor(375) - - @property - def MR_RT_RIGHT_FROM_HMID(self): - return self.apply_factor(-330) - - @property - def MR_RT_WIDTH(self): - return self.apply_factor(330) - - @property - def MR_RT_HEIGHT(self): - return self.apply_factor(75) - - @property - def TITLE_BOTTOM_FROM_VMID(self): - return self.apply_factor(-265) - - @property - def TITLE_FONT_PX(self): - return self.apply_factor(40) - - @property - def TITLE_WIDTH_RIGHT(self): - return self.apply_factor(275)