17 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
5e0642c832 impr: max recall masking 2023-10-10 01:45:21 +08:00
ede2b4ec51 refactor: DeviceOcr 2023-10-10 01:45:02 +08:00
6a19ead8d1 refactor: module structure 2023-10-10 01:30:12 +08:00
d13076c667 feat: split jacket & partner lookup in phash database 2023-10-10 01:28:36 +08:00
1aa71685ce fix: what is this shit 2023-10-09 00:16:54 +08:00
c009a28f92 refactor: use opencv to calculate image phash 2023-10-03 16:38:11 +08:00
d5ccbd5a01 refactor: module structure 2023-10-03 15:05:03 +08:00
33 changed files with 844 additions and 857 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 ....phash_db import ImagePHashDatabase
from ....types import Mat, cv2_ml_KNearest
from ....ocr import (
FixRects,
ocr_digits_by_contour_knn,
preprocess_hog,
resize_fill_square,
)
from ....phash_db import ImagePhashDatabase
from ....types import Mat
from ....utils import construct_int_xywh_rect
from ...shared import B30OcrResultItem
from .colors import *
@ -18,9 +22,9 @@ from .rois import ChieriBotV4Rois
class ChieriBotV4Ocr:
def __init__(
self,
score_knn: cv2_ml_KNearest,
pfl_knn: cv2_ml_KNearest,
phash_db: ImagePHashDatabase,
score_knn: cv2.ml.KNearest,
pfl_knn: cv2.ml.KNearest,
phash_db: ImagePhashDatabase,
factor: Optional[float] = 1.0,
):
self.__score_knn = score_knn
@ -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
@ -49,7 +53,7 @@ class ChieriBotV4Ocr:
return self.__phash_db
@phash_db.setter
def phash_db(self, phash_db: ImagePHashDatabase):
def phash_db(self, phash_db: ImagePhashDatabase):
self.__phash_db = phash_db
@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

@ -10,7 +10,9 @@ class DeviceOcrResult:
far: int
lost: int
score: int
max_recall: int
max_recall: Optional[int] = None
song_id: Optional[str] = None
title: Optional[str] = None
clear_type: 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,6 +1,5 @@
import cv2
import numpy as np
from PIL import Image
from ..crop import crop_xywh
from ..ocr import (
@ -10,25 +9,27 @@ from ..ocr import (
preprocess_hog,
resize_fill_square,
)
from ..phash_db import ImagePHashDatabase
from .roi.extractor import DeviceRoiExtractor
from .roi.masker import DeviceRoiMasker
from ..phash_db import ImagePhashDatabase
from ..types import Mat
from .common import DeviceOcrResult
from .rois.extractor import DeviceRoisExtractor
from .rois.masker import DeviceRoisMasker
class DeviceOcr:
def __init__(
self,
extractor: DeviceRoiExtractor,
masker: DeviceRoiMasker,
extractor: DeviceRoisExtractor,
masker: DeviceRoisMasker,
knn_model: cv2.ml.KNearest,
phash_db: ImagePHashDatabase,
phash_db: ImagePhashDatabase,
):
self.extractor = extractor
self.masker = masker
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)
@ -97,5 +97,64 @@ 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(self):
return self.phash_db.lookup_image(Image.fromarray(self.extractor.jacket))[0]
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)
h, w = img.shape[:2]
img = cv2.fillPoly(
img,
[
np.array([[0, 0], [round(w / 2), 0], [0, round(h / 2)]], np.int32),
np.array([[w, 0], [round(w / 2), 0], [w, round(h / 2)]], np.int32),
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),
)
return img
def lookup_partner_id(self):
return self.phash_db.lookup_partner_icon(
self.preprocess_char_icon(
cv2.cvtColor(self.extractor.partner_icon, cv2.COLOR_BGR2GRAY)
)
)
def partner_id(self):
return self.lookup_partner_id()[0]
def ocr(self) -> DeviceOcrResult:
rating_class = self.rating_class()
pure = self.pure()
far = self.far()
lost = self.lost()
score = self.score()
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()
return DeviceOcrResult(
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,
clear_status=clear_status,
partner_id=partner_id,
partner_id_possibility=1 - partner_id_distance / hash_len,
)

View File

@ -1 +0,0 @@
from .auto import *

View File

@ -1,3 +0,0 @@
from .common import DeviceAutoRoiSizes
from .t1 import DeviceAutoRoiSizesT1
from .t2 import DeviceAutoRoiSizesT2

View File

@ -1,7 +0,0 @@
from ..common import DeviceRoiSizes
class DeviceAutoRoiSizes(DeviceRoiSizes):
def __init__(self, w: int, h: int):
self.w = w
self.h = h

View File

@ -1,123 +0,0 @@
from .common import DeviceAutoRoiSizes
class DeviceAutoRoiSizesT1(DeviceAutoRoiSizes):
@property
def factor(self):
return (
((self.w / 16) * 9) / 720 if (self.w / self.h) < (16 / 9) else self.h / 720
)
@property
def w_mid(self):
return self.w / 2
@property
def h_mid(self):
return self.h / 2
@property
def top_bar(self):
return (0, 0, self.w, 50 * self.factor)
@property
def layout_area_h_mid(self):
return self.h / 2 + self.top_bar[3]
@property
def pfl_left_from_w_mid(self):
return 5 * self.factor
@property
def pfl_x(self):
return self.w_mid + self.pfl_left_from_w_mid
@property
def pfl_w(self):
return 76 * self.factor
@property
def pfl_h(self):
return 26 * self.factor
@property
def pure(self):
return (
self.pfl_x,
self.layout_area_h_mid + 110 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def far(self):
return (
self.pfl_x,
self.pure[1] + self.pure[3] + 12 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def lost(self):
return (
self.pfl_x,
self.far[1] + self.far[3] + 10 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def score(self):
w = 280 * self.factor
h = 45 * self.factor
return (
self.w_mid - w / 2,
self.layout_area_h_mid - 75 * self.factor - h,
w,
h,
)
@property
def rating_class(self):
return (
self.w_mid - 610 * self.factor,
self.layout_area_h_mid - 180 * self.factor,
265 * self.factor,
35 * self.factor,
)
@property
def max_recall(self):
return (
self.w_mid - 465 * self.factor,
self.layout_area_h_mid - 215 * self.factor,
150 * self.factor,
35 * self.factor,
)
@property
def jacket(self):
return (
self.w_mid - 610 * self.factor,
self.layout_area_h_mid - 143 * self.factor,
375 * self.factor,
375 * self.factor,
)
@property
def clear_status(self):
w = 550 * self.factor
h = 60 * self.factor
return (
self.w_mid - w / 2,
self.layout_area_h_mid - 155 * self.factor - h,
w,
h,
)
@property
def partner_icon(self):
w = 90 * self.factor
h = 75 * self.factor
return (self.w_mid - w / 2, 0, w, h)

View File

@ -1,125 +0,0 @@
from .common import DeviceAutoRoiSizes
class DeviceAutoRoiSizesT2(DeviceAutoRoiSizes):
@property
def factor(self):
return (
((self.w / 16) * 9) / 1080
if (self.w / self.h) < (16 / 9)
else self.h / 1080
)
@property
def w_mid(self):
return self.w / 2
@property
def h_mid(self):
return self.h / 2
@property
def top_bar(self):
return (0, 0, self.w, 75 * self.factor)
@property
def layout_area_h_mid(self):
return self.h / 2 + self.top_bar[3]
@property
def pfl_mid_from_w_mid(self):
return 60 * self.factor
@property
def pfl_x(self):
return self.w_mid + 10 * self.factor
@property
def pfl_w(self):
return 100 * self.factor
@property
def pfl_h(self):
return 24 * self.factor
@property
def pure(self):
return (
self.pfl_x,
self.layout_area_h_mid + 175 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def far(self):
return (
self.pfl_x,
self.pure[1] + self.pure[3] + 30 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def lost(self):
return (
self.pfl_x,
self.far[1] + self.far[3] + 35 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def score(self):
w = 420 * self.factor
h = 70 * self.factor
return (
self.w_mid - w / 2,
self.layout_area_h_mid - 110 * self.factor - h,
w,
h,
)
@property
def rating_class(self):
return (
max(0, self.w_mid - 965 * self.factor),
self.layout_area_h_mid - 330 * self.factor,
350 * self.factor,
110 * self.factor,
)
@property
def max_recall(self):
return (
self.w_mid - 625 * self.factor,
self.layout_area_h_mid - 275 * self.factor,
150 * self.factor,
50 * self.factor,
)
@property
def jacket(self):
return (
self.w_mid - 915 * self.factor,
self.layout_area_h_mid - 215 * self.factor,
565 * self.factor,
565 * self.factor,
)
@property
def clear_status(self):
w = 825 * self.factor
h = 90 * self.factor
return (
self.w_mid - w / 2,
self.layout_area_h_mid - 235 * self.factor - h,
w,
h,
)
@property
def partner_icon(self):
w = 135 * self.factor
h = 110 * self.factor
return (self.w_mid - w / 2, 0, w, h)

View File

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

View File

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

View File

@ -1,3 +0,0 @@
from .common import DeviceAutoRoiMasker
from .t1 import DeviceAutoRoiMaskerT1
from .t2 import DeviceAutoRoiMaskerT2

View File

@ -1,5 +0,0 @@
from ..common import DeviceRoiMasker
class DeviceAutoRoiMasker(DeviceRoiMasker):
...

View File

@ -1,123 +0,0 @@
import cv2
import numpy as np
from .common import DeviceAutoRoiMasker
GRAY_BGR_MIN = np.array([50] * 3, np.uint8)
GRAY_BGR_MAX = np.array([160] * 3, np.uint8)
WHITE_HSV_MIN = np.array([0, 0, 240], np.uint8)
WHITE_HSV_MAX = np.array([179, 10, 255], np.uint8)
PST_HSV_MIN = np.array([100, 50, 80], np.uint8)
PST_HSV_MAX = np.array([100, 255, 255], np.uint8)
PRS_HSV_MIN = np.array([43, 40, 75], np.uint8)
PRS_HSV_MAX = np.array([50, 155, 190], np.uint8)
FTR_HSV_MIN = np.array([149, 30, 0], np.uint8)
FTR_HSV_MAX = np.array([155, 181, 150], np.uint8)
BYD_HSV_MIN = np.array([170, 50, 50], np.uint8)
BYD_HSV_MAX = np.array([179, 210, 198], np.uint8)
TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8)
TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8)
TRACK_COMPLETE_HSV_MIN = np.array([140, 0, 50], np.uint8)
TRACK_COMPLETE_HSV_MAX = np.array([145, 50, 130], np.uint8)
FULL_RECALL_HSV_MIN = np.array([140, 60, 80], np.uint8)
FULL_RECALL_HSV_MAX = np.array([150, 130, 145], np.uint8)
PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8)
PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8)
class DeviceAutoRoiMaskerT1(DeviceAutoRoiMasker):
@classmethod
def gray(cls, roi_bgr: cv2.Mat) -> cv2.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)
img_bgr = cv2.erode(img_bgr, cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2)))
img_bgr = cv2.dilate(img_bgr, cv2.getStructuringElement(cv2.MORPH_RECT, (1, 1)))
return cv2.inRange(img_bgr, GRAY_BGR_MIN, GRAY_BGR_MAX)
@classmethod
def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cls.gray(roi_bgr)
@classmethod
def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cls.gray(roi_bgr)
@classmethod
def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cls.gray(roi_bgr)
@classmethod
def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), WHITE_HSV_MIN, WHITE_HSV_MAX
)
@classmethod
def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PST_HSV_MIN, PST_HSV_MAX
)
@classmethod
def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PRS_HSV_MIN, PRS_HSV_MAX
)
@classmethod
def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), FTR_HSV_MIN, FTR_HSV_MAX
)
@classmethod
def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), BYD_HSV_MIN, BYD_HSV_MAX
)
@classmethod
def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cls.gray(roi_bgr)
@classmethod
def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
TRACK_LOST_HSV_MIN,
TRACK_LOST_HSV_MAX,
)
@classmethod
def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
TRACK_COMPLETE_HSV_MIN,
TRACK_COMPLETE_HSV_MAX,
)
@classmethod
def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
FULL_RECALL_HSV_MIN,
FULL_RECALL_HSV_MAX,
)
@classmethod
def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
PURE_MEMORY_HSV_MIN,
PURE_MEMORY_HSV_MAX,
)

View File

@ -1,128 +0,0 @@
import cv2
import numpy as np
from .common import DeviceAutoRoiMasker
PFL_HSV_MIN = np.array([0, 0, 248], np.uint8)
PFL_HSV_MAX = np.array([179, 10, 255], np.uint8)
WHITE_HSV_MIN = np.array([0, 0, 240], np.uint8)
WHITE_HSV_MAX = np.array([179, 10, 255], np.uint8)
PST_HSV_MIN = np.array([100, 50, 80], np.uint8)
PST_HSV_MAX = np.array([100, 255, 255], np.uint8)
PRS_HSV_MIN = np.array([43, 40, 75], np.uint8)
PRS_HSV_MAX = np.array([50, 155, 190], np.uint8)
FTR_HSV_MIN = np.array([149, 30, 0], np.uint8)
FTR_HSV_MAX = np.array([155, 181, 150], np.uint8)
BYD_HSV_MIN = np.array([170, 50, 50], np.uint8)
BYD_HSV_MAX = np.array([179, 210, 198], np.uint8)
MAX_RECALL_HSV_MIN = np.array([125, 0, 0], np.uint8)
MAX_RECALL_HSV_MAX = np.array([130, 100, 150], np.uint8)
TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8)
TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8)
TRACK_COMPLETE_HSV_MIN = np.array([140, 0, 50], np.uint8)
TRACK_COMPLETE_HSV_MAX = np.array([145, 50, 130], np.uint8)
FULL_RECALL_HSV_MIN = np.array([140, 60, 80], np.uint8)
FULL_RECALL_HSV_MAX = np.array([150, 130, 145], np.uint8)
PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8)
PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8)
class DeviceAutoRoiMaskerT2(DeviceAutoRoiMasker):
@classmethod
def pfl(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PFL_HSV_MIN, PFL_HSV_MAX
)
@classmethod
def pure(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cls.pfl(roi_bgr)
@classmethod
def far(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cls.pfl(roi_bgr)
@classmethod
def lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cls.pfl(roi_bgr)
@classmethod
def score(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), WHITE_HSV_MIN, WHITE_HSV_MAX
)
@classmethod
def rating_class_pst(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PST_HSV_MIN, PST_HSV_MAX
)
@classmethod
def rating_class_prs(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), PRS_HSV_MIN, PRS_HSV_MAX
)
@classmethod
def rating_class_ftr(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), FTR_HSV_MIN, FTR_HSV_MAX
)
@classmethod
def rating_class_byd(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV), BYD_HSV_MIN, BYD_HSV_MAX
)
@classmethod
def max_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
MAX_RECALL_HSV_MIN,
MAX_RECALL_HSV_MAX,
)
@classmethod
def clear_status_track_lost(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
TRACK_LOST_HSV_MIN,
TRACK_LOST_HSV_MAX,
)
@classmethod
def clear_status_track_complete(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
TRACK_COMPLETE_HSV_MIN,
TRACK_COMPLETE_HSV_MAX,
)
@classmethod
def clear_status_full_recall(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
FULL_RECALL_HSV_MIN,
FULL_RECALL_HSV_MAX,
)
@classmethod
def clear_status_pure_memory(cls, roi_bgr: cv2.Mat) -> cv2.Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
PURE_MEMORY_HSV_MIN,
PURE_MEMORY_HSV_MAX,
)

View File

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

View File

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

View File

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

View File

@ -0,0 +1,255 @@
from .common import DeviceRois
__all__ = ["DeviceRoisAuto", "DeviceRoisAutoT1", "DeviceRoisAutoT2"]
class DeviceRoisAuto(DeviceRois):
def __init__(self, w: int, h: int):
self.w = w
self.h = h
class DeviceRoisAutoT1(DeviceRoisAuto):
@property
def factor(self):
return (
((self.w / 16) * 9) / 720 if (self.w / self.h) < (16 / 9) else self.h / 720
)
@property
def w_mid(self):
return self.w / 2
@property
def h_mid(self):
return self.h / 2
@property
def top_bar(self):
return (0, 0, self.w, 50 * self.factor)
@property
def layout_area_h_mid(self):
return self.h / 2 + self.top_bar[3]
@property
def pfl_left_from_w_mid(self):
return 5 * self.factor
@property
def pfl_x(self):
return self.w_mid + self.pfl_left_from_w_mid
@property
def pfl_w(self):
return 76 * self.factor
@property
def pfl_h(self):
return 26 * self.factor
@property
def pure(self):
return (
self.pfl_x,
self.layout_area_h_mid + 110 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def far(self):
return (
self.pfl_x,
self.pure[1] + self.pure[3] + 12 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def lost(self):
return (
self.pfl_x,
self.far[1] + self.far[3] + 10 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def score(self):
w = 280 * self.factor
h = 45 * self.factor
return (
self.w_mid - w / 2,
self.layout_area_h_mid - 75 * self.factor - h,
w,
h,
)
@property
def rating_class(self):
return (
self.w_mid - 610 * self.factor,
self.layout_area_h_mid - 180 * self.factor,
265 * self.factor,
35 * self.factor,
)
@property
def max_recall(self):
return (
self.w_mid - 465 * self.factor,
self.layout_area_h_mid - 215 * self.factor,
150 * self.factor,
35 * self.factor,
)
@property
def jacket(self):
return (
self.w_mid - 610 * self.factor,
self.layout_area_h_mid - 143 * self.factor,
375 * self.factor,
375 * self.factor,
)
@property
def clear_status(self):
w = 550 * self.factor
h = 60 * self.factor
return (
self.w_mid - w / 2,
self.layout_area_h_mid - 155 * self.factor - h,
w,
h,
)
@property
def partner_icon(self):
w = 90 * self.factor
h = 75 * self.factor
return (self.w_mid - w / 2, 0, w, h)
class DeviceRoisAutoT2(DeviceRoisAuto):
@property
def factor(self):
return (
((self.w / 16) * 9) / 1080
if (self.w / self.h) < (16 / 9)
else self.h / 1080
)
@property
def w_mid(self):
return self.w / 2
@property
def h_mid(self):
return self.h / 2
@property
def top_bar(self):
return (0, 0, self.w, 75 * self.factor)
@property
def layout_area_h_mid(self):
return self.h / 2 + self.top_bar[3]
@property
def pfl_mid_from_w_mid(self):
return 60 * self.factor
@property
def pfl_x(self):
return self.w_mid + 10 * self.factor
@property
def pfl_w(self):
return 100 * self.factor
@property
def pfl_h(self):
return 24 * self.factor
@property
def pure(self):
return (
self.pfl_x,
self.layout_area_h_mid + 175 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def far(self):
return (
self.pfl_x,
self.pure[1] + self.pure[3] + 30 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def lost(self):
return (
self.pfl_x,
self.far[1] + self.far[3] + 35 * self.factor,
self.pfl_w,
self.pfl_h,
)
@property
def score(self):
w = 420 * self.factor
h = 70 * self.factor
return (
self.w_mid - w / 2,
self.layout_area_h_mid - 110 * self.factor - h,
w,
h,
)
@property
def rating_class(self):
return (
max(0, self.w_mid - 965 * self.factor),
self.layout_area_h_mid - 330 * self.factor,
350 * self.factor,
110 * self.factor,
)
@property
def max_recall(self):
return (
self.w_mid - 625 * self.factor,
self.layout_area_h_mid - 275 * self.factor,
150 * self.factor,
50 * self.factor,
)
@property
def jacket(self):
return (
self.w_mid - 915 * self.factor,
self.layout_area_h_mid - 215 * self.factor,
565 * self.factor,
565 * self.factor,
)
@property
def clear_status(self):
w = 825 * self.factor
h = 90 * self.factor
return (
self.w_mid - w / 2,
self.layout_area_h_mid - 235 * self.factor - h,
w,
h,
)
@property
def partner_icon(self):
w = 135 * self.factor
h = 110 * self.factor
return (self.w_mid - w / 2, 0, w, h)

View File

@ -3,7 +3,7 @@ from typing import Tuple
Rect = Tuple[int, int, int, int]
class DeviceRoiSizes:
class DeviceRois:
pure: Rect
far: Rect
lost: Rect

View File

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

View File

@ -1,13 +1,12 @@
import cv2
from ....crop import crop_xywh
from ..definitions.common import DeviceRoiSizes
from ....types import Mat
from ..definition.common import DeviceRois
class DeviceRoiExtractor:
def __init__(self, img: cv2.Mat, sizes: DeviceRoiSizes):
class DeviceRoisExtractor:
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

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

View File

@ -0,0 +1,255 @@
import cv2
import numpy as np
from ....types import Mat
from .common import DeviceRoisMasker
class DeviceRoisMaskerAuto(DeviceRoisMasker):
...
class DeviceRoisMaskerAutoT1(DeviceRoisMaskerAuto):
GRAY_BGR_MIN = np.array([50] * 3, np.uint8)
GRAY_BGR_MAX = np.array([160] * 3, np.uint8)
WHITE_HSV_MIN = np.array([0, 0, 240], np.uint8)
WHITE_HSV_MAX = np.array([179, 10, 255], np.uint8)
PST_HSV_MIN = np.array([100, 50, 80], np.uint8)
PST_HSV_MAX = np.array([100, 255, 255], np.uint8)
PRS_HSV_MIN = np.array([43, 40, 75], np.uint8)
PRS_HSV_MAX = np.array([50, 155, 190], np.uint8)
FTR_HSV_MIN = np.array([149, 30, 0], np.uint8)
FTR_HSV_MAX = np.array([155, 181, 150], np.uint8)
BYD_HSV_MIN = np.array([170, 50, 50], np.uint8)
BYD_HSV_MAX = np.array([179, 210, 198], np.uint8)
TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8)
TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8)
TRACK_COMPLETE_HSV_MIN = np.array([140, 0, 50], np.uint8)
TRACK_COMPLETE_HSV_MAX = np.array([145, 50, 130], np.uint8)
FULL_RECALL_HSV_MIN = np.array([140, 60, 80], np.uint8)
FULL_RECALL_HSV_MAX = np.array([150, 130, 145], np.uint8)
PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8)
PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8)
@classmethod
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)
img_bgr = cv2.erode(img_bgr, cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2)))
img_bgr = cv2.dilate(img_bgr, cv2.getStructuringElement(cv2.MORPH_RECT, (1, 1)))
return cv2.inRange(img_bgr, cls.GRAY_BGR_MIN, cls.GRAY_BGR_MAX)
@classmethod
def pure(cls, roi_bgr: Mat) -> Mat:
return cls.gray(roi_bgr)
@classmethod
def far(cls, roi_bgr: Mat) -> Mat:
return cls.gray(roi_bgr)
@classmethod
def lost(cls, roi_bgr: Mat) -> Mat:
return cls.gray(roi_bgr)
@classmethod
def score(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.WHITE_HSV_MIN,
cls.WHITE_HSV_MAX,
)
@classmethod
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: 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: 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: 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: Mat) -> Mat:
return cls.gray(roi_bgr)
@classmethod
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,
cls.TRACK_LOST_HSV_MAX,
)
@classmethod
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,
cls.TRACK_COMPLETE_HSV_MAX,
)
@classmethod
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,
cls.FULL_RECALL_HSV_MAX,
)
@classmethod
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,
cls.PURE_MEMORY_HSV_MAX,
)
class DeviceRoisMaskerAutoT2(DeviceRoisMaskerAuto):
PFL_HSV_MIN = np.array([0, 0, 248], np.uint8)
PFL_HSV_MAX = np.array([179, 10, 255], np.uint8)
WHITE_HSV_MIN = np.array([0, 0, 240], np.uint8)
WHITE_HSV_MAX = np.array([179, 10, 255], np.uint8)
PST_HSV_MIN = np.array([100, 50, 80], np.uint8)
PST_HSV_MAX = np.array([100, 255, 255], np.uint8)
PRS_HSV_MIN = np.array([43, 40, 75], np.uint8)
PRS_HSV_MAX = np.array([50, 155, 190], np.uint8)
FTR_HSV_MIN = np.array([149, 30, 0], np.uint8)
FTR_HSV_MAX = np.array([155, 181, 150], np.uint8)
BYD_HSV_MIN = np.array([170, 50, 50], np.uint8)
BYD_HSV_MAX = np.array([179, 210, 198], np.uint8)
MAX_RECALL_HSV_MIN = np.array([125, 0, 0], np.uint8)
MAX_RECALL_HSV_MAX = np.array([145, 100, 150], np.uint8)
TRACK_LOST_HSV_MIN = np.array([170, 75, 90], np.uint8)
TRACK_LOST_HSV_MAX = np.array([175, 170, 160], np.uint8)
TRACK_COMPLETE_HSV_MIN = np.array([140, 0, 50], np.uint8)
TRACK_COMPLETE_HSV_MAX = np.array([145, 50, 130], np.uint8)
FULL_RECALL_HSV_MIN = np.array([140, 60, 80], np.uint8)
FULL_RECALL_HSV_MAX = np.array([150, 130, 145], np.uint8)
PURE_MEMORY_HSV_MIN = np.array([90, 70, 80], np.uint8)
PURE_MEMORY_HSV_MAX = np.array([110, 200, 175], np.uint8)
@classmethod
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: Mat) -> Mat:
return cls.pfl(roi_bgr)
@classmethod
def far(cls, roi_bgr: Mat) -> Mat:
return cls.pfl(roi_bgr)
@classmethod
def lost(cls, roi_bgr: Mat) -> Mat:
return cls.pfl(roi_bgr)
@classmethod
def score(cls, roi_bgr: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.WHITE_HSV_MIN,
cls.WHITE_HSV_MAX,
)
@classmethod
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: 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: 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: 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: Mat) -> Mat:
return cv2.inRange(
cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV),
cls.MAX_RECALL_HSV_MIN,
cls.MAX_RECALL_HSV_MAX,
)
@classmethod
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,
cls.TRACK_LOST_HSV_MAX,
)
@classmethod
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,
cls.TRACK_COMPLETE_HSV_MAX,
)
@classmethod
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,
cls.FULL_RECALL_HSV_MAX,
)
@classmethod
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,
cls.PURE_MEMORY_HSV_MAX,
)

View File

@ -0,0 +1,55 @@
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 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,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

@ -1,8 +1,34 @@
import sqlite3
from typing import List, Union
import imagehash
import cv2
import numpy as np
from PIL import Image
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:
@ -11,7 +37,7 @@ def hamming_distance_sql_function(user_input, db_entry) -> int:
)
class ImagePHashDatabase:
class ImagePhashDatabase:
def __init__(self, db_path: str):
with sqlite3.connect(db_path) as conn:
self.hash_size = int(
@ -30,36 +56,63 @@ class ImagePHashDatabase:
).fetchone()[0]
)
# self.conn.create_function(
# "HAMMING_DISTANCE",
# 2,
# hamming_distance_sql_function,
# deterministic=True,
# )
self.ids = [i[0] for i in conn.execute("SELECT id FROM hashes").fetchall()]
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.hashes_slice_size = round(len(self.hashes_byte[0]) * 0.25)
self.hashes_head = [h[: self.hashes_slice_size] for h in self.hashes]
self.hashes_tail = [h[-self.hashes_slice_size :] for h in self.hashes]
def lookup_hash(self, image_hash: imagehash.ImageHash, *, limit: int = 5):
image_hash = image_hash.hash.flatten()
# image_hash_head = image_hash[: self.hashes_slice_size]
# image_hash_tail = image_hash[-self.hashes_slice_size :]
# head_xor_results = [image_hash_head ^ h for h in self.hashes]
# tail_xor_results = [image_hash_head ^ h for h in self.hashes]
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, pil_image: Image.Image):
image_hash = imagehash.phash(
pil_image, hash_size=self.hash_size, highfreq_factor=self.highfreq_factor
)
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

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