mirror of
https://github.com/283375/arcaea-offline-ocr.git
synced 2025-04-22 15:00:18 +00:00
Compare commits
No commits in common. "07df8e3b5693c8ff74598ca3052e1b6d3f955e75" and "8389586ccdc840fbb353a03bb2bb6e3884b7f741" have entirely different histories.
07df8e3b56
...
8389586ccd
@ -3,11 +3,9 @@ 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 ....sift_db import SIFTDatabase
|
||||
from ....types import Mat, cv2_ml_KNearest
|
||||
from ....utils import construct_int_xywh_rect
|
||||
@ -21,12 +19,12 @@ class ChieriBotV4Ocr:
|
||||
self,
|
||||
score_knn: cv2_ml_KNearest,
|
||||
pfl_knn: cv2_ml_KNearest,
|
||||
phash_db: ImagePHashDatabase,
|
||||
sift_db: SIFTDatabase,
|
||||
factor: Optional[float] = 1.0,
|
||||
):
|
||||
self.__score_knn = score_knn
|
||||
self.__pfl_knn = pfl_knn
|
||||
self.__phash_db = phash_db
|
||||
self.__sift_db = sift_db
|
||||
self.__rois = ChieriBotV4Rois(factor)
|
||||
|
||||
@property
|
||||
@ -46,12 +44,12 @@ class ChieriBotV4Ocr:
|
||||
self.__pfl_knn = knn_digits_model
|
||||
|
||||
@property
|
||||
def phash_db(self):
|
||||
return self.__phash_db
|
||||
def sift_db(self):
|
||||
return self.__sift_db
|
||||
|
||||
@phash_db.setter
|
||||
def phash_db(self, phash_db: ImagePHashDatabase):
|
||||
self.__phash_db = phash_db
|
||||
@sift_db.setter
|
||||
def sift_db(self, sift_db: SIFTDatabase):
|
||||
self.__sift_db = sift_db
|
||||
|
||||
@property
|
||||
def rois(self):
|
||||
@ -100,7 +98,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]
|
||||
return self.sift_db.lookup_img(jacket_roi)[0]
|
||||
|
||||
# def ocr_component_score_paddle(self, component_bgr: Mat) -> int:
|
||||
# # sourcery skip: inline-immediately-returned-variable
|
||||
|
||||
@ -4,19 +4,9 @@ from typing import Sequence
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from ...crop import crop_xywh
|
||||
from ...mask import (
|
||||
mask_byd,
|
||||
mask_ftr,
|
||||
mask_gray,
|
||||
mask_max_recall_purple,
|
||||
mask_pfl_white,
|
||||
mask_prs,
|
||||
mask_pst,
|
||||
mask_white,
|
||||
)
|
||||
from ...mask import mask_byd, mask_ftr, mask_gray, mask_prs, mask_pst, mask_white
|
||||
from ...ocr import (
|
||||
FixRects,
|
||||
ocr_digit_samples_knn,
|
||||
@ -24,20 +14,18 @@ from ...ocr import (
|
||||
preprocess_hog,
|
||||
resize_fill_square,
|
||||
)
|
||||
from ...phash_db import ImagePHashDatabase
|
||||
from ...sift_db import SIFTDatabase
|
||||
from ...types import Mat, cv2_ml_KNearest
|
||||
from ..shared import DeviceOcrResult
|
||||
from .preprocess import find_digits_preprocess
|
||||
from .rois import DeviceV2Rois
|
||||
from .shared import MAX_RECALL_CLOSE_KERNEL
|
||||
from .sizes import SizesV2
|
||||
|
||||
|
||||
class DeviceV2Ocr:
|
||||
def __init__(self, knn_model: cv2_ml_KNearest, phash_db: ImagePHashDatabase):
|
||||
def __init__(self, knn_model: cv2_ml_KNearest, sift_db: SIFTDatabase):
|
||||
self.__knn_model = knn_model
|
||||
self.__phash_db = phash_db
|
||||
self.__sift_db = sift_db
|
||||
|
||||
@property
|
||||
def knn_model(self):
|
||||
@ -50,14 +38,14 @@ class DeviceV2Ocr:
|
||||
self.__knn_model = value
|
||||
|
||||
@property
|
||||
def phash_db(self):
|
||||
if not self.__phash_db:
|
||||
raise ValueError("`phash_db` unset.")
|
||||
return self.__phash_db
|
||||
def sift_db(self):
|
||||
if not self.__sift_db:
|
||||
raise ValueError("`sift_db` unset.")
|
||||
return self.__sift_db
|
||||
|
||||
@phash_db.setter
|
||||
def phash_db(self, value: SIFTDatabase):
|
||||
self.__phash_db = value
|
||||
@sift_db.setter
|
||||
def sift_db(self, value: SIFTDatabase):
|
||||
self.__sift_db = value
|
||||
|
||||
@lru_cache
|
||||
def _get_digit_widths(self, num_list: Sequence[int], factor: float):
|
||||
@ -98,7 +86,7 @@ class DeviceV2Ocr:
|
||||
|
||||
def ocr_song_id(self, rois: DeviceV2Rois):
|
||||
jacket = cv2.cvtColor(rois.jacket, cv2.COLOR_BGR2GRAY)
|
||||
return self.phash_db.lookup_image(Image.fromarray(jacket))[0]
|
||||
return self.sift_db.lookup_img(jacket)[0]
|
||||
|
||||
def ocr_rating_class(self, rois: DeviceV2Rois):
|
||||
roi = cv2.cvtColor(rois.max_recall_rating_class, cv2.COLOR_BGR2HSV)
|
||||
@ -115,33 +103,20 @@ class DeviceV2Ocr:
|
||||
roi = cv2.fillPoly(roi, [contour], [0])
|
||||
return ocr_digits_by_contour_knn(roi, self.knn_model)
|
||||
|
||||
def mask_pfl(self, pfl_roi: Mat, rois: DeviceV2Rois):
|
||||
return (
|
||||
mask_pfl_white(cv2.cvtColor(pfl_roi, cv2.COLOR_BGR2HSV))
|
||||
if isinstance(rois.sizes, SizesV2)
|
||||
else mask_gray(pfl_roi)
|
||||
)
|
||||
|
||||
def ocr_pure(self, rois: DeviceV2Rois):
|
||||
roi = self.mask_pfl(rois.pure, rois)
|
||||
roi = mask_gray(rois.pure)
|
||||
return self._base_ocr_pfl(roi, rois.sizes.factor)
|
||||
|
||||
def ocr_far(self, rois: DeviceV2Rois):
|
||||
roi = self.mask_pfl(rois.far, rois)
|
||||
roi = mask_gray(rois.far)
|
||||
return self._base_ocr_pfl(roi, rois.sizes.factor)
|
||||
|
||||
def ocr_lost(self, rois: DeviceV2Rois):
|
||||
roi = self.mask_pfl(rois.lost, rois)
|
||||
roi = mask_gray(rois.lost)
|
||||
return self._base_ocr_pfl(roi, rois.sizes.factor)
|
||||
|
||||
def ocr_max_recall(self, rois: DeviceV2Rois):
|
||||
roi = (
|
||||
mask_max_recall_purple(
|
||||
cv2.cvtColor(rois.max_recall_rating_class, cv2.COLOR_BGR2HSV)
|
||||
)
|
||||
if isinstance(rois.sizes, SizesV2)
|
||||
else mask_gray(rois.max_recall_rating_class)
|
||||
)
|
||||
roi = mask_gray(rois.max_recall_rating_class)
|
||||
roi_closed = cv2.morphologyEx(roi, cv2.MORPH_CLOSE, MAX_RECALL_CLOSE_KERNEL)
|
||||
contours, _ = cv2.findContours(
|
||||
roi_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
|
||||
|
||||
@ -1,15 +1,183 @@
|
||||
from typing import Union
|
||||
from typing import Tuple, Union
|
||||
|
||||
from ...crop import crop_black_edges, crop_xywh
|
||||
from ...types import Mat, XYWHRect
|
||||
from .definition import DeviceV2
|
||||
from .sizes import Sizes, SizesV1
|
||||
|
||||
|
||||
def to_int(num: Union[int, float]) -> int:
|
||||
return round(num)
|
||||
|
||||
|
||||
def apply_factor(num: Union[int, float], factor: float):
|
||||
return num * factor
|
||||
|
||||
|
||||
class Sizes:
|
||||
def __init__(self, factor: float):
|
||||
raise NotImplementedError()
|
||||
|
||||
@property
|
||||
def TOP_BAR_HEIGHT(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def SCORE_PANEL(self) -> Tuple[int, int]:
|
||||
...
|
||||
|
||||
@property
|
||||
def PFL_TOP_FROM_VMID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def PFL_LEFT_FROM_HMID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def PFL_WIDTH(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def PFL_FONT_PX(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def PURE_FAR_GAP(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def FAR_LOST_GAP(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def SCORE_BOTTOM_FROM_VMID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def SCORE_FONT_PX(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def SCORE_WIDTH(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def JACKET_RIGHT_FROM_HOR_MID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def JACKET_WIDTH(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def MR_RT_RIGHT_FROM_HMID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def MR_RT_WIDTH(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def MR_RT_HEIGHT(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def TITLE_BOTTOM_FROM_VMID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def TITLE_FONT_PX(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def TITLE_WIDTH_RIGHT(self):
|
||||
...
|
||||
|
||||
|
||||
class SizesV1(Sizes):
|
||||
def __init__(self, factor: float):
|
||||
self.factor = factor
|
||||
|
||||
def apply_factor(self, num):
|
||||
return apply_factor(num, self.factor)
|
||||
|
||||
@property
|
||||
def TOP_BAR_HEIGHT(self):
|
||||
return self.apply_factor(50)
|
||||
|
||||
@property
|
||||
def SCORE_PANEL(self) -> Tuple[int, int]:
|
||||
return tuple(self.apply_factor(num) for num in [485, 239])
|
||||
|
||||
@property
|
||||
def PFL_TOP_FROM_VMID(self):
|
||||
return self.apply_factor(135)
|
||||
|
||||
@property
|
||||
def PFL_LEFT_FROM_HMID(self):
|
||||
return self.apply_factor(5)
|
||||
|
||||
@property
|
||||
def PFL_WIDTH(self):
|
||||
return self.apply_factor(76)
|
||||
|
||||
@property
|
||||
def PFL_FONT_PX(self):
|
||||
return self.apply_factor(26)
|
||||
|
||||
@property
|
||||
def PURE_FAR_GAP(self):
|
||||
return self.apply_factor(12)
|
||||
|
||||
@property
|
||||
def FAR_LOST_GAP(self):
|
||||
return self.apply_factor(10)
|
||||
|
||||
@property
|
||||
def SCORE_BOTTOM_FROM_VMID(self):
|
||||
return self.apply_factor(-50)
|
||||
|
||||
@property
|
||||
def SCORE_FONT_PX(self):
|
||||
return self.apply_factor(45)
|
||||
|
||||
@property
|
||||
def SCORE_WIDTH(self):
|
||||
return self.apply_factor(280)
|
||||
|
||||
@property
|
||||
def JACKET_RIGHT_FROM_HOR_MID(self):
|
||||
return self.apply_factor(-235)
|
||||
|
||||
@property
|
||||
def JACKET_WIDTH(self):
|
||||
return self.apply_factor(375)
|
||||
|
||||
@property
|
||||
def MR_RT_RIGHT_FROM_HMID(self):
|
||||
return self.apply_factor(-300)
|
||||
|
||||
@property
|
||||
def MR_RT_WIDTH(self):
|
||||
return self.apply_factor(275)
|
||||
|
||||
@property
|
||||
def MR_RT_HEIGHT(self):
|
||||
return self.apply_factor(75)
|
||||
|
||||
@property
|
||||
def TITLE_BOTTOM_FROM_VMID(self):
|
||||
return self.apply_factor(-265)
|
||||
|
||||
@property
|
||||
def TITLE_FONT_PX(self):
|
||||
return self.apply_factor(40)
|
||||
|
||||
@property
|
||||
def TITLE_WIDTH_RIGHT(self):
|
||||
return self.apply_factor(275)
|
||||
|
||||
|
||||
class DeviceV2Rois:
|
||||
def __init__(self, device: DeviceV2, img: Mat):
|
||||
self.device = device
|
||||
|
||||
@ -1,254 +0,0 @@
|
||||
from typing import Tuple, Union
|
||||
|
||||
|
||||
def apply_factor(num: Union[int, float], factor: float):
|
||||
return num * factor
|
||||
|
||||
|
||||
class Sizes:
|
||||
def __init__(self, factor: float):
|
||||
raise NotImplementedError()
|
||||
|
||||
@property
|
||||
def TOP_BAR_HEIGHT(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def SCORE_PANEL(self) -> Tuple[int, int]:
|
||||
...
|
||||
|
||||
@property
|
||||
def PFL_TOP_FROM_VMID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def PFL_LEFT_FROM_HMID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def PFL_WIDTH(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def PFL_FONT_PX(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def PURE_FAR_GAP(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def FAR_LOST_GAP(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def SCORE_BOTTOM_FROM_VMID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def SCORE_FONT_PX(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def SCORE_WIDTH(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def JACKET_RIGHT_FROM_HOR_MID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def JACKET_WIDTH(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def MR_RT_RIGHT_FROM_HMID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def MR_RT_WIDTH(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def MR_RT_HEIGHT(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def TITLE_BOTTOM_FROM_VMID(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def TITLE_FONT_PX(self):
|
||||
...
|
||||
|
||||
@property
|
||||
def TITLE_WIDTH_RIGHT(self):
|
||||
...
|
||||
|
||||
|
||||
class SizesV1(Sizes):
|
||||
def __init__(self, factor: float):
|
||||
self.factor = factor
|
||||
|
||||
def apply_factor(self, num):
|
||||
return apply_factor(num, self.factor)
|
||||
|
||||
@property
|
||||
def TOP_BAR_HEIGHT(self):
|
||||
return self.apply_factor(50)
|
||||
|
||||
@property
|
||||
def SCORE_PANEL(self) -> Tuple[int, int]:
|
||||
return tuple(self.apply_factor(num) for num in [485, 239])
|
||||
|
||||
@property
|
||||
def PFL_TOP_FROM_VMID(self):
|
||||
return self.apply_factor(135)
|
||||
|
||||
@property
|
||||
def PFL_LEFT_FROM_HMID(self):
|
||||
return self.apply_factor(5)
|
||||
|
||||
@property
|
||||
def PFL_WIDTH(self):
|
||||
return self.apply_factor(76)
|
||||
|
||||
@property
|
||||
def PFL_FONT_PX(self):
|
||||
return self.apply_factor(26)
|
||||
|
||||
@property
|
||||
def PURE_FAR_GAP(self):
|
||||
return self.apply_factor(12)
|
||||
|
||||
@property
|
||||
def FAR_LOST_GAP(self):
|
||||
return self.apply_factor(10)
|
||||
|
||||
@property
|
||||
def SCORE_BOTTOM_FROM_VMID(self):
|
||||
return self.apply_factor(-50)
|
||||
|
||||
@property
|
||||
def SCORE_FONT_PX(self):
|
||||
return self.apply_factor(45)
|
||||
|
||||
@property
|
||||
def SCORE_WIDTH(self):
|
||||
return self.apply_factor(280)
|
||||
|
||||
@property
|
||||
def JACKET_RIGHT_FROM_HOR_MID(self):
|
||||
return self.apply_factor(-235)
|
||||
|
||||
@property
|
||||
def JACKET_WIDTH(self):
|
||||
return self.apply_factor(375)
|
||||
|
||||
@property
|
||||
def MR_RT_RIGHT_FROM_HMID(self):
|
||||
return self.apply_factor(-300)
|
||||
|
||||
@property
|
||||
def MR_RT_WIDTH(self):
|
||||
return self.apply_factor(275)
|
||||
|
||||
@property
|
||||
def MR_RT_HEIGHT(self):
|
||||
return self.apply_factor(75)
|
||||
|
||||
@property
|
||||
def TITLE_BOTTOM_FROM_VMID(self):
|
||||
return self.apply_factor(-265)
|
||||
|
||||
@property
|
||||
def TITLE_FONT_PX(self):
|
||||
return self.apply_factor(40)
|
||||
|
||||
@property
|
||||
def TITLE_WIDTH_RIGHT(self):
|
||||
return self.apply_factor(275)
|
||||
|
||||
|
||||
class SizesV2(Sizes):
|
||||
def __init__(self, factor: float):
|
||||
self.factor = factor
|
||||
|
||||
def apply_factor(self, num):
|
||||
return apply_factor(num, self.factor)
|
||||
|
||||
@property
|
||||
def TOP_BAR_HEIGHT(self):
|
||||
return self.apply_factor(50)
|
||||
|
||||
@property
|
||||
def SCORE_PANEL(self) -> Tuple[int, int]:
|
||||
return tuple(self.apply_factor(num) for num in [447, 233])
|
||||
|
||||
@property
|
||||
def PFL_TOP_FROM_VMID(self):
|
||||
return self.apply_factor(142)
|
||||
|
||||
@property
|
||||
def PFL_LEFT_FROM_HMID(self):
|
||||
return self.apply_factor(10)
|
||||
|
||||
@property
|
||||
def PFL_WIDTH(self):
|
||||
return self.apply_factor(60)
|
||||
|
||||
@property
|
||||
def PFL_FONT_PX(self):
|
||||
return self.apply_factor(16)
|
||||
|
||||
@property
|
||||
def PURE_FAR_GAP(self):
|
||||
return self.apply_factor(20)
|
||||
|
||||
@property
|
||||
def FAR_LOST_GAP(self):
|
||||
return self.apply_factor(23)
|
||||
|
||||
@property
|
||||
def SCORE_BOTTOM_FROM_VMID(self):
|
||||
return self.apply_factor(-50)
|
||||
|
||||
@property
|
||||
def SCORE_FONT_PX(self):
|
||||
return self.apply_factor(45)
|
||||
|
||||
@property
|
||||
def SCORE_WIDTH(self):
|
||||
return self.apply_factor(280)
|
||||
|
||||
@property
|
||||
def JACKET_RIGHT_FROM_HOR_MID(self):
|
||||
return self.apply_factor(-235)
|
||||
|
||||
@property
|
||||
def JACKET_WIDTH(self):
|
||||
return self.apply_factor(375)
|
||||
|
||||
@property
|
||||
def MR_RT_RIGHT_FROM_HMID(self):
|
||||
return self.apply_factor(-330)
|
||||
|
||||
@property
|
||||
def MR_RT_WIDTH(self):
|
||||
return self.apply_factor(330)
|
||||
|
||||
@property
|
||||
def MR_RT_HEIGHT(self):
|
||||
return self.apply_factor(75)
|
||||
|
||||
@property
|
||||
def TITLE_BOTTOM_FROM_VMID(self):
|
||||
return self.apply_factor(-265)
|
||||
|
||||
@property
|
||||
def TITLE_FONT_PX(self):
|
||||
return self.apply_factor(40)
|
||||
|
||||
@property
|
||||
def TITLE_WIDTH_RIGHT(self):
|
||||
return self.apply_factor(275)
|
||||
@ -8,8 +8,6 @@ __all__ = [
|
||||
"GRAY_MAX_HSV",
|
||||
"WHITE_MIN_HSV",
|
||||
"WHITE_MAX_HSV",
|
||||
"PFL_WHITE_MIN_HSV",
|
||||
"PFL_WHITE_MAX_HSV",
|
||||
"PST_MIN_HSV",
|
||||
"PST_MAX_HSV",
|
||||
"PRS_MIN_HSV",
|
||||
@ -18,17 +16,13 @@ __all__ = [
|
||||
"FTR_MAX_HSV",
|
||||
"BYD_MIN_HSV",
|
||||
"BYD_MAX_HSV",
|
||||
"MAX_RECALL_PURPLE_MIN_HSV",
|
||||
"MAX_RECALL_PURPLE_MAX_HSV",
|
||||
"mask_gray",
|
||||
"mask_white",
|
||||
"mask_pfl_white",
|
||||
"mask_pst",
|
||||
"mask_prs",
|
||||
"mask_ftr",
|
||||
"mask_byd",
|
||||
"mask_rating_class",
|
||||
"mask_max_recall_purple",
|
||||
]
|
||||
|
||||
GRAY_MIN_HSV = np.array([0, 0, 70], np.uint8)
|
||||
@ -40,9 +34,6 @@ GRAY_MAX_BGR = np.array([160] * 3, np.uint8)
|
||||
WHITE_MIN_HSV = np.array([0, 0, 240], np.uint8)
|
||||
WHITE_MAX_HSV = np.array([179, 10, 255], np.uint8)
|
||||
|
||||
PFL_WHITE_MIN_HSV = np.array([0, 0, 248], np.uint8)
|
||||
PFL_WHITE_MAX_HSV = np.array([179, 10, 255], np.uint8)
|
||||
|
||||
PST_MIN_HSV = np.array([100, 50, 80], np.uint8)
|
||||
PST_MAX_HSV = np.array([100, 255, 255], np.uint8)
|
||||
|
||||
@ -55,9 +46,6 @@ FTR_MAX_HSV = np.array([155, 181, 150], np.uint8)
|
||||
BYD_MIN_HSV = np.array([170, 50, 50], np.uint8)
|
||||
BYD_MAX_HSV = np.array([179, 210, 198], np.uint8)
|
||||
|
||||
MAX_RECALL_PURPLE_MIN_HSV = np.array([125, 0, 0], np.uint8)
|
||||
MAX_RECALL_PURPLE_MAX_HSV = np.array([130, 100, 150], np.uint8)
|
||||
|
||||
|
||||
def mask_gray(__img_bgr: Mat):
|
||||
# bgr_value_equal_mask = all(__img_bgr[:, 1:] == __img_bgr[:, :-1], axis=1)
|
||||
@ -75,12 +63,6 @@ def mask_white(img_hsv: Mat):
|
||||
return mask
|
||||
|
||||
|
||||
def mask_pfl_white(img_hsv: Mat):
|
||||
mask = cv2.inRange(img_hsv, PFL_WHITE_MIN_HSV, PFL_WHITE_MAX_HSV)
|
||||
mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2)))
|
||||
return mask
|
||||
|
||||
|
||||
def mask_pst(img_hsv: Mat):
|
||||
mask = cv2.inRange(img_hsv, PST_MIN_HSV, PST_MAX_HSV)
|
||||
mask = cv2.dilate(mask, (1, 1))
|
||||
@ -111,9 +93,3 @@ def mask_rating_class(img_hsv: Mat):
|
||||
ftr = mask_ftr(img_hsv)
|
||||
byd = mask_byd(img_hsv)
|
||||
return cv2.bitwise_or(byd, cv2.bitwise_or(ftr, cv2.bitwise_or(pst, prs)))
|
||||
|
||||
|
||||
def mask_max_recall_purple(img_hsv: Mat):
|
||||
mask = cv2.inRange(img_hsv, MAX_RECALL_PURPLE_MIN_HSV, MAX_RECALL_PURPLE_MAX_HSV)
|
||||
mask = cv2.dilate(mask, (2, 2))
|
||||
return mask
|
||||
|
||||
@ -1,65 +0,0 @@
|
||||
import sqlite3
|
||||
|
||||
import imagehash
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def hamming_distance_sql_function(user_input, db_entry) -> int:
|
||||
return np.count_nonzero(
|
||||
np.frombuffer(user_input, bool) ^ np.frombuffer(db_entry, bool)
|
||||
)
|
||||
|
||||
|
||||
class ImagePHashDatabase:
|
||||
def __init__(self, db_path: str):
|
||||
with sqlite3.connect(db_path) as conn:
|
||||
self.hash_size = int(
|
||||
conn.execute(
|
||||
"SELECT value FROM properties WHERE key = 'hash_size'"
|
||||
).fetchone()[0]
|
||||
)
|
||||
self.highfreq_factor = int(
|
||||
conn.execute(
|
||||
"SELECT value FROM properties WHERE key = 'highfreq_factor'"
|
||||
).fetchone()[0]
|
||||
)
|
||||
self.built_timestamp = int(
|
||||
conn.execute(
|
||||
"SELECT value FROM properties WHERE key = 'built_timestamp'"
|
||||
).fetchone()[0]
|
||||
)
|
||||
|
||||
# self.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.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]
|
||||
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
|
||||
)
|
||||
return self.lookup_hash(image_hash)[0]
|
||||
Loading…
x
Reference in New Issue
Block a user