10 Commits

Author SHA1 Message Date
f9c867e180 chore: v0.0.95 2023-10-25 18:08:29 +08:00
f2f854040f fix: cv2.Mat typing annotations 2023-10-23 13:29:17 +08:00
cf46bf7a59 chore: update dependencies 2023-10-23 13:28:59 +08:00
122a546174 refactor: CropBlackEdges 2023-10-22 01:55:20 +08:00
42bcd7b430 refactor: B30 ocr 2023-10-13 18:32:33 +08:00
3400df2d52 refactor!: remove utils.convert_to_srgb 2023-10-12 18:08:54 +08:00
4fd31b1e9b impr: module __init__ entries 2023-10-12 17:53:51 +08:00
82229b8b5c chore: remove custom cv2 type annotations (#8) 2023-10-12 01:50:27 +08:00
2895eb7233 refactor!: hog computing 2023-10-12 01:37:59 +08:00
02599780e3 impr: minor improvements 2023-10-12 01:36:57 +08:00
16 changed files with 174 additions and 291 deletions

View File

@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
[project]
name = "arcaea-offline-ocr"
version = "0.1.0"
version = "0.0.95"
authors = [{ name = "283375", email = "log_283375@163.com" }]
description = "Extract your Arcaea play result from screenshot."
readme = "README.md"
requires-python = ">=3.8"
dependencies = ["attrs==23.1.0", "numpy==1.25.2", "opencv-python==4.8.0.76"]
dependencies = ["attrs==23.1.0", "numpy==1.26.1", "opencv-python==4.8.1.78"]
classifiers = [
"Development Status :: 3 - Alpha",
"Programming Language :: Python :: 3",

View File

@ -1,3 +1,3 @@
attrs==23.1.0
numpy==1.25.2
opencv-python==4.8.0.76
numpy==1.26.1
opencv-python==4.8.1.78

View File

@ -3,12 +3,16 @@ from typing import List, Optional, Tuple
import cv2
import numpy as np
from PIL import Image
from ....crop import crop_xywh
from ....ocr import FixRects, ocr_digits_by_contour_knn, preprocess_hog
from ....ocr import (
FixRects,
ocr_digits_by_contour_knn,
preprocess_hog,
resize_fill_square,
)
from ....phash_db import ImagePhashDatabase
from ....types import Mat, cv2_ml_KNearest
from ....types import Mat
from ....utils import construct_int_xywh_rect
from ...shared import B30OcrResultItem
from .colors import *
@ -18,8 +22,8 @@ from .rois import ChieriBotV4Rois
class ChieriBotV4Ocr:
def __init__(
self,
score_knn: cv2_ml_KNearest,
pfl_knn: cv2_ml_KNearest,
score_knn: cv2.ml.KNearest,
pfl_knn: cv2.ml.KNearest,
phash_db: ImagePhashDatabase,
factor: Optional[float] = 1.0,
):
@ -33,7 +37,7 @@ class ChieriBotV4Ocr:
return self.__score_knn
@score_knn.setter
def score_knn(self, knn_digits_model: Mat):
def score_knn(self, knn_digits_model: cv2.ml.KNearest):
self.__score_knn = knn_digits_model
@property
@ -41,7 +45,7 @@ class ChieriBotV4Ocr:
return self.__pfl_knn
@pfl_knn.setter
def pfl_knn(self, knn_digits_model: Mat):
def pfl_knn(self, knn_digits_model: cv2.ml.KNearest):
self.__pfl_knn = knn_digits_model
@property
@ -84,14 +88,6 @@ class ChieriBotV4Ocr:
else:
return max(enumerate(rating_class_results), key=lambda i: i[1])[0] + 1
# def ocr_component_title(self, component_bgr: Mat) -> str:
# # sourcery skip: inline-immediately-returned-variable
# title_rect = construct_int_xywh_rect(self.rois.component_rois.title_rect)
# title_roi = crop_xywh(component_bgr, title_rect)
# ocr_result = self.sift_db.ocr(title_roi, cls=False)
# title = ocr_result[0][-1][1][0] if ocr_result and ocr_result[0] else ""
# return title
def ocr_component_song_id(self, component_bgr: Mat):
jacket_rect = construct_int_xywh_rect(
self.rois.component_rois.jacket_rect, floor
@ -99,20 +95,7 @@ class ChieriBotV4Ocr:
jacket_roi = cv2.cvtColor(
crop_xywh(component_bgr, jacket_rect), cv2.COLOR_BGR2GRAY
)
return self.phash_db.lookup_image(Image.fromarray(jacket_roi))[0]
# def ocr_component_score_paddle(self, component_bgr: Mat) -> int:
# # sourcery skip: inline-immediately-returned-variable
# score_rect = construct_int_xywh_rect(self.rois.component_rois.score_rect)
# score_roi = cv2.cvtColor(
# crop_xywh(component_bgr, score_rect), cv2.COLOR_BGR2GRAY
# )
# _, score_roi = cv2.threshold(
# score_roi, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU
# )
# score_str = self.sift_db.ocr(score_roi, cls=False)[0][-1][1][0]
# score = int(score_str.replace("'", "").replace(" ", ""))
# return score
return self.phash_db.lookup_jacket(jacket_roi)[0]
def ocr_component_score_knn(self, component_bgr: Mat) -> int:
# sourcery skip: inline-immediately-returned-variable
@ -222,7 +205,7 @@ class ChieriBotV4Ocr:
digits = []
for digit_rect in digit_rects:
digit = crop_xywh(roi, digit_rect)
digit = cv2.resize(digit, (20, 20))
digit = resize_fill_square(digit, 20)
digits.append(digit)
samples = preprocess_hog(digits)
@ -233,15 +216,6 @@ class ChieriBotV4Ocr:
except Exception:
return (None, None, None)
# def ocr_component_date(self, component_bgr: Mat):
# date_rect = construct_int_xywh_rect(self.rois.component_rois.date_rect)
# date_roi = cv2.cvtColor(crop_xywh(component_bgr, date_rect), cv2.COLOR_BGR2GRAY)
# _, date_roi = cv2.threshold(
# date_roi, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU
# )
# date_str = self.sift_db.ocr(date_roi, cls=False)[0][-1][1][0]
# return date_str
def ocr_component(self, component_bgr: Mat) -> B30OcrResultItem:
component_blur = cv2.GaussianBlur(component_bgr, (5, 5), 0)
rating_class = self.ocr_component_rating_class(component_blur)

View File

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

View File

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

View File

@ -13,6 +13,6 @@ class DeviceOcrResult:
max_recall: Optional[int] = None
song_id: Optional[str] = None
song_id_possibility: Optional[float] = None
clear_status: Optional[str] = None
clear_status: Optional[int] = None
partner_id: Optional[str] = None
partner_id_possibility: Optional[float] = None

View File

@ -10,6 +10,7 @@ from ..ocr import (
resize_fill_square,
)
from ..phash_db import ImagePhashDatabase
from ..types import Mat
from .common import DeviceOcrResult
from .rois.extractor import DeviceRoisExtractor
from .rois.masker import DeviceRoisMasker
@ -28,7 +29,7 @@ class DeviceOcr:
self.knn_model = knn_model
self.phash_db = phash_db
def pfl(self, roi_gray: cv2.Mat, factor: float = 1.25):
def pfl(self, roi_gray: Mat, factor: float = 1.25):
contours, _ = cv2.findContours(
roi_gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
)
@ -47,8 +48,7 @@ class DeviceOcr:
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])
resize_fill_square(crop_xywh(roi_ocr, r), 20) for r in filtered_rects
]
samples = preprocess_hog(digit_rois)
@ -106,7 +106,7 @@ class DeviceOcr:
return self.lookup_song_id()[0]
@staticmethod
def preprocess_char_icon(img_gray: cv2.Mat):
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)
h, w = img.shape[:2]

View File

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

View File

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

View File

@ -1,13 +1,12 @@
import cv2
from ....crop import crop_xywh
from ....types import Mat
from ..definition.common import DeviceRois
class DeviceRoisExtractor:
def __init__(self, img: cv2.Mat, sizes: DeviceRois):
def __init__(self, img: Mat, rois: DeviceRois):
self.img = img
self.sizes = sizes
self.sizes = rois
def __construct_int_rect(self, rect):
return tuple(round(r) for r in rect)

View File

@ -1,6 +1,7 @@
import cv2
import numpy as np
from ....types import Mat
from .common import DeviceRoisMasker
@ -40,7 +41,7 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto):
PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8)
@classmethod
def gray(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def gray(cls, roi_bgr: Mat) -> Mat:
bgr_value_equal_mask = np.max(roi_bgr, axis=2) - np.min(roi_bgr, axis=2) <= 5
img_bgr = roi_bgr.copy()
img_bgr[~bgr_value_equal_mask] = np.array([0, 0, 0], roi_bgr.dtype)
@ -49,19 +50,19 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto):
return cv2.inRange(img_bgr, cls.GRAY_BGR_MIN, cls.GRAY_BGR_MAX)
@classmethod
def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def pure(cls, roi_bgr: Mat) -> Mat:
return cls.gray(roi_bgr)
@classmethod
def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def far(cls, roi_bgr: Mat) -> Mat:
return cls.gray(roi_bgr)
@classmethod
def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def lost(cls, roi_bgr: Mat) -> Mat:
return cls.gray(roi_bgr)
@classmethod
def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def score(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.WHITE_HSV_MIN,
@ -69,35 +70,35 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto):
)
@classmethod
def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def rating_class_pst(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PST_HSV_MIN, cls.PST_HSV_MAX
)
@classmethod
def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def rating_class_prs(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PRS_HSV_MIN, cls.PRS_HSV_MAX
)
@classmethod
def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def rating_class_ftr(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.FTR_HSV_MIN, cls.FTR_HSV_MAX
)
@classmethod
def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def rating_class_byd(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.BYD_HSV_MIN, cls.BYD_HSV_MAX
)
@classmethod
def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def max_recall(cls, roi_bgr: Mat) -> Mat:
return cls.gray(roi_bgr)
@classmethod
def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def clear_status_track_lost(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.TRACK_LOST_HSV_MIN,
@ -105,7 +106,7 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto):
)
@classmethod
def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def clear_status_track_complete(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.TRACK_COMPLETE_HSV_MIN,
@ -113,7 +114,7 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto):
)
@classmethod
def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def clear_status_full_recall(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.FULL_RECALL_HSV_MIN,
@ -121,7 +122,7 @@ class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto):
)
@classmethod
def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def clear_status_pure_memory(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.PURE_MEMORY_HSV_MIN,
@ -164,25 +165,25 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto):
PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8)
@classmethod
def pfl(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def pfl(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PFL_HSV_MIN, cls.PFL_HSV_MAX
)
@classmethod
def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def pure(cls, roi_bgr: Mat) -> Mat:
return cls.pfl(roi_bgr)
@classmethod
def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def far(cls, roi_bgr: Mat) -> Mat:
return cls.pfl(roi_bgr)
@classmethod
def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def lost(cls, roi_bgr: Mat) -> Mat:
return cls.pfl(roi_bgr)
@classmethod
def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def score(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.WHITE_HSV_MIN,
@ -190,31 +191,31 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto):
)
@classmethod
def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def rating_class_pst(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PST_HSV_MIN, cls.PST_HSV_MAX
)
@classmethod
def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def rating_class_prs(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.PRS_HSV_MIN, cls.PRS_HSV_MAX
)
@classmethod
def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def rating_class_ftr(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.FTR_HSV_MIN, cls.FTR_HSV_MAX
)
@classmethod
def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def rating_class_byd(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), cls.BYD_HSV_MIN, cls.BYD_HSV_MAX
)
@classmethod
def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def max_recall(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.MAX_RECALL_HSV_MIN,
@ -222,7 +223,7 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto):
)
@classmethod
def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def clear_status_track_lost(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.TRACK_LOST_HSV_MIN,
@ -230,7 +231,7 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto):
)
@classmethod
def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def clear_status_track_complete(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.TRACK_COMPLETE_HSV_MIN,
@ -238,7 +239,7 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto):
)
@classmethod
def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def clear_status_full_recall(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.FULL_RECALL_HSV_MIN,
@ -246,7 +247,7 @@ class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto):
)
@classmethod
def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
def clear_status_pure_memory(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.PURE_MEMORY_HSV_MIN,

View File

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

View File

@ -1,13 +1,11 @@
import math
from copy import deepcopy
from typing import Optional, Sequence, Tuple
import cv2
import numpy as np
from numpy.linalg import norm
from .crop import crop_xywh
from .types import Mat, cv2_ml_KNearest
from .types import Mat
__all__ = [
"FixRects",
@ -64,8 +62,7 @@ class FixRects:
new_h = new_bottom - new_y
new_rects.append((new_x, new_y, new_w, new_h))
return_rects = deepcopy(rects)
return_rects = [r for r in return_rects if r not in consumed_rects]
return_rects = [r for r in rects if r not in consumed_rects]
return_rects.extend(new_rects)
return return_rects
@ -80,42 +77,42 @@ class FixRects:
new_rects = []
for rect in rects:
rx, ry, rw, rh = rect
if rw / rh > rect_wh_ratio:
# consider this is a connected contour
connected_rects.append(rect)
if rw / rh <= rect_wh_ratio:
continue
# find the thinnest part
border_ignore = round(rw * width_range_ratio)
img_cropped = crop_xywh(
img_masked,
(border_ignore, ry, rw - border_ignore, rh),
)
white_pixels = {} # dict[x, white_pixel_number]
for i in range(img_cropped.shape[1]):
col = img_cropped[:, i]
white_pixels[rx + border_ignore + i] = np.count_nonzero(col > 200)
least_white_pixels = min(v for v in white_pixels.values() if v > 0)
x_values = [
x
for x, pixel in white_pixels.items()
if pixel == least_white_pixels
]
# select only middle values
x_mean = np.mean(x_values)
x_std = np.std(x_values)
x_values = [
x
for x in x_values
if x_mean - x_std * 1.5 <= x <= x_mean + x_std * 1.5
]
x_mid = round(np.median(x_values))
connected_rects.append(rect)
# split the rect
new_rects.extend(
[(rx, ry, x_mid - rx, rh), (x_mid, ry, rx + rw - x_mid, rh)]
)
# find the thinnest part
border_ignore = round(rw * width_range_ratio)
img_cropped = crop_xywh(
img_masked,
(border_ignore, ry, rw - border_ignore, rh),
)
white_pixels = {} # dict[x, white_pixel_number]
for i in range(img_cropped.shape[1]):
col = img_cropped[:, i]
white_pixels[rx + border_ignore + i] = np.count_nonzero(col > 200)
if all(v == 0 for v in white_pixels.values()):
return rects
least_white_pixels = min(v for v in white_pixels.values() if v > 0)
x_values = [
x for x, pixel in white_pixels.items() if pixel == least_white_pixels
]
# select only middle values
x_mean = np.mean(x_values)
x_std = np.std(x_values)
x_values = [
x for x in x_values if x_mean - x_std * 1.5 <= x <= x_mean + x_std * 1.5
]
x_mid = round(np.median(x_values))
# split the rect
new_rects.extend(
[(rx, ry, x_mid - rx, rh), (x_mid, ry, rx + rw - x_mid, rh)]
)
return_rects = deepcopy(rects)
return_rects = [r for r in rects if r not in connected_rects]
return_rects.extend(new_rects)
return return_rects
@ -144,33 +141,16 @@ def resize_fill_square(img: Mat, target: int = 20):
def preprocess_hog(digit_rois):
# https://github.com/opencv/opencv/blob/f834736307c8328340aea48908484052170c9224/samples/python/digits.py
# https://learnopencv.com/handwritten-digits-classification-an-opencv-c-python-tutorial/
samples = []
for digit in digit_rois:
gx = cv2.Sobel(digit, cv2.CV_32F, 1, 0)
gy = cv2.Sobel(digit, cv2.CV_32F, 0, 1)
mag, ang = cv2.cartToPolar(gx, gy)
bin_n = 16
_bin = np.int32(bin_n * ang / (2 * np.pi))
bin_cells = _bin[:10, :10], _bin[10:, :10], _bin[:10, 10:], _bin[10:, 10:]
mag_cells = mag[:10, :10], mag[10:, :10], mag[:10, 10:], mag[10:, 10:]
hists = [
np.bincount(b.ravel(), m.ravel(), bin_n)
for b, m in zip(bin_cells, mag_cells)
]
hist = np.hstack(hists)
# transform to Hellinger kernel
eps = 1e-7
hist /= hist.sum() + eps
hist = np.sqrt(hist)
hist /= norm(hist) + eps
hog = cv2.HOGDescriptor((20, 20), (10, 10), (5, 5), (10, 10), 9)
hist = hog.compute(digit)
samples.append(hist)
return np.float32(samples)
def ocr_digit_samples_knn(__samples, knn_model: cv2_ml_KNearest, k: int = 4):
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)
@ -191,7 +171,7 @@ def ocr_digits_by_contour_get_samples(__roi_gray: Mat, size: int):
def ocr_digits_by_contour_knn(
__roi_gray: Mat,
knn_model: cv2_ml_KNearest,
knn_model: cv2.ml.KNearest,
*,
k=4,
size: int = 20,

View File

@ -4,9 +4,11 @@ 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[cv2.Mat, np.ndarray], int, int) -> np.ndarray
# type: (Union[Mat, np.ndarray], int, int) -> np.ndarray
"""
Perceptual Hash computation.
@ -76,7 +78,7 @@ class ImagePhashDatabase:
self.jacket_ids.append(id)
self.jacket_hashes.append(hash)
def calculate_phash(self, img_gray: cv2.Mat):
def calculate_phash(self, img_gray: Mat):
return phash_opencv(
img_gray, hash_size=self.hash_size, highfreq_factor=self.highfreq_factor
)
@ -89,11 +91,11 @@ class ImagePhashDatabase:
]
return sorted(xor_results, key=lambda r: r[1])[:limit]
def lookup_image(self, img_gray: cv2.Mat):
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: cv2.Mat, *, limit: int = 5):
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))
@ -101,10 +103,10 @@ class ImagePhashDatabase:
]
return sorted(xor_results, key=lambda r: r[1])[:limit]
def lookup_jacket(self, img_gray: cv2.Mat):
def lookup_jacket(self, img_gray: Mat):
return self.lookup_jackets(img_gray)[0]
def lookup_partner_icons(self, img_gray: cv2.Mat, *, limit: int = 5):
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))
@ -112,5 +114,5 @@ class ImagePhashDatabase:
]
return sorted(xor_results, key=lambda r: r[1])[:limit]
def lookup_partner_icon(self, img_gray: cv2.Mat):
def lookup_partner_icon(self, img_gray: Mat):
return self.lookup_partner_icons(img_gray)[0]

View File

@ -1,10 +1,9 @@
from collections.abc import Iterable
from typing import Any, NamedTuple, Protocol, Tuple, Union
from typing import NamedTuple, Tuple, Union
import numpy as np
# from pylance
Mat = np.ndarray[int, np.dtype[np.generic]]
Mat = np.ndarray
class XYWHRect(NamedTuple):
@ -24,19 +23,3 @@ class XYWHRect(NamedTuple):
raise ValueError()
return self.__class__(*[a - b for a, b in zip(self, other)])
class cv2_ml_StatModel(Protocol):
def predict(self, samples: np.ndarray, results: np.ndarray, flags: int = 0):
...
def train(self, samples: np.ndarray, layout: int, responses: np.ndarray):
...
class cv2_ml_KNearest(cv2_ml_StatModel, Protocol):
def findNearest(
self, samples: np.ndarray, k: int
) -> Tuple[Any, np.ndarray, np.ndarray, np.ndarray]:
"""cv.ml.KNearest.findNearest(samples, k[, results[, neighborResponses[, dist]]]) -> retval, results, neighborResponses, dist"""
...

View File

@ -1,17 +1,15 @@
import io
from collections.abc import Iterable
from typing import Callable, Tuple, TypeVar, Union, overload
from typing import Callable, TypeVar, Union, overload
import cv2
import numpy as np
from PIL import Image, ImageCms
from .types import Mat, XYWHRect
from .types import XYWHRect
__all__ = ["imread_unicode"]
def imread_unicode(filepath: str, flags: int = cv2.IMREAD_UNCHANGED) -> Mat:
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)
@ -46,25 +44,3 @@ def apply_factor(item, factor: float):
return item * factor
elif isinstance(item, Iterable):
return item.__class__([i * factor for i in item])
def convert_to_srgb(pil_img: Image.Image):
"""
Convert PIL image to sRGB color space (if possible)
and save the converted file.
https://stackoverflow.com/a/65667797/16484891
CC BY-SA 4.0
"""
icc = pil_img.info.get("icc_profile", "")
icc_conv = ""
if icc:
io_handle = io.BytesIO(icc) # virtual file
src_profile = ImageCms.ImageCmsProfile(io_handle)
dst_profile = ImageCms.createProfile("sRGB")
img_conv = ImageCms.profileToProfile(pil_img, src_profile, dst_profile)
icc_conv = img_conv.info.get("icc_profile", "")
return img_conv if icc != icc_conv else pil_img