19 Commits

33 changed files with 844 additions and 857 deletions

View File

@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
[project] [project]
name = "arcaea-offline-ocr" name = "arcaea-offline-ocr"
version = "0.1.0" version = "0.0.95"
authors = [{ name = "283375", email = "log_283375@163.com" }] authors = [{ name = "283375", email = "log_283375@163.com" }]
description = "Extract your Arcaea play result from screenshot." description = "Extract your Arcaea play result from screenshot."
readme = "README.md" readme = "README.md"
requires-python = ">=3.8" 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 = [ classifiers = [
"Development Status :: 3 - Alpha", "Development Status :: 3 - Alpha",
"Programming Language :: Python :: 3", "Programming Language :: Python :: 3",

View File

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

View File

@ -3,12 +3,16 @@ from typing import List, Optional, Tuple
import cv2 import cv2
import numpy as np import numpy as np
from PIL import Image
from ....crop import crop_xywh from ....crop import crop_xywh
from ....ocr import FixRects, ocr_digits_by_contour_knn, preprocess_hog from ....ocr import (
from ....phash_db import ImagePHashDatabase FixRects,
from ....types import Mat, cv2_ml_KNearest 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 ....utils import construct_int_xywh_rect
from ...shared import B30OcrResultItem from ...shared import B30OcrResultItem
from .colors import * from .colors import *
@ -18,9 +22,9 @@ from .rois import ChieriBotV4Rois
class ChieriBotV4Ocr: class ChieriBotV4Ocr:
def __init__( def __init__(
self, self,
score_knn: cv2_ml_KNearest, score_knn: cv2.ml.KNearest,
pfl_knn: cv2_ml_KNearest, pfl_knn: cv2.ml.KNearest,
phash_db: ImagePHashDatabase, phash_db: ImagePhashDatabase,
factor: Optional[float] = 1.0, factor: Optional[float] = 1.0,
): ):
self.__score_knn = score_knn self.__score_knn = score_knn
@ -33,7 +37,7 @@ class ChieriBotV4Ocr:
return self.__score_knn return self.__score_knn
@score_knn.setter @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 self.__score_knn = knn_digits_model
@property @property
@ -41,7 +45,7 @@ class ChieriBotV4Ocr:
return self.__pfl_knn return self.__pfl_knn
@pfl_knn.setter @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 self.__pfl_knn = knn_digits_model
@property @property
@ -49,7 +53,7 @@ class ChieriBotV4Ocr:
return self.__phash_db return self.__phash_db
@phash_db.setter @phash_db.setter
def phash_db(self, phash_db: ImagePHashDatabase): def phash_db(self, phash_db: ImagePhashDatabase):
self.__phash_db = phash_db self.__phash_db = phash_db
@property @property
@ -84,14 +88,6 @@ class ChieriBotV4Ocr:
else: else:
return max(enumerate(rating_class_results), key=lambda i: i[1])[0] + 1 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): def ocr_component_song_id(self, component_bgr: Mat):
jacket_rect = construct_int_xywh_rect( jacket_rect = construct_int_xywh_rect(
self.rois.component_rois.jacket_rect, floor self.rois.component_rois.jacket_rect, floor
@ -99,20 +95,7 @@ class ChieriBotV4Ocr:
jacket_roi = cv2.cvtColor( jacket_roi = cv2.cvtColor(
crop_xywh(component_bgr, jacket_rect), cv2.COLOR_BGR2GRAY crop_xywh(component_bgr, jacket_rect), cv2.COLOR_BGR2GRAY
) )
return self.phash_db.lookup_image(Image.fromarray(jacket_roi))[0] return self.phash_db.lookup_jacket(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
def ocr_component_score_knn(self, component_bgr: Mat) -> int: def ocr_component_score_knn(self, component_bgr: Mat) -> int:
# sourcery skip: inline-immediately-returned-variable # sourcery skip: inline-immediately-returned-variable
@ -222,7 +205,7 @@ class ChieriBotV4Ocr:
digits = [] digits = []
for digit_rect in digit_rects: for digit_rect in digit_rects:
digit = crop_xywh(roi, digit_rect) digit = crop_xywh(roi, digit_rect)
digit = cv2.resize(digit, (20, 20)) digit = resize_fill_square(digit, 20)
digits.append(digit) digits.append(digit)
samples = preprocess_hog(digits) samples = preprocess_hog(digits)
@ -233,15 +216,6 @@ class ChieriBotV4Ocr:
except Exception: except Exception:
return (None, None, None) 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: def ocr_component(self, component_bgr: Mat) -> B30OcrResultItem:
component_blur = cv2.GaussianBlur(component_bgr, (5, 5), 0) component_blur = cv2.GaussianBlur(component_bgr, (5, 5), 0)
rating_class = self.ocr_component_rating_class(component_blur) 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 from typing import Tuple
import cv2
import numpy as np import numpy as np
from .types import Mat 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]): def crop_xywh(mat: Mat, rect: Tuple[int, int, int, int]):
@ -13,62 +14,16 @@ def crop_xywh(mat: Mat, rect: Tuple[int, int, int, int]):
return mat[y : y + h, x : x + w] return mat[y : y + h, x : x + w]
def is_black_edge(list_of_pixels: Mat, black_pixel: Mat, ratio: float = 0.6): class CropBlackEdges:
pixels = list_of_pixels.reshape([-1, 3]) @staticmethod
return np.count_nonzero(np.all(pixels < black_pixel, axis=1)) > floor( def is_black_edge(__img_gray_slice: Mat, black_pixel: int, ratio: float = 0.6):
len(pixels) * ratio pixels_compared = __img_gray_slice < black_pixel
return np.count_nonzero(pixels_compared) > math.floor(
__img_gray_slice.size * ratio
) )
@classmethod
def crop_black_edges(img_bgr: Mat, black_threshold: int = 50): def get_crop_rect(cls, img_gray: Mat, black_threshold: int = 25):
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 = cropped[:, i]
if not is_black_edge(column, black_pixel):
break
left += 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 = cropped[i]
if not is_black_edge(row, black_pixel):
break
top += 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
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] height, width = img_gray.shape[:2]
left = 0 left = 0
right = width right = width
@ -77,28 +32,35 @@ def crop_black_edges_grayscale(
for i in range(width): for i in range(width):
column = img_gray[:, i] column = img_gray[:, i]
if not is_black_edge_grayscale(column, black_threshold): if not cls.is_black_edge(column, black_threshold):
break break
left += 1 left += 1
for i in sorted(range(width), reverse=True): for i in sorted(range(width), reverse=True):
column = img_gray[:, i] column = img_gray[:, i]
if i <= left + 1 or not is_black_edge_grayscale(column, black_threshold): if i <= left + 1 or not cls.is_black_edge(column, black_threshold):
break break
right -= 1 right -= 1
for i in range(height): for i in range(height):
row = img_gray[i] row = img_gray[i]
if not is_black_edge_grayscale(row, black_threshold): if not cls.is_black_edge(row, black_threshold):
break break
top += 1 top += 1
for i in sorted(range(height), reverse=True): for i in sorted(range(height), reverse=True):
row = img_gray[i] row = img_gray[i]
if i <= top + 1 or not is_black_edge_grayscale(row, black_threshold): if i <= top + 1 or not cls.is_black_edge(row, black_threshold):
break break
bottom -= 1 bottom -= 1
assert right > left, "cropped width > 0" assert right > left, "cropped width < 0"
assert bottom > top, "cropped height > 0" assert bottom > top, "cropped height < 0"
return (left, top, right - left, bottom - top) 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 far: int
lost: int lost: int
score: int score: int
max_recall: int max_recall: Optional[int] = None
song_id: Optional[str] = None song_id: Optional[str] = None
title: Optional[str] = None song_id_possibility: Optional[float] = None
clear_type: Optional[str] = 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 cv2
import numpy as np import numpy as np
from PIL import Image
from ..crop import crop_xywh from ..crop import crop_xywh
from ..ocr import ( from ..ocr import (
@ -10,25 +9,27 @@ from ..ocr import (
preprocess_hog, preprocess_hog,
resize_fill_square, resize_fill_square,
) )
from ..phash_db import ImagePHashDatabase from ..phash_db import ImagePhashDatabase
from .roi.extractor import DeviceRoiExtractor from ..types import Mat
from .roi.masker import DeviceRoiMasker from .common import DeviceOcrResult
from .rois.extractor import DeviceRoisExtractor
from .rois.masker import DeviceRoisMasker
class DeviceOcr: class DeviceOcr:
def __init__( def __init__(
self, self,
extractor: DeviceRoiExtractor, extractor: DeviceRoisExtractor,
masker: DeviceRoiMasker, masker: DeviceRoisMasker,
knn_model: cv2.ml.KNearest, knn_model: cv2.ml.KNearest,
phash_db: ImagePHashDatabase, phash_db: ImagePhashDatabase,
): ):
self.extractor = extractor self.extractor = extractor
self.masker = masker self.masker = masker
self.knn_model = knn_model self.knn_model = knn_model
self.phash_db = phash_db 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( contours, _ = cv2.findContours(
roi_gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE roi_gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
) )
@ -47,8 +48,7 @@ class DeviceOcr:
continue continue
roi_ocr = cv2.fillPoly(roi_ocr, [contour], [0]) roi_ocr = cv2.fillPoly(roi_ocr, [contour], [0])
digit_rois = [ digit_rois = [
resize_fill_square(crop_xywh(roi_ocr, r), 20) resize_fill_square(crop_xywh(roi_ocr, r), 20) for r in filtered_rects
for r in sorted(filtered_rects, key=lambda r: r[0])
] ]
samples = preprocess_hog(digit_rois) samples = preprocess_hog(digit_rois)
@ -97,5 +97,64 @@ class DeviceOcr:
] ]
return max(enumerate(results), key=lambda i: np.count_nonzero(i[1]))[0] 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): 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] Rect = Tuple[int, int, int, int]
class DeviceRoiSizes: class DeviceRois:
pure: Rect pure: Rect
far: Rect far: Rect
lost: 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 ....crop import crop_xywh
from ..definitions.common import DeviceRoiSizes from ....types import Mat
from ..definition.common import DeviceRois
class DeviceRoiExtractor: class DeviceRoisExtractor:
def __init__(self, img: cv2.Mat, sizes: DeviceRoiSizes): def __init__(self, img: Mat, rois: DeviceRois):
self.img = img self.img = img
self.sizes = sizes self.sizes = rois
def __construct_int_rect(self, rect): def __construct_int_rect(self, rect):
return tuple(round(r) for r in 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 import math
from copy import deepcopy
from typing import Optional, Sequence, Tuple from typing import Optional, Sequence, Tuple
import cv2 import cv2
import numpy as np import numpy as np
from numpy.linalg import norm
from .crop import crop_xywh from .crop import crop_xywh
from .types import Mat, cv2_ml_KNearest from .types import Mat
__all__ = [ __all__ = [
"FixRects", "FixRects",
@ -64,8 +62,7 @@ class FixRects:
new_h = new_bottom - new_y new_h = new_bottom - new_y
new_rects.append((new_x, new_y, new_w, new_h)) new_rects.append((new_x, new_y, new_w, new_h))
return_rects = deepcopy(rects) return_rects = [r for r in rects if r not in consumed_rects]
return_rects = [r for r in return_rects if r not in consumed_rects]
return_rects.extend(new_rects) return_rects.extend(new_rects)
return return_rects return return_rects
@ -80,8 +77,9 @@ class FixRects:
new_rects = [] new_rects = []
for rect in rects: for rect in rects:
rx, ry, rw, rh = rect rx, ry, rw, rh = rect
if rw / rh > rect_wh_ratio: if rw / rh <= rect_wh_ratio:
# consider this is a connected contour continue
connected_rects.append(rect) connected_rects.append(rect)
# find the thinnest part # find the thinnest part
@ -94,19 +92,19 @@ class FixRects:
for i in range(img_cropped.shape[1]): for i in range(img_cropped.shape[1]):
col = img_cropped[:, i] col = img_cropped[:, i]
white_pixels[rx + border_ignore + i] = np.count_nonzero(col > 200) 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) least_white_pixels = min(v for v in white_pixels.values() if v > 0)
x_values = [ x_values = [
x x for x, pixel in white_pixels.items() if pixel == least_white_pixels
for x, pixel in white_pixels.items()
if pixel == least_white_pixels
] ]
# select only middle values # select only middle values
x_mean = np.mean(x_values) x_mean = np.mean(x_values)
x_std = np.std(x_values) x_std = np.std(x_values)
x_values = [ x_values = [
x x for x in x_values if x_mean - x_std * 1.5 <= x <= x_mean + x_std * 1.5
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)) x_mid = round(np.median(x_values))
@ -115,7 +113,6 @@ class FixRects:
[(rx, ry, x_mid - rx, rh), (x_mid, ry, rx + rw - x_mid, rh)] [(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 = [r for r in rects if r not in connected_rects]
return_rects.extend(new_rects) return_rects.extend(new_rects)
return return_rects return return_rects
@ -144,33 +141,16 @@ def resize_fill_square(img: Mat, target: int = 20):
def preprocess_hog(digit_rois): 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 = [] samples = []
for digit in digit_rois: for digit in digit_rois:
gx = cv2.Sobel(digit, cv2.CV_32F, 1, 0) hog = cv2.HOGDescriptor((20, 20), (10, 10), (5, 5), (10, 10), 9)
gy = cv2.Sobel(digit, cv2.CV_32F, 0, 1) hist = hog.compute(digit)
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
samples.append(hist) samples.append(hist)
return np.float32(samples) 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) _, results, _, _ = knn_model.findNearest(__samples, k)
result_list = [int(r) for r in results.ravel()] result_list = [int(r) for r in results.ravel()]
result_str = "".join(str(r) for r in result_list if r > -1) 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( def ocr_digits_by_contour_knn(
__roi_gray: Mat, __roi_gray: Mat,
knn_model: cv2_ml_KNearest, knn_model: cv2.ml.KNearest,
*, *,
k=4, k=4,
size: int = 20, size: int = 20,

View File

@ -1,8 +1,34 @@
import sqlite3 import sqlite3
from typing import List, Union
import imagehash import cv2
import numpy as np 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: 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): def __init__(self, db_path: str):
with sqlite3.connect(db_path) as conn: with sqlite3.connect(db_path) as conn:
self.hash_size = int( self.hash_size = int(
@ -30,36 +56,63 @@ class ImagePHashDatabase:
).fetchone()[0] ).fetchone()[0]
) )
# self.conn.create_function( self.ids: List[str] = [
# "HAMMING_DISTANCE", i[0] for i in conn.execute("SELECT id FROM hashes").fetchall()
# 2, ]
# hamming_distance_sql_function,
# deterministic=True,
# )
self.ids = [i[0] for i in conn.execute("SELECT id FROM hashes").fetchall()]
self.hashes_byte = [ self.hashes_byte = [
i[0] for i in conn.execute("SELECT hash FROM hashes").fetchall() 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 = [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): self.jacket_ids: List[str] = []
image_hash = image_hash.hash.flatten() self.jacket_hashes = []
# image_hash_head = image_hash[: self.hashes_slice_size] self.partner_icon_ids: List[str] = []
# image_hash_tail = image_hash[-self.hashes_slice_size :] self.partner_icon_hashes = []
# head_xor_results = [image_hash_head ^ h for h in self.hashes]
# tail_xor_results = [image_hash_head ^ h for h in self.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 = [ xor_results = [
(id, np.count_nonzero(image_hash ^ h)) (id, np.count_nonzero(image_hash ^ h))
for id, h in zip(self.ids, self.hashes) for id, h in zip(self.ids, self.hashes)
] ]
return sorted(xor_results, key=lambda r: r[1])[:limit] return sorted(xor_results, key=lambda r: r[1])[:limit]
def lookup_image(self, pil_image: Image.Image): def lookup_image(self, img_gray: Mat):
image_hash = imagehash.phash( image_hash = self.calculate_phash(img_gray)
pil_image, hash_size=self.hash_size, highfreq_factor=self.highfreq_factor
)
return self.lookup_hash(image_hash)[0] 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 collections.abc import Iterable
from typing import Any, NamedTuple, Protocol, Tuple, Union from typing import NamedTuple, Tuple, Union
import numpy as np import numpy as np
# from pylance Mat = np.ndarray
Mat = np.ndarray[int, np.dtype[np.generic]]
class XYWHRect(NamedTuple): class XYWHRect(NamedTuple):
@ -24,19 +23,3 @@ class XYWHRect(NamedTuple):
raise ValueError() raise ValueError()
return self.__class__(*[a - b for a, b in zip(self, other)]) 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 collections.abc import Iterable
from typing import Callable, Tuple, TypeVar, Union, overload from typing import Callable, TypeVar, Union, overload
import cv2 import cv2
import numpy as np import numpy as np
from PIL import Image, ImageCms
from .types import Mat, XYWHRect from .types import XYWHRect
__all__ = ["imread_unicode"] __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 # https://stackoverflow.com/a/57872297/16484891
# CC BY-SA 4.0 # CC BY-SA 4.0
return cv2.imdecode(np.fromfile(filepath, dtype=np.uint8), flags) return cv2.imdecode(np.fromfile(filepath, dtype=np.uint8), flags)
@ -46,25 +44,3 @@ def apply_factor(item, factor: float):
return item * factor return item * factor
elif isinstance(item, Iterable): elif isinstance(item, Iterable):
return item.__class__([i * factor for i in item]) 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