Add: Add argument direct_match to Button methods

This commit is contained in:
LmeSzinc 2023-11-06 22:29:49 +08:00
parent 119c731ef3
commit acc0be8c05

View File

@ -74,7 +74,7 @@ class Button(Resource):
threshold=threshold threshold=threshold
) )
def match_template(self, image, similarity=0.85) -> bool: def match_template(self, image, similarity=0.85, direct_match=False) -> bool:
""" """
Detects assets by template matching. Detects assets by template matching.
@ -83,29 +83,33 @@ class Button(Resource):
Args: Args:
image: Screenshot. image: Screenshot.
similarity (float): 0-1. similarity (float): 0-1.
direct_match: True to ignore `self.search`
Returns: Returns:
bool. bool.
""" """
image = crop(image, self.search, copy=False) if not direct_match:
image = crop(image, self.search, copy=False)
res = cv2.matchTemplate(self.image, image, cv2.TM_CCOEFF_NORMED) res = cv2.matchTemplate(self.image, image, cv2.TM_CCOEFF_NORMED)
_, sim, _, point = cv2.minMaxLoc(res) _, sim, _, point = cv2.minMaxLoc(res)
self._button_offset = np.array(point) + self.search[:2] - self.area[:2] self._button_offset = np.array(point) + self.search[:2] - self.area[:2]
return sim > similarity return sim > similarity
def match_multi_template(self, image, similarity=0.85): def match_multi_template(self, image, similarity=0.85, direct_match=False):
""" """
Detects assets by template matching, return multiple reults Detects assets by template matching, return multiple reults
Args: Args:
image: Screenshot. image: Screenshot.
similarity (float): 0-1. similarity (float): 0-1.
direct_match: True to ignore `self.search`
Returns: Returns:
list: list:
""" """
image = crop(image, self.search, copy=False) if not direct_match:
image = crop(image, self.search, copy=False)
res = cv2.matchTemplate(self.image, image, cv2.TM_CCOEFF_NORMED) res = cv2.matchTemplate(self.image, image, cv2.TM_CCOEFF_NORMED)
res = cv2.inRange(res, similarity, 1.) res = cv2.inRange(res, similarity, 1.)
try: try:
@ -117,7 +121,7 @@ class Button(Resource):
# IndexError: too many indices for array: array is 0-dimensional, but 3 were indexed # IndexError: too many indices for array: array is 0-dimensional, but 3 were indexed
return [] return []
def match_template_color(self, image, similarity=0.85, threshold=30) -> bool: def match_template_color(self, image, similarity=0.85, threshold=30, direct_match=False) -> bool:
""" """
Template match first, color match then Template match first, color match then
@ -125,11 +129,12 @@ class Button(Resource):
image: Screenshot. image: Screenshot.
similarity (float): 0-1. similarity (float): 0-1.
threshold (int): Default to 10. threshold (int): Default to 10.
direct_match: True to ignore `self.search`
Returns: Returns:
bool.
""" """
matched = self.match_template(image, similarity=similarity) matched = self.match_template(image, similarity=similarity, direct_match=direct_match)
if not matched: if not matched:
return False return False
@ -196,43 +201,45 @@ class ButtonWrapper(Resource):
return True return True
return False return False
def match_template(self, image, similarity=0.85) -> bool: def match_template(self, image, similarity=0.85, direct_match=False) -> bool:
for assets in self.buttons: for assets in self.buttons:
if assets.match_template(image, similarity=similarity): if assets.match_template(image, similarity=similarity, direct_match=direct_match):
self._matched_button = assets self._matched_button = assets
return True return True
return False return False
def match_multi_template(self, image, similarity=0.85, threshold=5): def match_multi_template(self, image, similarity=0.85, threshold=5, direct_match=False):
""" """
Detects assets by template matching, return multiple reults Detects assets by template matching, return multiple results
Args: Args:
image: Screenshot. image: Screenshot.
similarity (float): 0-1. similarity (float): 0-1.
threshold: threshold:
direct_match: True to ignore `self.search`
Returns: Returns:
list[ClickButton]: list[ClickButton]:
""" """
points = [] ps = []
for assets in self.buttons: for assets in self.buttons:
points += assets.match_multi_template(image, similarity=similarity) ps += assets.match_multi_template(image, similarity=similarity, direct_match=direct_match)
if not points: if not ps:
return [] return []
from module.base.utils.points import Points from module.base.utils.points import Points
points = Points(points).group(threshold=threshold) ps = Points(ps).group(threshold=threshold)
area_list = [area_offset(self.area, p - self.area[:2]) for p in points] area_list = [area_offset(self.area, p - self.area[:2]) for p in ps]
button_list = [area_offset(self.button, p - self.area[:2]) for p in points] button_list = [area_offset(self.button, p - self.area[:2]) for p in ps]
return [ return [
ClickButton(area=info[0], button=info[1], name=f'{self.name}_result{i}') ClickButton(area=info[0], button=info[1], name=f'{self.name}_result{i}')
for i, info in enumerate(zip(area_list, button_list)) for i, info in enumerate(zip(area_list, button_list))
] ]
def match_template_color(self, image, similarity=0.85, threshold=30) -> bool: def match_template_color(self, image, similarity=0.85, threshold=30, direct_match=False) -> bool:
for assets in self.buttons: for assets in self.buttons:
if assets.match_template_color(image, similarity=similarity, threshold=threshold): if assets.match_template_color(
image, similarity=similarity, threshold=threshold, direct_match=direct_match):
self._matched_button = assets self._matched_button = assets
return True return True
return False return False
@ -301,12 +308,9 @@ class ButtonWrapper(Resource):
class ClickButton: class ClickButton:
def __init__(self, area, button=None, name='CLICK_BUTTON'): def __init__(self, button, name='CLICK_BUTTON'):
self.area = area self.area = button
if button is None: self.button = button
self.button = area
else:
self.button = button
self.name = name self.name = name
def __str__(self): def __str__(self):