mirror of
https://github.com/283375/arcaea-offline-ocr.git
synced 2025-07-01 04:16:27 +00:00
refactor!: chieri v4 b30 scenario
- Remove useless `.utils` code
This commit is contained in:
@ -1,15 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class B30OcrResultItem:
|
||||
rating_class: int
|
||||
score: int
|
||||
pure: Optional[int] = None
|
||||
far: Optional[int] = None
|
||||
lost: Optional[int] = None
|
||||
date: Optional[datetime] = None
|
||||
title: Optional[str] = None
|
||||
song_id: Optional[str] = None
|
3
src/arcaea_offline_ocr/scenarios/b30/__init__.py
Normal file
3
src/arcaea_offline_ocr/scenarios/b30/__init__.py
Normal file
@ -0,0 +1,3 @@
|
||||
from .chieri import ChieriBotV4Best30Scenario
|
||||
|
||||
__all__ = ["ChieriBotV4Best30Scenario"]
|
3
src/arcaea_offline_ocr/scenarios/b30/chieri/__init__.py
Normal file
3
src/arcaea_offline_ocr/scenarios/b30/chieri/__init__.py
Normal file
@ -0,0 +1,3 @@
|
||||
from .v4 import ChieriBotV4Best30Scenario
|
||||
|
||||
__all__ = ["ChieriBotV4Best30Scenario"]
|
@ -0,0 +1,3 @@
|
||||
from .impl import ChieriBotV4Best30Scenario
|
||||
|
||||
__all__ = ["ChieriBotV4Best30Scenario"]
|
@ -27,11 +27,11 @@ FAR_BG_MAX_HSV = np.array([20, 255, 255], np.uint8)
|
||||
LOST_BG_MIN_HSV = np.array([115, 60, 150], np.uint8)
|
||||
LOST_BG_MAX_HSV = np.array([140, 255, 255], np.uint8)
|
||||
|
||||
BYD_MIN_HSV = (158, 120, 0)
|
||||
BYD_MAX_HSV = (172, 255, 255)
|
||||
BYD_MIN_HSV = np.array([158, 120, 0], np.uint8)
|
||||
BYD_MAX_HSV = np.array([172, 255, 255], np.uint8)
|
||||
|
||||
FTR_MIN_HSV = (145, 70, 0)
|
||||
FTR_MAX_HSV = (160, 255, 255)
|
||||
FTR_MIN_HSV = np.array([145, 70, 0], np.uint8)
|
||||
FTR_MAX_HSV = np.array([160, 255, 255], np.uint8)
|
||||
|
||||
PRS_MIN_HSV = (45, 60, 0)
|
||||
PRS_MAX_HSV = (70, 255, 255)
|
||||
PRS_MIN_HSV = np.array([45, 60, 0], np.uint8)
|
||||
PRS_MAX_HSV = np.array([70, 255, 255], np.uint8)
|
@ -3,10 +3,16 @@ from typing import List, Optional, Tuple
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ....crop import crop_xywh
|
||||
from ....phash_db import ImagePhashDatabase
|
||||
from ....types import Mat
|
||||
from ...shared import B30OcrResultItem
|
||||
from arcaea_offline_ocr.crop import crop_xywh
|
||||
from arcaea_offline_ocr.providers import (
|
||||
ImageCategory,
|
||||
ImageIdProvider,
|
||||
OcrKNearestTextProvider,
|
||||
)
|
||||
from arcaea_offline_ocr.scenarios.b30.base import Best30Scenario
|
||||
from arcaea_offline_ocr.scenarios.base import OcrScenarioResult
|
||||
from arcaea_offline_ocr.types import Mat
|
||||
|
||||
from .colors import (
|
||||
BYD_MAX_HSV,
|
||||
BYD_MIN_HSV,
|
||||
@ -22,29 +28,20 @@ from .colors import (
|
||||
PURE_BG_MIN_HSV,
|
||||
)
|
||||
from .rois import ChieriBotV4Rois
|
||||
from ....providers.knn import OcrKNearestTextProvider
|
||||
|
||||
|
||||
class ChieriBotV4Ocr:
|
||||
class ChieriBotV4Best30Scenario(Best30Scenario):
|
||||
def __init__(
|
||||
self,
|
||||
score_knn_provider: OcrKNearestTextProvider,
|
||||
pfl_knn_provider: OcrKNearestTextProvider,
|
||||
phash_db: ImagePhashDatabase,
|
||||
image_id_provider: ImageIdProvider,
|
||||
factor: float = 1.0,
|
||||
):
|
||||
self.__phash_db = phash_db
|
||||
self.__rois = ChieriBotV4Rois(factor)
|
||||
self.pfl_knn_provider = pfl_knn_provider
|
||||
self.score_knn_provider = score_knn_provider
|
||||
|
||||
@property
|
||||
def phash_db(self):
|
||||
return self.__phash_db
|
||||
|
||||
@phash_db.setter
|
||||
def phash_db(self, phash_db: ImagePhashDatabase):
|
||||
self.__phash_db = phash_db
|
||||
self.image_id_provider = image_id_provider
|
||||
|
||||
@property
|
||||
def rois(self):
|
||||
@ -77,12 +74,12 @@ class ChieriBotV4Ocr:
|
||||
else:
|
||||
return max(enumerate(rating_class_results), key=lambda i: i[1])[0] + 1
|
||||
|
||||
def ocr_component_song_id(self, component_bgr: Mat):
|
||||
def ocr_component_song_id_results(self, component_bgr: Mat):
|
||||
jacket_rect = self.rois.component_rois.jacket_rect.floored()
|
||||
jacket_roi = cv2.cvtColor(
|
||||
crop_xywh(component_bgr, jacket_rect), cv2.COLOR_BGR2GRAY
|
||||
)
|
||||
return self.phash_db.lookup_jacket(jacket_roi)[0]
|
||||
return self.image_id_provider.results(jacket_roi, ImageCategory.JACKET)
|
||||
|
||||
def ocr_component_score_knn(self, component_bgr: Mat) -> int:
|
||||
# sourcery skip: inline-immediately-returned-variable
|
||||
@ -191,28 +188,36 @@ class ChieriBotV4Ocr:
|
||||
except Exception:
|
||||
return (None, None, None)
|
||||
|
||||
def ocr_component(self, component_bgr: Mat) -> B30OcrResultItem:
|
||||
def ocr_component(self, component_bgr: Mat) -> OcrScenarioResult:
|
||||
component_blur = cv2.GaussianBlur(component_bgr, (5, 5), 0)
|
||||
rating_class = self.ocr_component_rating_class(component_blur)
|
||||
song_id = self.ocr_component_song_id(component_bgr)
|
||||
# title = self.ocr_component_title(component_blur)
|
||||
song_id_results = self.ocr_component_song_id_results(component_bgr)
|
||||
# score = self.ocr_component_score(component_blur)
|
||||
score = self.ocr_component_score_knn(component_bgr)
|
||||
pure, far, lost = self.ocr_component_pfl(component_bgr)
|
||||
return B30OcrResultItem(
|
||||
song_id=song_id,
|
||||
return OcrScenarioResult(
|
||||
song_id=song_id_results[0].image_id,
|
||||
song_id_results=song_id_results,
|
||||
rating_class=rating_class,
|
||||
# title=title,
|
||||
score=score,
|
||||
pure=pure,
|
||||
far=far,
|
||||
lost=lost,
|
||||
date=None,
|
||||
played_at=None,
|
||||
)
|
||||
|
||||
def ocr(self, img_bgr: Mat) -> List[B30OcrResultItem]:
|
||||
self.set_factor(img_bgr)
|
||||
return [
|
||||
self.ocr_component(component_bgr)
|
||||
for component_bgr in self.rois.components(img_bgr)
|
||||
]
|
||||
def components(self, img: Mat, /):
|
||||
"""
|
||||
:param img: BGR format image
|
||||
"""
|
||||
self.set_factor(img)
|
||||
return self.rois.components(img)
|
||||
|
||||
def result(self, component_img: Mat, /):
|
||||
return self.ocr_component(component_img)
|
||||
|
||||
def results(self, img: Mat, /) -> List[OcrScenarioResult]:
|
||||
"""
|
||||
:param img: BGR format image
|
||||
"""
|
||||
return [self.ocr_component(component) for component in self.components(img)]
|
@ -1,8 +1,7 @@
|
||||
from typing import List
|
||||
|
||||
from ....crop import crop_xywh
|
||||
from ....types import Mat, XYWHRect
|
||||
from ....utils import apply_factor
|
||||
from arcaea_offline_ocr.crop import crop_xywh
|
||||
from arcaea_offline_ocr.types import Mat, XYWHRect
|
||||
|
||||
|
||||
class ChieriBotV4ComponentRois:
|
||||
@ -19,39 +18,39 @@ class ChieriBotV4ComponentRois:
|
||||
|
||||
@property
|
||||
def top_font_color_detect(self):
|
||||
return apply_factor(XYWHRect(35, 10, 120, 100), self.factor)
|
||||
return XYWHRect(35, 10, 120, 100), self.factor
|
||||
|
||||
@property
|
||||
def bottom_font_color_detect(self):
|
||||
return apply_factor(XYWHRect(30, 125, 175, 110), self.factor)
|
||||
return XYWHRect(30, 125, 175, 110) * self.factor
|
||||
|
||||
@property
|
||||
def bg_point(self):
|
||||
return apply_factor((75, 10), self.factor)
|
||||
return (75 * self.factor, 10 * self.factor)
|
||||
|
||||
@property
|
||||
def rating_class_rect(self):
|
||||
return apply_factor(XYWHRect(21, 40, 7, 20), self.factor)
|
||||
return XYWHRect(21, 40, 7, 20) * self.factor
|
||||
|
||||
@property
|
||||
def title_rect(self):
|
||||
return apply_factor(XYWHRect(35, 10, 430, 50), self.factor)
|
||||
return XYWHRect(35, 10, 430, 50) * self.factor
|
||||
|
||||
@property
|
||||
def jacket_rect(self):
|
||||
return apply_factor(XYWHRect(263, 0, 239, 239), self.factor)
|
||||
return XYWHRect(263, 0, 239, 239) * self.factor
|
||||
|
||||
@property
|
||||
def score_rect(self):
|
||||
return apply_factor(XYWHRect(30, 60, 270, 55), self.factor)
|
||||
return XYWHRect(30, 60, 270, 55) * self.factor
|
||||
|
||||
@property
|
||||
def pfl_rect(self):
|
||||
return apply_factor(XYWHRect(50, 125, 80, 100), self.factor)
|
||||
return XYWHRect(50, 125, 80, 100) * self.factor
|
||||
|
||||
@property
|
||||
def date_rect(self):
|
||||
return apply_factor(XYWHRect(205, 200, 225, 25), self.factor)
|
||||
return XYWHRect(205, 200, 225, 25) * self.factor
|
||||
|
||||
|
||||
class ChieriBotV4Rois:
|
||||
@ -74,27 +73,27 @@ class ChieriBotV4Rois:
|
||||
|
||||
@property
|
||||
def top(self):
|
||||
return apply_factor(823, self.factor)
|
||||
return 823 * self.factor
|
||||
|
||||
@property
|
||||
def left(self):
|
||||
return apply_factor(107, self.factor)
|
||||
return 107 * self.factor
|
||||
|
||||
@property
|
||||
def width(self):
|
||||
return apply_factor(502, self.factor)
|
||||
return 502 * self.factor
|
||||
|
||||
@property
|
||||
def height(self):
|
||||
return apply_factor(240, self.factor)
|
||||
return 240 * self.factor
|
||||
|
||||
@property
|
||||
def vertical_gap(self):
|
||||
return apply_factor(74, self.factor)
|
||||
return 74 * self.factor
|
||||
|
||||
@property
|
||||
def horizontal_gap(self):
|
||||
return apply_factor(40, self.factor)
|
||||
return 40 * self.factor
|
||||
|
||||
@property
|
||||
def horizontal_items(self):
|
||||
@ -104,7 +103,7 @@ class ChieriBotV4Rois:
|
||||
|
||||
@property
|
||||
def b33_vertical_gap(self):
|
||||
return apply_factor(121, self.factor)
|
||||
return 121 * self.factor
|
||||
|
||||
def components(self, img_bgr: Mat) -> List[Mat]:
|
||||
first_rect = XYWHRect(x=self.left, y=self.top, w=self.width, h=self.height)
|
@ -1,11 +1,6 @@
|
||||
from collections.abc import Iterable
|
||||
from typing import TypeVar, overload
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from .types import XYWHRect
|
||||
|
||||
__all__ = ["imread_unicode"]
|
||||
|
||||
|
||||
@ -13,27 +8,3 @@ def imread_unicode(filepath: str, flags: int = cv2.IMREAD_UNCHANGED):
|
||||
# https://stackoverflow.com/a/57872297/16484891
|
||||
# CC BY-SA 4.0
|
||||
return cv2.imdecode(np.fromfile(filepath, dtype=np.uint8), flags)
|
||||
|
||||
|
||||
@overload
|
||||
def apply_factor(item: int, factor: float) -> float: ...
|
||||
|
||||
|
||||
@overload
|
||||
def apply_factor(item: float, factor: float) -> float: ...
|
||||
|
||||
|
||||
T = TypeVar("T", bound=Iterable)
|
||||
|
||||
|
||||
@overload
|
||||
def apply_factor(item: T, factor: float) -> T: ...
|
||||
|
||||
|
||||
def apply_factor(item, factor: float):
|
||||
if isinstance(item, (int, float)):
|
||||
return item * factor
|
||||
if isinstance(item, XYWHRect):
|
||||
return item.__class__(*[i * factor for i in item])
|
||||
if isinstance(item, Iterable):
|
||||
return item.__class__([i * factor for i in item])
|
||||
|
Reference in New Issue
Block a user