20 Commits

Author SHA1 Message Date
0b53682398 ci: fix tag condition
- ok i didnt expect github actions does not support regex wtf
2025-06-26 00:15:02 +08:00
b117346b46 chore: setuptools-scm integration 2025-06-26 00:07:08 +08:00
ad0a33daad ci: tag regex 2025-06-25 23:43:46 +08:00
c08a1332a7 chore!: removing unused code 2025-06-25 23:37:21 +08:00
0055d9e8da refactor!: device scenario
- Correct abstract class annotations
2025-06-25 23:35:38 +08:00
06156db9c2 refactor!: chieri v4 b30 scenario
- Remove useless `.utils` code
2025-06-25 23:27:15 +08:00
c65798a02d feat: XYWHRect __mul__ 2025-06-25 23:19:52 +08:00
f11dc6e38f refactor: scenario base 2025-06-25 23:11:45 +08:00
2b18906935 refactor!: image hash database provider 2025-06-22 01:28:59 +08:00
abfd37dbef refactor!: OCR text result provider 2025-06-22 00:32:31 +08:00
3ebb058cdf refactor: XYWHRect and b30 ocr 2025-06-21 16:06:49 +08:00
b545c5b6bf refactor!: ImageHashType -> ImageHashCategory 2025-06-17 22:28:16 +08:00
212afa32db chore: update dependencies 2025-06-17 18:13:59 +08:00
2264e90b8e refactor!: replace attrs with dataclass 2025-06-17 18:05:44 +08:00
619bff2ea4 feat: image hashes database 2025-01-10 23:55:37 +08:00
413188d86a feat: core hashers 2025-01-10 23:54:37 +08:00
cfe8de043c chore: pre-commit hooks update 2025-01-10 23:54:00 +08:00
3f6c08b2ad fix: preprocess_char_icon copyMakeBorder edge case 2025-01-10 23:46:50 +08:00
854d5558cf chore: v0.0.99 2024-06-19 22:23:30 +08:00
df77421a34 impr: DeviceRoisMaskerAutoT2 pfl color range (#11) 2024-06-05 18:46:16 +08:00
50 changed files with 974 additions and 603 deletions

View File

@ -4,7 +4,7 @@ on:
workflow_dispatch:
push:
tags:
- "v[0-9]+.[0-9]+.[0-9]+"
- '*.*.*'
permissions:
contents: write
@ -29,14 +29,6 @@ jobs:
pip install build
python -m build
- name: Remove `v` in tag name
uses: mad9000/actions-find-and-replace-string@5
id: tagNameReplaced
with:
source: ${{ github.ref_name }}
find: "v"
replace: ""
- name: Draft a release
uses: softprops/action-gh-release@v2
with:
@ -44,5 +36,4 @@ jobs:
draft: true
generate_release_notes: true
files: |
dist/arcaea_offline_ocr-${{ steps.tagNameReplaced.outputs.value }}*.whl
dist/arcaea-offline-ocr-${{ steps.tagNameReplaced.outputs.value }}.tar.gz
dist/*

View File

@ -4,11 +4,10 @@ repos:
hooks:
- id: end-of-file-fixer
- id: trailing-whitespace
- repo: https://github.com/psf/black
rev: 23.1.0
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.11.13
hooks:
- id: black
- repo: https://github.com/PyCQA/isort
rev: 5.12.0
hooks:
- id: isort
- id: ruff
args: ["--fix"]
- id: ruff-format

View File

@ -1,10 +1,11 @@
[build-system]
requires = ["setuptools>=61.0"]
requires = ["setuptools>=64", "setuptools-scm>=8"]
build-backend = "setuptools.build_meta"
[project]
dynamic = ["version"]
name = "arcaea-offline-ocr"
version = "0.0.98"
authors = [{ name = "283375", email = "log_283375@163.com" }]
description = "Extract your Arcaea play result from screenshot."
readme = "README.md"
@ -19,6 +20,8 @@ classifiers = [
"Homepage" = "https://github.com/ArcaeaOffline/core-ocr"
"Bug Tracker" = "https://github.com/ArcaeaOffline/core-ocr/issues"
[tool.setuptools_scm]
[tool.isort]
profile = "black"
src_paths = ["src/arcaea_offline_ocr"]
@ -34,5 +37,5 @@ generated-members = ["cv2.*"]
disable = [
"missing-module-docstring",
"missing-class-docstring",
"missing-function-docstring"
"missing-function-docstring",
]

View File

@ -1,3 +1,2 @@
attrs==23.1.0
numpy==1.26.1
opencv-python==4.8.1.78
numpy~=2.3
opencv-python~=4.11

View File

@ -1,4 +0,0 @@
from .crop import *
from .device import *
from .ocr import *
from .utils import *

View File

@ -1,16 +0,0 @@
from datetime import datetime
from typing import Optional
import attrs
@attrs.define
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

View File

@ -0,0 +1,6 @@
from .ihdb import ImageHashDatabaseBuildTask, ImageHashesDatabaseBuilder
__all__ = [
"ImageHashDatabaseBuildTask",
"ImageHashesDatabaseBuilder",
]

View File

@ -0,0 +1,112 @@
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Callable, List
import cv2
from arcaea_offline_ocr.core import hashers
from arcaea_offline_ocr.providers import ImageCategory
from arcaea_offline_ocr.providers.ihdb import (
PROP_KEY_BUILT_AT,
PROP_KEY_HASH_SIZE,
PROP_KEY_HIGH_FREQ_FACTOR,
ImageHashDatabaseIdProvider,
ImageHashType,
)
if TYPE_CHECKING:
from sqlite3 import Connection
from arcaea_offline_ocr.types import Mat
def _default_imread_gray(image_path: str):
return cv2.cvtColor(cv2.imread(image_path, cv2.IMREAD_COLOR), cv2.COLOR_BGR2GRAY)
@dataclass
class ImageHashDatabaseBuildTask:
image_path: str
image_id: str
category: ImageCategory
imread_function: Callable[[str], "Mat"] = _default_imread_gray
@dataclass
class _ImageHash:
image_id: str
category: ImageCategory
image_hash_type: ImageHashType
hash: bytes
class ImageHashesDatabaseBuilder:
@staticmethod
def __insert_property(conn: "Connection", key: str, value: str):
return conn.execute(
"INSERT INTO properties (key, value) VALUES (?, ?)",
(key, value),
)
@classmethod
def build(
cls,
conn: "Connection",
tasks: List[ImageHashDatabaseBuildTask],
*,
hash_size: int = 16,
high_freq_factor: int = 4,
):
hashes: List[_ImageHash] = []
for task in tasks:
img_gray = task.imread_function(task.image_path)
for hash_type, hash_mat in [
(
ImageHashType.AVERAGE,
hashers.average(img_gray, hash_size),
),
(
ImageHashType.DCT,
hashers.dct(img_gray, hash_size, high_freq_factor),
),
(
ImageHashType.DIFFERENCE,
hashers.difference(img_gray, hash_size),
),
]:
hashes.append(
_ImageHash(
image_id=task.image_id,
image_hash_type=hash_type,
category=task.category,
hash=ImageHashDatabaseIdProvider.hash_mat_to_bytes(hash_mat),
)
)
conn.execute("CREATE TABLE properties (`key` VARCHAR, `value` VARCHAR)")
conn.execute(
"""CREATE TABLE hashes (
`id` VARCHAR,
`category` INTEGER,
`hash_type` INTEGER,
`hash` BLOB
)"""
)
now = datetime.now(tz=timezone.utc)
timestamp = int(now.timestamp() * 1000)
cls.__insert_property(conn, PROP_KEY_HASH_SIZE, str(hash_size))
cls.__insert_property(conn, PROP_KEY_HIGH_FREQ_FACTOR, str(high_freq_factor))
cls.__insert_property(conn, PROP_KEY_BUILT_AT, str(timestamp))
conn.executemany(
"INSERT INTO hashes (`id`, `category`, `hash_type`, `hash`) VALUES (?, ?, ?, ?)",
[
(it.image_id, it.category.value, it.image_hash_type.value, it.hash)
for it in hashes
],
)
conn.commit()

View File

@ -0,0 +1,3 @@
from .index import average, dct, difference
__all__ = ["average", "dct", "difference"]

View File

@ -0,0 +1,7 @@
import cv2
from arcaea_offline_ocr.types import Mat
def _resize_image(src: Mat, dsize: ...) -> Mat:
return cv2.resize(src, dsize, fx=0, fy=0, interpolation=cv2.INTER_AREA)

View File

@ -0,0 +1,35 @@
import cv2
import numpy as np
from arcaea_offline_ocr.types import Mat
from ._common import _resize_image
def average(img_gray: Mat, hash_size: int) -> Mat:
img_resized = _resize_image(img_gray, (hash_size, hash_size))
diff = img_resized > img_resized.mean()
return diff.flatten()
def difference(img_gray: Mat, hash_size: int) -> Mat:
img_size = (hash_size + 1, hash_size)
img_resized = _resize_image(img_gray, img_size)
previous = img_resized[:, :-1]
current = img_resized[:, 1:]
diff = previous > current
return diff.flatten()
def dct(img_gray: Mat, hash_size: int = 16, high_freq_factor: int = 4) -> Mat:
# TODO: consistency?
img_size_base = hash_size * high_freq_factor
img_size = (img_size_base, img_size_base)
img_resized = _resize_image(img_gray, img_size)
img_resized = img_resized.astype(np.float32)
dct_mat = cv2.dct(img_resized)
hash_mat = dct_mat[:hash_size, :hash_size]
return hash_mat > hash_mat.mean()

View File

@ -1,2 +0,0 @@
from .common import DeviceOcrResult
from .ocr import DeviceOcr

View File

@ -1,18 +0,0 @@
from typing import Optional
import attrs
@attrs.define
class DeviceOcrResult:
rating_class: int
pure: int
far: int
lost: int
score: int
max_recall: Optional[int] = None
song_id: Optional[str] = None
song_id_possibility: Optional[float] = None
clear_status: Optional[int] = None
partner_id: Optional[str] = None
partner_id_possibility: Optional[float] = None

View File

@ -1,3 +0,0 @@
from .definition import *
from .extractor import *
from .masker import *

View File

@ -1,2 +0,0 @@
from .auto import *
from .common import DeviceRois

View File

@ -1,15 +0,0 @@
from typing import Tuple
Rect = Tuple[int, int, int, int]
class DeviceRois:
pure: Rect
far: Rect
lost: Rect
score: Rect
rating_class: Rect
max_recall: Rect
jacket: Rect
clear_status: Rect
partner_icon: Rect

View File

@ -1 +0,0 @@
from .common import DeviceRoisExtractor

View File

@ -1,48 +0,0 @@
from ....crop import crop_xywh
from ....types import Mat
from ..definition.common import DeviceRois
class DeviceRoisExtractor:
def __init__(self, img: Mat, rois: DeviceRois):
self.img = img
self.sizes = rois
def __construct_int_rect(self, rect):
return tuple(round(r) for r in rect)
@property
def pure(self):
return crop_xywh(self.img, self.__construct_int_rect(self.sizes.pure))
@property
def far(self):
return crop_xywh(self.img, self.__construct_int_rect(self.sizes.far))
@property
def lost(self):
return crop_xywh(self.img, self.__construct_int_rect(self.sizes.lost))
@property
def score(self):
return crop_xywh(self.img, self.__construct_int_rect(self.sizes.score))
@property
def jacket(self):
return crop_xywh(self.img, self.__construct_int_rect(self.sizes.jacket))
@property
def rating_class(self):
return crop_xywh(self.img, self.__construct_int_rect(self.sizes.rating_class))
@property
def max_recall(self):
return crop_xywh(self.img, self.__construct_int_rect(self.sizes.max_recall))
@property
def clear_status(self):
return crop_xywh(self.img, self.__construct_int_rect(self.sizes.clear_status))
@property
def partner_icon(self):
return crop_xywh(self.img, self.__construct_int_rect(self.sizes.partner_icon))

View File

@ -1,2 +0,0 @@
from .auto import *
from .common import DeviceRoisMasker

View File

@ -1,59 +0,0 @@
from ....types import Mat
class DeviceRoisMasker:
@classmethod
def pure(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def far(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def lost(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def score(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def rating_class_pst(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def rating_class_prs(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def rating_class_ftr(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def rating_class_byd(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def rating_class_etr(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def max_recall(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def clear_status_track_lost(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def clear_status_track_complete(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def clear_status_full_recall(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()
@classmethod
def clear_status_pure_memory(cls, roi_bgr: Mat) -> Mat:
raise NotImplementedError()

View File

@ -1,119 +0,0 @@
import sqlite3
from typing import List, Union
import cv2
import numpy as np
from .types import Mat
def phash_opencv(img_gray, hash_size=8, highfreq_factor=4):
# type: (Union[Mat, np.ndarray], int, int) -> np.ndarray
"""
Perceptual Hash computation.
Implementation follows
http://www.hackerfactor.com/blog/index.php?/archives/432-Looks-Like-It.html
Adapted from `imagehash.phash`, pure opencv implementation
The result is slightly different from `imagehash.phash`.
"""
if hash_size < 2:
raise ValueError("Hash size must be greater than or equal to 2")
img_size = hash_size * highfreq_factor
image = cv2.resize(img_gray, (img_size, img_size), interpolation=cv2.INTER_LANCZOS4)
image = np.float32(image)
dct = cv2.dct(image)
dctlowfreq = dct[:hash_size, :hash_size]
med = np.median(dctlowfreq)
diff = dctlowfreq > med
return diff
def hamming_distance_sql_function(user_input, db_entry) -> int:
return np.count_nonzero(
np.frombuffer(user_input, bool) ^ np.frombuffer(db_entry, bool)
)
class ImagePhashDatabase:
def __init__(self, db_path: str):
with sqlite3.connect(db_path) as conn:
self.hash_size = int(
conn.execute(
"SELECT value FROM properties WHERE key = 'hash_size'"
).fetchone()[0]
)
self.highfreq_factor = int(
conn.execute(
"SELECT value FROM properties WHERE key = 'highfreq_factor'"
).fetchone()[0]
)
self.built_timestamp = int(
conn.execute(
"SELECT value FROM properties WHERE key = 'built_timestamp'"
).fetchone()[0]
)
self.ids: List[str] = [
i[0] for i in conn.execute("SELECT id FROM hashes").fetchall()
]
self.hashes_byte = [
i[0] for i in conn.execute("SELECT hash FROM hashes").fetchall()
]
self.hashes = [np.frombuffer(hb, bool) for hb in self.hashes_byte]
self.jacket_ids: List[str] = []
self.jacket_hashes = []
self.partner_icon_ids: List[str] = []
self.partner_icon_hashes = []
for _id, _hash in zip(self.ids, self.hashes):
id_splitted = _id.split("||")
if len(id_splitted) > 1 and id_splitted[0] == "partner_icon":
self.partner_icon_ids.append(id_splitted[1])
self.partner_icon_hashes.append(_hash)
else:
self.jacket_ids.append(_id)
self.jacket_hashes.append(_hash)
def calculate_phash(self, img_gray: Mat):
return phash_opencv(
img_gray, hash_size=self.hash_size, highfreq_factor=self.highfreq_factor
)
def lookup_hash(self, image_hash: np.ndarray, *, limit: int = 5):
image_hash = image_hash.flatten()
xor_results = [
(id, np.count_nonzero(image_hash ^ h))
for id, h in zip(self.ids, self.hashes)
]
return sorted(xor_results, key=lambda r: r[1])[:limit]
def lookup_image(self, img_gray: Mat):
image_hash = self.calculate_phash(img_gray)
return self.lookup_hash(image_hash)[0]
def lookup_jackets(self, img_gray: Mat, *, limit: int = 5):
image_hash = self.calculate_phash(img_gray).flatten()
xor_results = [
(id, np.count_nonzero(image_hash ^ h))
for id, h in zip(self.jacket_ids, self.jacket_hashes)
]
return sorted(xor_results, key=lambda r: r[1])[:limit]
def lookup_jacket(self, img_gray: Mat):
return self.lookup_jackets(img_gray)[0]
def lookup_partner_icons(self, img_gray: Mat, *, limit: int = 5):
image_hash = self.calculate_phash(img_gray).flatten()
xor_results = [
(id, np.count_nonzero(image_hash ^ h))
for id, h in zip(self.partner_icon_ids, self.partner_icon_hashes)
]
return sorted(xor_results, key=lambda r: r[1])[:limit]
def lookup_partner_icon(self, img_gray: Mat):
return self.lookup_partner_icons(img_gray)[0]

View File

@ -0,0 +1,12 @@
from .base import ImageCategory, ImageIdProvider, ImageIdProviderResult, OcrTextProvider
from .ihdb import ImageHashDatabaseIdProvider
from .knn import OcrKNearestTextProvider
__all__ = [
"ImageCategory",
"ImageHashDatabaseIdProvider",
"OcrKNearestTextProvider",
"ImageIdProvider",
"OcrTextProvider",
"ImageIdProviderResult",
]

View File

@ -0,0 +1,38 @@
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import IntEnum
from typing import TYPE_CHECKING, Any, Sequence, Optional
if TYPE_CHECKING:
from ..types import Mat
class OcrTextProvider(ABC):
@abstractmethod
def result_raw(self, img: "Mat", /, *args, **kwargs) -> Any: ...
@abstractmethod
def result(self, img: "Mat", /, *args, **kwargs) -> Optional[str]: ...
class ImageCategory(IntEnum):
JACKET = 0
PARTNER_ICON = 1
@dataclass(kw_only=True)
class ImageIdProviderResult:
image_id: str
category: ImageCategory
confidence: float
class ImageIdProvider(ABC):
@abstractmethod
def result(
self, img: "Mat", category: ImageCategory, /, *args, **kwargs
) -> ImageIdProviderResult: ...
@abstractmethod
def results(
self, img: "Mat", category: ImageCategory, /, *args, **kwargs
) -> Sequence[ImageIdProviderResult]: ...

View File

@ -0,0 +1,194 @@
import sqlite3
from dataclasses import dataclass
from datetime import datetime, timezone
from enum import IntEnum
from typing import TYPE_CHECKING, Any, Callable, List, Optional, TypeVar
from arcaea_offline_ocr.core import hashers
from .base import ImageCategory, ImageIdProvider, ImageIdProviderResult
if TYPE_CHECKING:
from arcaea_offline_ocr.types import Mat
T = TypeVar("T")
PROP_KEY_HASH_SIZE = "hash_size"
PROP_KEY_HIGH_FREQ_FACTOR = "high_freq_factor"
PROP_KEY_BUILT_AT = "built_at"
def _sql_hamming_distance(hash1: bytes, hash2: bytes):
assert len(hash1) == len(hash2), "hash size does not match!"
count = sum(1 for byte1, byte2 in zip(hash1, hash2) if byte1 != byte2)
return count
class ImageHashType(IntEnum):
AVERAGE = 0
DIFFERENCE = 1
DCT = 2
@dataclass(kw_only=True)
class ImageHashDatabaseIdProviderResult(ImageIdProviderResult):
image_hash_type: ImageHashType
class MissingPropertiesError(Exception):
keys: List[str]
def __init__(self, keys, *args):
super().__init__(*args)
self.keys = keys
class ImageHashDatabaseIdProvider(ImageIdProvider):
def __init__(self, conn: sqlite3.Connection):
self.conn = conn
self.conn.create_function("HAMMING_DISTANCE", 2, _sql_hamming_distance)
self.properties = {
PROP_KEY_HASH_SIZE: -1,
PROP_KEY_HIGH_FREQ_FACTOR: -1,
PROP_KEY_BUILT_AT: None,
}
self._hashes_count = {
ImageCategory.JACKET: 0,
ImageCategory.PARTNER_ICON: 0,
}
self._hash_length: int = -1
self._initialize()
@property
def hash_size(self) -> int:
return self.properties[PROP_KEY_HASH_SIZE]
@property
def high_freq_factor(self) -> int:
return self.properties[PROP_KEY_HIGH_FREQ_FACTOR]
@property
def built_at(self) -> Optional[datetime]:
return self.properties.get(PROP_KEY_BUILT_AT)
@property
def hash_length(self):
return self._hash_length
def _initialize(self):
def get_property(key, converter: Callable[[Any], T]) -> Optional[T]:
result = self.conn.execute(
"SELECT value FROM properties WHERE key = ?",
(key,),
).fetchone()
return converter(result[0]) if result is not None else None
def set_hashes_count(category: ImageCategory):
self._hashes_count[category] = self.conn.execute(
"SELECT COUNT(DISTINCT `id`) FROM hashes WHERE category = ?",
(category.value,),
).fetchone()[0]
properties_converter_map = {
PROP_KEY_HASH_SIZE: lambda x: int(x),
PROP_KEY_HIGH_FREQ_FACTOR: lambda x: int(x),
PROP_KEY_BUILT_AT: lambda ts: datetime.fromtimestamp(
int(ts) / 1000, tz=timezone.utc
),
}
required_properties = [PROP_KEY_HASH_SIZE, PROP_KEY_HIGH_FREQ_FACTOR]
missing_properties = []
for property_key, converter in properties_converter_map.items():
value = get_property(property_key, converter)
if value is None:
if property_key in required_properties:
missing_properties.append(property_key)
continue
self.properties[property_key] = value
if missing_properties:
raise MissingPropertiesError(keys=missing_properties)
set_hashes_count(ImageCategory.JACKET)
set_hashes_count(ImageCategory.PARTNER_ICON)
self._hash_length = self.hash_size**2
def lookup_hash(
self, category: ImageCategory, hash_type: ImageHashType, hash: bytes
) -> List[ImageHashDatabaseIdProviderResult]:
cursor = self.conn.execute(
"""
SELECT
`id`,
HAMMING_DISTANCE(hash, ?) AS distance
FROM hashes
WHERE category = ? AND hash_type = ?
ORDER BY distance ASC LIMIT 10""",
(hash, category.value, hash_type.value),
)
results = []
for id_, distance in cursor.fetchall():
results.append(
ImageHashDatabaseIdProviderResult(
image_id=id_,
category=category,
confidence=(self.hash_length - distance) / self.hash_length,
image_hash_type=hash_type,
)
)
return results
@staticmethod
def hash_mat_to_bytes(hash: "Mat") -> bytes:
return bytes([255 if b else 0 for b in hash.flatten()])
def results(self, img: "Mat", category: ImageCategory, /):
results: List[ImageHashDatabaseIdProviderResult] = []
results.extend(
self.lookup_hash(
category,
ImageHashType.AVERAGE,
self.hash_mat_to_bytes(hashers.average(img, self.hash_size)),
)
)
results.extend(
self.lookup_hash(
category,
ImageHashType.DIFFERENCE,
self.hash_mat_to_bytes(hashers.difference(img, self.hash_size)),
)
)
results.extend(
self.lookup_hash(
category,
ImageHashType.DCT,
self.hash_mat_to_bytes(
hashers.dct(img, self.hash_size, self.high_freq_factor)
),
)
)
return results
def result(
self,
img: "Mat",
category: ImageCategory,
/,
*,
hash_type: ImageHashType = ImageHashType.DCT,
):
return [
it for it in self.results(img, category) if it.image_hash_type == hash_type
][0]

View File

@ -1,18 +1,19 @@
import logging
import math
from typing import Optional, Sequence, Tuple
from typing import TYPE_CHECKING, Callable, Optional, Sequence, Tuple
import cv2
import numpy as np
from .crop import crop_xywh
from .types import Mat
from ..crop import crop_xywh
from .base import OcrTextProvider
__all__ = [
"FixRects",
"preprocess_hog",
"ocr_digits_by_contour_get_samples",
"ocr_digits_by_contour_knn",
]
if TYPE_CHECKING:
from cv2.ml import KNearest
from ..types import Mat
logger = logging.getLogger(__name__)
class FixRects:
@ -68,7 +69,7 @@ class FixRects:
@staticmethod
def split_connected(
img_masked: Mat,
img_masked: "Mat",
rects: Sequence[Tuple[int, int, int, int]],
rect_wh_ratio: float = 1.05,
width_range_ratio: float = 0.1,
@ -118,7 +119,7 @@ class FixRects:
return return_rects
def resize_fill_square(img: Mat, target: int = 20):
def resize_fill_square(img: "Mat", target: int = 20):
h, w = img.shape[:2]
if h > w:
new_h = target
@ -152,29 +153,88 @@ def preprocess_hog(digit_rois):
def ocr_digit_samples_knn(__samples, knn_model: cv2.ml.KNearest, k: int = 4):
_, results, _, _ = knn_model.findNearest(__samples, k)
result_list = [int(r) for r in results.ravel()]
result_str = "".join(str(r) for r in result_list if r > -1)
return int(result_str) if result_str else 0
return [int(r) for r in results.ravel()]
def ocr_digits_by_contour_get_samples(__roi_gray: Mat, size: int):
roi = __roi_gray.copy()
contours, _ = cv2.findContours(roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
rects = [cv2.boundingRect(c) for c in contours]
rects = FixRects.connect_broken(rects, roi.shape[1], roi.shape[0])
rects = FixRects.split_connected(roi, rects)
rects = sorted(rects, key=lambda r: r[0])
# digit_rois = [cv2.resize(crop_xywh(roi, rect), size) for rect in rects]
digit_rois = [resize_fill_square(crop_xywh(roi, rect), size) for rect in rects]
return preprocess_hog(digit_rois)
class OcrKNearestTextProvider(OcrTextProvider):
_ContourFilter = Callable[["Mat"], bool]
_RectsFilter = Callable[[Sequence[int]], bool]
def __init__(self, model: "KNearest"):
self.model = model
def ocr_digits_by_contour_knn(
__roi_gray: Mat,
knn_model: cv2.ml.KNearest,
def contours(
self, img: "Mat", /, *, contours_filter: Optional[_ContourFilter] = None
):
cnts, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
if contours_filter:
cnts = list(filter(contours_filter, cnts))
return cnts
def result_raw(
self,
img: "Mat",
/,
*,
k=4,
size: int = 20,
) -> int:
samples = ocr_digits_by_contour_get_samples(__roi_gray, size)
return ocr_digit_samples_knn(samples, knn_model, k)
fix_rects: bool = True,
contours_filter: Optional[_ContourFilter] = None,
rects_filter: Optional[_RectsFilter] = None,
):
"""
:param img: grayscaled roi
"""
try:
cnts, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if contours_filter:
cnts = list(filter(contours_filter, cnts))
rects = [cv2.boundingRect(cnt) for cnt in cnts]
if fix_rects and rects_filter:
rects = FixRects.connect_broken(rects, img.shape[1], img.shape[0]) # type: ignore
rects = list(filter(rects_filter, rects))
rects = FixRects.split_connected(img, rects)
elif fix_rects:
rects = FixRects.connect_broken(rects, img.shape[1], img.shape[0]) # type: ignore
rects = FixRects.split_connected(img, rects)
elif rects_filter:
rects = list(filter(rects_filter, rects))
rects = sorted(rects, key=lambda r: r[0])
digits = []
for rect in rects:
digit = crop_xywh(img, rect)
digit = resize_fill_square(digit, 20)
digits.append(digit)
samples = preprocess_hog(digits)
return ocr_digit_samples_knn(samples, self.model)
except Exception:
logger.exception("Error occurred during KNearest OCR")
return None
def result(
self,
img: "Mat",
/,
*,
fix_rects: bool = True,
contours_filter: Optional[_ContourFilter] = None,
rects_filter: Optional[_RectsFilter] = None,
):
"""
:param img: grayscaled roi
"""
raw = self.result_raw(
img,
fix_rects=fix_rects,
contours_filter=contours_filter,
rects_filter=rects_filter,
)
return (
"".join(["".join(str(r) for r in raw if r > -1)])
if raw is not None
else None
)

View File

@ -0,0 +1,3 @@
from .chieri import ChieriBotV4Best30Scenario
__all__ = ["ChieriBotV4Best30Scenario"]

View File

@ -0,0 +1,22 @@
from abc import abstractmethod
from typing import TYPE_CHECKING, List
from arcaea_offline_ocr.scenarios.base import OcrScenario, OcrScenarioResult
if TYPE_CHECKING:
from arcaea_offline_ocr.types import Mat
class Best30Scenario(OcrScenario):
@abstractmethod
def components(self, img: "Mat", /) -> List["Mat"]: ...
@abstractmethod
def result(self, component_img: "Mat", /, *args, **kwargs) -> OcrScenarioResult: ...
@abstractmethod
def results(self, img: "Mat", /, *args, **kwargs) -> List[OcrScenarioResult]:
"""
Commonly a shorthand for `[self.result(comp) for comp in self.components(img)]`
"""
...

View File

@ -0,0 +1,3 @@
from .v4 import ChieriBotV4Best30Scenario
__all__ = ["ChieriBotV4Best30Scenario"]

View File

@ -0,0 +1,3 @@
from .impl import ChieriBotV4Best30Scenario
__all__ = ["ChieriBotV4Best30Scenario"]

View File

@ -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)

View File

@ -1,60 +1,47 @@
from math import floor
from typing import List, Optional, Tuple
import cv2
import numpy as np
from ....crop import crop_xywh
from ....ocr import (
FixRects,
ocr_digits_by_contour_knn,
preprocess_hog,
resize_fill_square,
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,
FAR_BG_MAX_HSV,
FAR_BG_MIN_HSV,
FTR_MAX_HSV,
FTR_MIN_HSV,
LOST_BG_MAX_HSV,
LOST_BG_MIN_HSV,
PRS_MAX_HSV,
PRS_MIN_HSV,
PURE_BG_MAX_HSV,
PURE_BG_MIN_HSV,
)
from ....phash_db import ImagePhashDatabase
from ....types import Mat
from ....utils import construct_int_xywh_rect
from ...shared import B30OcrResultItem
from .colors import *
from .rois import ChieriBotV4Rois
class ChieriBotV4Ocr:
class ChieriBotV4Best30Scenario(Best30Scenario):
def __init__(
self,
score_knn: cv2.ml.KNearest,
pfl_knn: cv2.ml.KNearest,
phash_db: ImagePhashDatabase,
factor: Optional[float] = 1.0,
score_knn_provider: OcrKNearestTextProvider,
pfl_knn_provider: OcrKNearestTextProvider,
image_id_provider: ImageIdProvider,
factor: float = 1.0,
):
self.__score_knn = score_knn
self.__pfl_knn = pfl_knn
self.__phash_db = phash_db
self.__rois = ChieriBotV4Rois(factor)
@property
def score_knn(self):
return self.__score_knn
@score_knn.setter
def score_knn(self, knn_digits_model: cv2.ml.KNearest):
self.__score_knn = knn_digits_model
@property
def pfl_knn(self):
return self.__pfl_knn
@pfl_knn.setter
def pfl_knn(self, knn_digits_model: cv2.ml.KNearest):
self.__pfl_knn = knn_digits_model
@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.pfl_knn_provider = pfl_knn_provider
self.score_knn_provider = score_knn_provider
self.image_id_provider = image_id_provider
@property
def rois(self):
@ -72,9 +59,8 @@ class ChieriBotV4Ocr:
self.factor = img.shape[0] / 4400
def ocr_component_rating_class(self, component_bgr: Mat) -> int:
rating_class_rect = construct_int_xywh_rect(
self.rois.component_rois.rating_class_rect
)
rating_class_rect = self.rois.component_rois.rating_class_rect.rounded()
rating_class_roi = crop_xywh(component_bgr, rating_class_rect)
rating_class_roi = cv2.cvtColor(rating_class_roi, cv2.COLOR_BGR2HSV)
rating_class_masks = [
@ -88,18 +74,16 @@ 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):
jacket_rect = construct_int_xywh_rect(
self.rois.component_rois.jacket_rect, floor
)
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
score_rect = construct_int_xywh_rect(self.rois.component_rois.score_rect)
score_rect = self.rois.component_rois.score_rect.rounded()
score_roi = cv2.cvtColor(
crop_xywh(component_bgr, score_rect), cv2.COLOR_BGR2GRAY
)
@ -117,9 +101,13 @@ class ChieriBotV4Ocr:
if rect[3] > score_roi.shape[0] * 0.5:
continue
score_roi = cv2.fillPoly(score_roi, [contour], 0)
return ocr_digits_by_contour_knn(score_roi, self.score_knn)
def find_pfl_rects(self, component_pfl_processed: Mat) -> List[List[int]]:
ocr_result = self.score_knn_provider.result(score_roi)
return int(ocr_result) if ocr_result else 0
def find_pfl_rects(
self, component_pfl_processed: Mat
) -> List[Tuple[int, int, int, int]]:
# sourcery skip: inline-immediately-returned-variable
pfl_roi_find = cv2.morphologyEx(
component_pfl_processed,
@ -146,7 +134,7 @@ class ChieriBotV4Ocr:
return pfl_rects_adjusted
def preprocess_component_pfl(self, component_bgr: Mat) -> Mat:
pfl_rect = construct_int_xywh_rect(self.rois.component_rois.pfl_rect)
pfl_rect = self.rois.component_rois.pfl_rect.rounded()
pfl_roi = crop_xywh(component_bgr, pfl_rect)
pfl_roi_hsv = cv2.cvtColor(pfl_roi, cv2.COLOR_BGR2HSV)
@ -193,51 +181,43 @@ class ChieriBotV4Ocr:
pure_far_lost = []
for pfl_roi_rect in pfl_rects:
roi = crop_xywh(pfl_roi, pfl_roi_rect)
digit_contours, _ = cv2.findContours(
roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
)
digit_rects = [cv2.boundingRect(c) for c in digit_contours]
digit_rects = FixRects.connect_broken(
digit_rects, roi.shape[1], roi.shape[0]
)
digit_rects = FixRects.split_connected(roi, digit_rects)
digit_rects = sorted(digit_rects, key=lambda r: r[0])
digits = []
for digit_rect in digit_rects:
digit = crop_xywh(roi, digit_rect)
digit = resize_fill_square(digit, 20)
digits.append(digit)
samples = preprocess_hog(digits)
result = self.pfl_knn_provider.result(roi)
pure_far_lost.append(int(result) if result else None)
_, results, _, _ = self.pfl_knn.findNearest(samples, 4)
results = [str(int(i)) for i in results.ravel()]
pure_far_lost.append(int("".join(results)))
return tuple(pure_far_lost)
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)]

View File

@ -1,12 +1,11 @@
from typing import List, Optional
from typing import List
from ....crop import crop_xywh
from ....types import Mat, XYWHRect
from ....utils import apply_factor, construct_int_xywh_rect
from arcaea_offline_ocr.crop import crop_xywh
from arcaea_offline_ocr.types import Mat, XYWHRect
class ChieriBotV4ComponentRois:
def __init__(self, factor: Optional[float] = 1.0):
def __init__(self, factor: float = 1.0):
self.__factor = factor
@property
@ -19,43 +18,43 @@ class ChieriBotV4ComponentRois:
@property
def top_font_color_detect(self):
return apply_factor((35, 10, 120, 100), self.factor)
return XYWHRect(35, 10, 120, 100), self.factor
@property
def bottom_font_color_detect(self):
return apply_factor((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((21, 40, 7, 20), self.factor)
return XYWHRect(21, 40, 7, 20) * self.factor
@property
def title_rect(self):
return apply_factor((35, 10, 430, 50), self.factor)
return XYWHRect(35, 10, 430, 50) * self.factor
@property
def jacket_rect(self):
return apply_factor((263, 0, 239, 239), self.factor)
return XYWHRect(263, 0, 239, 239) * self.factor
@property
def score_rect(self):
return apply_factor((30, 60, 270, 55), self.factor)
return XYWHRect(30, 60, 270, 55) * self.factor
@property
def pfl_rect(self):
return apply_factor((50, 125, 80, 100), self.factor)
return XYWHRect(50, 125, 80, 100) * self.factor
@property
def date_rect(self):
return apply_factor((205, 200, 225, 25), self.factor)
return XYWHRect(205, 200, 225, 25) * self.factor
class ChieriBotV4Rois:
def __init__(self, factor: Optional[float] = 1.0):
def __init__(self, factor: float = 1.0):
self.__factor = factor
self.__component_rois = ChieriBotV4ComponentRois(factor)
@ -74,54 +73,53 @@ 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):
return 3
@property
def vertical_items(self):
return 10
vertical_items = 10
@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)
results = []
last_rect = first_rect
for vi in range(self.vertical_items):
rect = XYWHRect(*first_rect)
rect += (0, (self.vertical_gap + self.height) * vi, 0, 0)
for hi in range(self.horizontal_items):
if hi > 0:
rect += ((self.width + self.horizontal_gap), 0, 0, 0)
int_rect = construct_int_xywh_rect(rect)
results.append(crop_xywh(img_bgr, int_rect))
results.append(crop_xywh(img_bgr, rect.rounded()))
last_rect = rect
rect += (
last_rect += (
-(self.width + self.horizontal_gap) * 2,
self.height + self.b33_vertical_gap,
0,
@ -129,8 +127,7 @@ class ChieriBotV4Rois:
)
for hi in range(self.horizontal_items):
if hi > 0:
rect += ((self.width + self.horizontal_gap), 0, 0, 0)
int_rect = construct_int_xywh_rect(rect)
results.append(crop_xywh(img_bgr, int_rect))
last_rect += ((self.width + self.horizontal_gap), 0, 0, 0)
results.append(crop_xywh(img_bgr, last_rect.rounded()))
return results

View File

@ -0,0 +1,38 @@
from abc import ABC
from dataclasses import dataclass, field
from datetime import datetime
from typing import Sequence, Optional
from arcaea_offline_ocr.providers import ImageIdProviderResult
@dataclass(kw_only=True)
class OcrScenarioResult:
song_id: str
rating_class: int
score: int
song_id_results: Sequence[ImageIdProviderResult] = field(default_factory=lambda: [])
partner_id_results: Sequence[ImageIdProviderResult] = field(
default_factory=lambda: []
)
pure: Optional[int] = None
pure_inaccurate: Optional[int] = None
pure_early: Optional[int] = None
pure_late: Optional[int] = None
far: Optional[int] = None
far_inaccurate: Optional[int] = None
far_early: Optional[int] = None
far_late: Optional[int] = None
lost: Optional[int] = None
played_at: Optional[datetime] = None
max_recall: Optional[int] = None
clear_status: Optional[int] = None
clear_type: Optional[int] = None
modifier: Optional[int] = None
class OcrScenario(ABC):
pass

View File

@ -0,0 +1,13 @@
from .extractor import DeviceRoisExtractor
from .impl import DeviceScenario
from .masker import DeviceRoisMaskerAutoT1, DeviceRoisMaskerAutoT2
from .rois import DeviceRoisAutoT1, DeviceRoisAutoT2
__all__ = [
"DeviceRoisMaskerAutoT1",
"DeviceRoisMaskerAutoT2",
"DeviceRoisAutoT1",
"DeviceRoisAutoT2",
"DeviceRoisExtractor",
"DeviceScenario",
]

View File

@ -0,0 +1,8 @@
from abc import abstractmethod
from arcaea_offline_ocr.scenarios.base import OcrScenario, OcrScenarioResult
class DeviceScenarioBase(OcrScenario):
@abstractmethod
def result(self) -> OcrScenarioResult: ...

View File

@ -0,0 +1,3 @@
from .base import DeviceRoisExtractor
__all__ = ["DeviceRoisExtractor"]

View File

@ -0,0 +1,46 @@
from arcaea_offline_ocr.crop import crop_xywh
from arcaea_offline_ocr.types import Mat
from ..rois.base import DeviceRois
class DeviceRoisExtractor:
def __init__(self, img: Mat, rois: DeviceRois):
self.img = img
self.sizes = rois
@property
def pure(self):
return crop_xywh(self.img, self.sizes.pure.rounded())
@property
def far(self):
return crop_xywh(self.img, self.sizes.far.rounded())
@property
def lost(self):
return crop_xywh(self.img, self.sizes.lost.rounded())
@property
def score(self):
return crop_xywh(self.img, self.sizes.score.rounded())
@property
def jacket(self):
return crop_xywh(self.img, self.sizes.jacket.rounded())
@property
def rating_class(self):
return crop_xywh(self.img, self.sizes.rating_class.rounded())
@property
def max_recall(self):
return crop_xywh(self.img, self.sizes.max_recall.rounded())
@property
def clear_status(self):
return crop_xywh(self.img, self.sizes.clear_status.rounded())
@property
def partner_icon(self):
return crop_xywh(self.img, self.sizes.partner_icon.rounded())

View File

@ -1,58 +1,55 @@
import cv2
import numpy as np
from ..crop import crop_xywh
from ..ocr import (
FixRects,
ocr_digit_samples_knn,
ocr_digits_by_contour_knn,
preprocess_hog,
resize_fill_square,
from arcaea_offline_ocr.providers import (
ImageCategory,
ImageIdProvider,
OcrKNearestTextProvider,
)
from ..phash_db import ImagePhashDatabase
from ..types import Mat
from .common import DeviceOcrResult
from .rois.extractor import DeviceRoisExtractor
from .rois.masker import DeviceRoisMasker
from arcaea_offline_ocr.scenarios.base import OcrScenarioResult
from arcaea_offline_ocr.types import Mat
from .base import DeviceScenarioBase
from .extractor import DeviceRoisExtractor
from .masker import DeviceRoisMasker
class DeviceOcr:
class DeviceScenario(DeviceScenarioBase):
def __init__(
self,
extractor: DeviceRoisExtractor,
masker: DeviceRoisMasker,
knn_model: cv2.ml.KNearest,
phash_db: ImagePhashDatabase,
knn_provider: OcrKNearestTextProvider,
image_id_provider: ImageIdProvider,
):
self.extractor = extractor
self.masker = masker
self.knn_model = knn_model
self.phash_db = phash_db
self.knn_provider = knn_provider
self.image_id_provider = image_id_provider
def pfl(self, roi_gray: Mat, factor: float = 1.25):
contours, _ = cv2.findContours(
roi_gray, 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_gray.shape[1], roi_gray.shape[0])
def contour_filter(cnt):
return cv2.contourArea(cnt) >= 5 * factor
filtered_rects = [r for r in rects if r[2] >= 5 * factor and r[3] >= 6 * factor]
filtered_rects = FixRects.split_connected(roi_gray, filtered_rects)
filtered_rects = sorted(filtered_rects, key=lambda r: r[0])
contours = self.knn_provider.contours(roi_gray)
contours_filtered = self.knn_provider.contours(
roi_gray, contours_filter=contour_filter
)
roi_ocr = roi_gray.copy()
filtered_contours_flattened = {tuple(c.flatten()) for c in filtered_contours}
contours_filtered_flattened = {tuple(c.flatten()) for c in contours_filtered}
for contour in contours:
if tuple(contour.flatten()) in filtered_contours_flattened:
if tuple(contour.flatten()) in contours_filtered_flattened:
continue
roi_ocr = cv2.fillPoly(roi_ocr, [contour], [0])
digit_rois = [
resize_fill_square(crop_xywh(roi_ocr, r), 20) for r in filtered_rects
]
samples = preprocess_hog(digit_rois)
return ocr_digit_samples_knn(samples, self.knn_model)
ocr_result = self.knn_provider.result(
roi_ocr,
contours_filter=lambda cnt: cv2.contourArea(cnt) >= 5 * factor,
rects_filter=lambda rect: rect[2] >= 5 * factor and rect[3] >= 6 * factor,
)
return int(ocr_result) if ocr_result else 0
def pure(self):
return self.pfl(self.masker.pure(self.extractor.pure))
@ -65,13 +62,14 @@ class DeviceOcr:
def score(self):
roi = self.masker.score(self.extractor.score)
contours, _ = cv2.findContours(roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
contours = self.knn_provider.contours(roi)
for contour in contours:
if (
cv2.boundingRect(contour)[3] < roi.shape[0] * 0.6
): # h < score_component_h * 0.6
roi = cv2.fillPoly(roi, [contour], [0])
return ocr_digits_by_contour_knn(roi, self.knn_model)
ocr_result = self.knn_provider.result(roi)
return int(ocr_result) if ocr_result else 0
def rating_class(self):
roi = self.extractor.rating_class
@ -85,9 +83,10 @@ class DeviceOcr:
return max(enumerate(results), key=lambda i: np.count_nonzero(i[1]))[0]
def max_recall(self):
return ocr_digits_by_contour_knn(
self.masker.max_recall(self.extractor.max_recall), self.knn_model
ocr_result = self.knn_provider.result(
self.masker.max_recall(self.extractor.max_recall)
)
return int(ocr_result) if ocr_result else None
def clear_status(self):
roi = self.extractor.clear_status
@ -99,18 +98,16 @@ class DeviceOcr:
]
return max(enumerate(results), key=lambda i: np.count_nonzero(i[1]))[0]
def lookup_song_id(self):
return self.phash_db.lookup_jacket(
cv2.cvtColor(self.extractor.jacket, cv2.COLOR_BGR2GRAY)
def song_id_results(self):
return self.image_id_provider.results(
cv2.cvtColor(self.extractor.jacket, cv2.COLOR_BGR2GRAY),
ImageCategory.JACKET,
)
def song_id(self):
return self.lookup_song_id()[0]
@staticmethod
def preprocess_char_icon(img_gray: Mat):
h, w = img_gray.shape[:2]
img = cv2.copyMakeBorder(img_gray, w - h, 0, 0, 0, cv2.BORDER_REPLICATE)
img = cv2.copyMakeBorder(img_gray, max(w - h, 0), 0, 0, 0, cv2.BORDER_REPLICATE)
h, w = img.shape[:2]
img = cv2.fillPoly(
img,
@ -120,21 +117,19 @@ class DeviceOcr:
np.array([[0, h], [round(w / 2), h], [0, round(h / 2)]], np.int32),
np.array([[w, h], [round(w / 2), h], [w, round(h / 2)]], np.int32),
],
(128),
(128,),
)
return img
def lookup_partner_id(self):
return self.phash_db.lookup_partner_icon(
def partner_id_results(self):
return self.image_id_provider.results(
self.preprocess_char_icon(
cv2.cvtColor(self.extractor.partner_icon, cv2.COLOR_BGR2GRAY)
)
),
ImageCategory.PARTNER_ICON,
)
def partner_id(self):
return self.lookup_partner_id()[0]
def ocr(self) -> DeviceOcrResult:
def result(self):
rating_class = self.rating_class()
pure = self.pure()
far = self.far()
@ -143,20 +138,18 @@ class DeviceOcr:
max_recall = self.max_recall()
clear_status = self.clear_status()
hash_len = self.phash_db.hash_size**2
song_id, song_id_distance = self.lookup_song_id()
partner_id, partner_id_distance = self.lookup_partner_id()
song_id_results = self.song_id_results()
partner_id_results = self.partner_id_results()
return DeviceOcrResult(
return OcrScenarioResult(
song_id=song_id_results[0].image_id,
song_id_results=song_id_results,
rating_class=rating_class,
pure=pure,
far=far,
lost=lost,
score=score,
max_recall=max_recall,
song_id=song_id,
song_id_possibility=1 - song_id_distance / hash_len,
partner_id_results=partner_id_results,
clear_status=clear_status,
partner_id=partner_id,
partner_id_possibility=1 - partner_id_distance / hash_len,
)

View File

@ -0,0 +1,9 @@
from .auto import DeviceRoisMaskerAuto, DeviceRoisMaskerAutoT1, DeviceRoisMaskerAutoT2
from .base import DeviceRoisMasker
__all__ = [
"DeviceRoisMaskerAuto",
"DeviceRoisMaskerAutoT1",
"DeviceRoisMaskerAutoT2",
"DeviceRoisMasker",
]

View File

@ -1,13 +1,12 @@
import cv2
import numpy as np
from ....types import Mat
from .common import DeviceRoisMasker
from arcaea_offline_ocr.types import Mat
from .base import DeviceRoisMasker
class DeviceRoisMaskerAuto(DeviceRoisMasker):
# pylint: disable=abstract-method
@staticmethod
def mask_bgr_in_hsv(roi_bgr: Mat, hsv_lower: Mat, hsv_upper: Mat):
return cv2.inRange(
@ -125,7 +124,7 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto):
class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto):
PFL_HSV_MIN = np.array([0, 0, 248], np.uint8)
PFL_HSV_MAX = np.array([179, 10, 255], np.uint8)
PFL_HSV_MAX = np.array([179, 40, 255], np.uint8)
SCORE_HSV_MIN = np.array([0, 0, 180], np.uint8)
SCORE_HSV_MAX = np.array([179, 255, 255], np.uint8)

View File

@ -0,0 +1,61 @@
from abc import ABC, abstractmethod
from arcaea_offline_ocr.types import Mat
class DeviceRoisMasker(ABC):
@classmethod
@abstractmethod
def pure(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def far(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def lost(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def score(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def rating_class_pst(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def rating_class_prs(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def rating_class_ftr(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def rating_class_byd(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def rating_class_etr(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def max_recall(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def clear_status_track_lost(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def clear_status_track_complete(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def clear_status_full_recall(cls, roi_bgr: Mat) -> Mat: ...
@classmethod
@abstractmethod
def clear_status_pure_memory(cls, roi_bgr: Mat) -> Mat: ...

View File

@ -0,0 +1,9 @@
from .auto import DeviceRoisAuto, DeviceRoisAutoT1, DeviceRoisAutoT2
from .base import DeviceRois
__all__ = [
"DeviceRois",
"DeviceRoisAuto",
"DeviceRoisAutoT1",
"DeviceRoisAutoT2",
]

View File

@ -1,6 +1,6 @@
from .common import DeviceRois
from arcaea_offline_ocr.types import XYWHRect
__all__ = ["DeviceRoisAuto", "DeviceRoisAutoT1", "DeviceRoisAutoT2"]
from .base import DeviceRois
class DeviceRoisAuto(DeviceRois):
@ -50,7 +50,7 @@ class DeviceRoisAutoT1(DeviceRoisAuto):
@property
def pure(self):
return (
return XYWHRect(
self.pfl_x,
self.layout_area_h_mid + 110 * self.factor,
self.pfl_w,
@ -59,7 +59,7 @@ class DeviceRoisAutoT1(DeviceRoisAuto):
@property
def far(self):
return (
return XYWHRect(
self.pfl_x,
self.pure[1] + self.pure[3] + 12 * self.factor,
self.pfl_w,
@ -68,7 +68,7 @@ class DeviceRoisAutoT1(DeviceRoisAuto):
@property
def lost(self):
return (
return XYWHRect(
self.pfl_x,
self.far[1] + self.far[3] + 10 * self.factor,
self.pfl_w,
@ -79,7 +79,7 @@ class DeviceRoisAutoT1(DeviceRoisAuto):
def score(self):
w = 280 * self.factor
h = 45 * self.factor
return (
return XYWHRect(
self.w_mid - w / 2,
self.layout_area_h_mid - 75 * self.factor - h,
w,
@ -88,7 +88,7 @@ class DeviceRoisAutoT1(DeviceRoisAuto):
@property
def rating_class(self):
return (
return XYWHRect(
self.w_mid - 610 * self.factor,
self.layout_area_h_mid - 180 * self.factor,
265 * self.factor,
@ -97,7 +97,7 @@ class DeviceRoisAutoT1(DeviceRoisAuto):
@property
def max_recall(self):
return (
return XYWHRect(
self.w_mid - 465 * self.factor,
self.layout_area_h_mid - 215 * self.factor,
150 * self.factor,
@ -106,7 +106,7 @@ class DeviceRoisAutoT1(DeviceRoisAuto):
@property
def jacket(self):
return (
return XYWHRect(
self.w_mid - 610 * self.factor,
self.layout_area_h_mid - 143 * self.factor,
375 * self.factor,
@ -117,7 +117,7 @@ class DeviceRoisAutoT1(DeviceRoisAuto):
def clear_status(self):
w = 550 * self.factor
h = 60 * self.factor
return (
return XYWHRect(
self.w_mid - w / 2,
self.layout_area_h_mid - 155 * self.factor - h,
w * 0.4,
@ -128,7 +128,7 @@ class DeviceRoisAutoT1(DeviceRoisAuto):
def partner_icon(self):
w = 90 * self.factor
h = 75 * self.factor
return (self.w_mid - w / 2, 0, w, h)
return XYWHRect(self.w_mid - w / 2, 0, w, h)
class DeviceRoisAutoT2(DeviceRoisAuto):
@ -174,7 +174,7 @@ class DeviceRoisAutoT2(DeviceRoisAuto):
@property
def pure(self):
return (
return XYWHRect(
self.pfl_x,
self.layout_area_h_mid + 175 * self.factor,
self.pfl_w,
@ -183,7 +183,7 @@ class DeviceRoisAutoT2(DeviceRoisAuto):
@property
def far(self):
return (
return XYWHRect(
self.pfl_x,
self.pure[1] + self.pure[3] + 30 * self.factor,
self.pfl_w,
@ -192,7 +192,7 @@ class DeviceRoisAutoT2(DeviceRoisAuto):
@property
def lost(self):
return (
return XYWHRect(
self.pfl_x,
self.far[1] + self.far[3] + 35 * self.factor,
self.pfl_w,
@ -203,7 +203,7 @@ class DeviceRoisAutoT2(DeviceRoisAuto):
def score(self):
w = 420 * self.factor
h = 70 * self.factor
return (
return XYWHRect(
self.w_mid - w / 2,
self.layout_area_h_mid - 110 * self.factor - h,
w,
@ -212,7 +212,7 @@ class DeviceRoisAutoT2(DeviceRoisAuto):
@property
def rating_class(self):
return (
return XYWHRect(
max(0, self.w_mid - 965 * self.factor),
self.layout_area_h_mid - 330 * self.factor,
350 * self.factor,
@ -221,7 +221,7 @@ class DeviceRoisAutoT2(DeviceRoisAuto):
@property
def max_recall(self):
return (
return XYWHRect(
self.w_mid - 625 * self.factor,
self.layout_area_h_mid - 275 * self.factor,
150 * self.factor,
@ -230,7 +230,7 @@ class DeviceRoisAutoT2(DeviceRoisAuto):
@property
def jacket(self):
return (
return XYWHRect(
self.w_mid - 915 * self.factor,
self.layout_area_h_mid - 215 * self.factor,
565 * self.factor,
@ -241,7 +241,7 @@ class DeviceRoisAutoT2(DeviceRoisAuto):
def clear_status(self):
w = 825 * self.factor
h = 90 * self.factor
return (
return XYWHRect(
self.w_mid - w / 2,
self.layout_area_h_mid - 235 * self.factor - h,
w * 0.4,
@ -252,4 +252,4 @@ class DeviceRoisAutoT2(DeviceRoisAuto):
def partner_icon(self):
w = 135 * self.factor
h = 110 * self.factor
return (self.w_mid - w / 2, 0, w, h)
return XYWHRect(self.w_mid - w / 2, 0, w, h)

View File

@ -0,0 +1,33 @@
from abc import ABC, abstractmethod
from arcaea_offline_ocr.types import XYWHRect
class DeviceRois(ABC):
@property
@abstractmethod
def pure(self) -> XYWHRect: ...
@property
@abstractmethod
def far(self) -> XYWHRect: ...
@property
@abstractmethod
def lost(self) -> XYWHRect: ...
@property
@abstractmethod
def score(self) -> XYWHRect: ...
@property
@abstractmethod
def rating_class(self) -> XYWHRect: ...
@property
@abstractmethod
def max_recall(self) -> XYWHRect: ...
@property
@abstractmethod
def jacket(self) -> XYWHRect: ...
@property
@abstractmethod
def clear_status(self) -> XYWHRect: ...
@property
@abstractmethod
def partner_icon(self) -> XYWHRect: ...

View File

@ -1,25 +1,42 @@
from collections.abc import Iterable
from typing import NamedTuple, Tuple, Union
from math import floor
from typing import Callable, NamedTuple, Union
import numpy as np
Mat = np.ndarray
_IntOrFloat = Union[int, float]
class XYWHRect(NamedTuple):
x: int
y: int
w: int
h: int
x: _IntOrFloat
y: _IntOrFloat
w: _IntOrFloat
h: _IntOrFloat
def __add__(self, other: Union["XYWHRect", Tuple[int, int, int, int]]):
if not isinstance(other, Iterable) or len(other) != 4:
raise ValueError()
def _to_int(self, func: Callable[[_IntOrFloat], int]):
return (func(self.x), func(self.y), func(self.w), func(self.h))
def rounded(self):
return self._to_int(round)
def floored(self):
return self._to_int(floor)
def __add__(self, other):
if not isinstance(other, (list, tuple)) or len(other) != 4:
raise TypeError()
return self.__class__(*[a + b for a, b in zip(self, other)])
def __sub__(self, other: Union["XYWHRect", Tuple[int, int, int, int]]):
if not isinstance(other, Iterable) or len(other) != 4:
raise ValueError()
def __sub__(self, other):
if not isinstance(other, (list, tuple)) or len(other) != 4:
raise TypeError()
return self.__class__(*[a - b for a, b in zip(self, other)])
def __mul__(self, other):
if not isinstance(other, (int, float)):
raise TypeError()
return self.__class__(*[v * other for v in self])

View File

@ -1,11 +1,6 @@
from collections.abc import Iterable
from typing import Callable, TypeVar, Union, overload
import cv2
import numpy as np
from .types import XYWHRect
__all__ = ["imread_unicode"]
@ -13,34 +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)
def construct_int_xywh_rect(
rect: XYWHRect, func: Callable[[Union[int, float]], int] = round
):
return XYWHRect(*[func(num) for num in rect])
@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, Iterable):
return item.__class__([i * factor for i in item])