mirror of
https://github.com/LmeSzinc/StarRailCopilot.git
synced 2024-11-16 14:31:16 +00:00
377 lines
11 KiB
Python
377 lines
11 KiB
Python
import module.config.server as server
|
||
from module.base.decorator import cached_property, del_cached_property
|
||
from module.base.resource import Resource
|
||
from module.base.utils import *
|
||
from module.exception import ScriptError
|
||
|
||
|
||
class Button(Resource):
|
||
def __init__(self, file, area, search, color, button):
|
||
"""
|
||
Args:
|
||
file: Filepath to an assets
|
||
area: Area to crop template
|
||
search: Area to search from, 20px larger than `area` by default
|
||
color: Average color of assets
|
||
button: Area to click if assets appears on the image
|
||
"""
|
||
self.file: str = file
|
||
self.area: t.Tuple[int, int, int, int] = area
|
||
self.search: t.Tuple[int, int, int, int] = search
|
||
self.color: t.Tuple[int, int, int] = color
|
||
self._button: t.Tuple[int, int, int, int] = button
|
||
|
||
self.resource_add(self.file)
|
||
self._button_offset: t.Tuple[int, int] = (0, 0)
|
||
|
||
@property
|
||
def button(self):
|
||
return area_offset(self._button, self._button_offset)
|
||
|
||
def load_offset(self, button):
|
||
self._button_offset = button._button_offset
|
||
|
||
def clear_offset(self):
|
||
self._button_offset = (0, 0)
|
||
|
||
@cached_property
|
||
def image(self):
|
||
return load_image(self.file, self.area)
|
||
|
||
def resource_release(self):
|
||
del_cached_property(self, 'image')
|
||
self.clear_offset()
|
||
|
||
def __str__(self):
|
||
return self.file
|
||
|
||
__repr__ = __str__
|
||
|
||
def __eq__(self, other):
|
||
return str(self) == str(other)
|
||
|
||
def __hash__(self):
|
||
return hash(self.file)
|
||
|
||
def __bool__(self):
|
||
return True
|
||
|
||
def match_color(self, image, threshold=10) -> bool:
|
||
"""
|
||
Check if the button appears on the image, using average color
|
||
|
||
Args:
|
||
image (np.ndarray): Screenshot.
|
||
threshold (int): Default to 10.
|
||
|
||
Returns:
|
||
bool: True if button appears on screenshot.
|
||
"""
|
||
color = get_color(image, self.area)
|
||
return color_similar(
|
||
color1=color,
|
||
color2=self.color,
|
||
threshold=threshold
|
||
)
|
||
|
||
def match_template(self, image, similarity=0.85, direct_match=False) -> bool:
|
||
"""
|
||
Detects assets by template matching.
|
||
|
||
To Some buttons, its location may not be static, `_button_offset` will be set.
|
||
|
||
Args:
|
||
image: Screenshot.
|
||
similarity (float): 0-1.
|
||
direct_match: True to ignore `self.search`
|
||
|
||
Returns:
|
||
bool.
|
||
"""
|
||
if not direct_match:
|
||
image = crop(image, self.search, copy=False)
|
||
res = cv2.matchTemplate(self.image, image, cv2.TM_CCOEFF_NORMED)
|
||
_, sim, _, point = cv2.minMaxLoc(res)
|
||
|
||
self._button_offset = np.array(point) + self.search[:2] - self.area[:2]
|
||
return sim > similarity
|
||
|
||
def match_multi_template(self, image, similarity=0.85, direct_match=False):
|
||
"""
|
||
Detects assets by template matching, return multiple reults
|
||
|
||
Args:
|
||
image: Screenshot.
|
||
similarity (float): 0-1.
|
||
direct_match: True to ignore `self.search`
|
||
|
||
Returns:
|
||
list:
|
||
"""
|
||
if not direct_match:
|
||
image = crop(image, self.search, copy=False)
|
||
res = cv2.matchTemplate(self.image, image, cv2.TM_CCOEFF_NORMED)
|
||
res = cv2.inRange(res, similarity, 1.)
|
||
try:
|
||
points = np.array(cv2.findNonZero(res))[:, 0, :]
|
||
points += self.search[:2]
|
||
return points.tolist()
|
||
except IndexError:
|
||
# Empty result
|
||
# IndexError: too many indices for array: array is 0-dimensional, but 3 were indexed
|
||
return []
|
||
|
||
def match_template_color(self, image, similarity=0.85, threshold=30, direct_match=False) -> bool:
|
||
"""
|
||
Template match first, color match then
|
||
|
||
Args:
|
||
image: Screenshot.
|
||
similarity (float): 0-1.
|
||
threshold (int): Default to 10.
|
||
direct_match: True to ignore `self.search`
|
||
|
||
Returns:
|
||
bool.
|
||
"""
|
||
matched = self.match_template(image, similarity=similarity, direct_match=direct_match)
|
||
if not matched:
|
||
return False
|
||
|
||
area = area_offset(self.area, offset=self._button_offset)
|
||
color = get_color(image, area)
|
||
return color_similar(
|
||
color1=color,
|
||
color2=self.color,
|
||
threshold=threshold
|
||
)
|
||
|
||
|
||
class ButtonWrapper(Resource):
|
||
def __init__(self, name='MULTI_ASSETS', **kwargs):
|
||
self.name = name
|
||
self.data_buttons = kwargs
|
||
self._matched_button: t.Optional[Button] = None
|
||
self.resource_add(f'{name}:{next(self.iter_buttons(), None)}')
|
||
|
||
def resource_release(self):
|
||
del_cached_property(self, 'buttons')
|
||
self._matched_button = None
|
||
|
||
def __str__(self):
|
||
return self.name
|
||
|
||
__repr__ = __str__
|
||
|
||
def __eq__(self, other):
|
||
return str(self) == str(other)
|
||
|
||
def __hash__(self):
|
||
return hash(self.name)
|
||
|
||
def __bool__(self):
|
||
return True
|
||
|
||
def iter_buttons(self) -> t.Iterator[Button]:
|
||
for _, assets in self.data_buttons.items():
|
||
if isinstance(assets, Button):
|
||
yield assets
|
||
elif isinstance(assets, list):
|
||
for asset in assets:
|
||
yield asset
|
||
|
||
@cached_property
|
||
def buttons(self) -> t.List[Button]:
|
||
for trial in [server.lang, 'share', 'cn']:
|
||
try:
|
||
assets = self.data_buttons[trial]
|
||
if isinstance(assets, Button):
|
||
return [assets]
|
||
elif isinstance(assets, list):
|
||
return assets
|
||
except KeyError:
|
||
pass
|
||
|
||
raise ScriptError(f'ButtonWrapper({self}) on server {server.lang} has no fallback button')
|
||
|
||
def match_color(self, image, threshold=10) -> bool:
|
||
for assets in self.buttons:
|
||
if assets.match_color(image, threshold=threshold):
|
||
self._matched_button = assets
|
||
return True
|
||
return False
|
||
|
||
def match_template(self, image, similarity=0.85, direct_match=False) -> bool:
|
||
for assets in self.buttons:
|
||
if assets.match_template(image, similarity=similarity, direct_match=direct_match):
|
||
self._matched_button = assets
|
||
return True
|
||
return False
|
||
|
||
def match_multi_template(self, image, similarity=0.85, threshold=5, direct_match=False):
|
||
"""
|
||
Detects assets by template matching, return multiple results
|
||
|
||
Args:
|
||
image: Screenshot.
|
||
similarity (float): 0-1.
|
||
threshold:
|
||
direct_match: True to ignore `self.search`
|
||
|
||
Returns:
|
||
list[ClickButton]:
|
||
"""
|
||
ps = []
|
||
for assets in self.buttons:
|
||
ps += assets.match_multi_template(image, similarity=similarity, direct_match=direct_match)
|
||
if not ps:
|
||
return []
|
||
|
||
from module.base.utils.points import Points
|
||
ps = Points(ps).group(threshold=threshold)
|
||
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 ps]
|
||
return [
|
||
ClickButton(area=info[0], button=info[1], name=f'{self.name}_result{i}')
|
||
for i, info in enumerate(zip(area_list, button_list))
|
||
]
|
||
|
||
def match_template_color(self, image, similarity=0.85, threshold=30, direct_match=False) -> bool:
|
||
for assets in self.buttons:
|
||
if assets.match_template_color(
|
||
image, similarity=similarity, threshold=threshold, direct_match=direct_match):
|
||
self._matched_button = assets
|
||
return True
|
||
return False
|
||
|
||
@property
|
||
def matched_button(self) -> Button:
|
||
if self._matched_button is None:
|
||
return self.buttons[0]
|
||
else:
|
||
return self._matched_button
|
||
|
||
@property
|
||
def area(self) -> tuple[int, int, int, int]:
|
||
return self.matched_button.area
|
||
|
||
@property
|
||
def search(self) -> tuple[int, int, int, int]:
|
||
return self.matched_button.search
|
||
|
||
@property
|
||
def color(self) -> tuple[int, int, int]:
|
||
return self.matched_button.color
|
||
|
||
@property
|
||
def button(self) -> tuple[int, int, int, int]:
|
||
return self.matched_button.button
|
||
|
||
@property
|
||
def button_offset(self) -> tuple[int, int]:
|
||
return self.matched_button._button_offset
|
||
|
||
@property
|
||
def width(self) -> int:
|
||
return area_size(self.area)[0]
|
||
|
||
@property
|
||
def height(self) -> int:
|
||
return area_size(self.area)[1]
|
||
|
||
def load_offset(self, button):
|
||
"""
|
||
Load offset from another button.
|
||
|
||
Args:
|
||
button (Button, ButtonWrapper):
|
||
"""
|
||
if isinstance(button, ButtonWrapper):
|
||
button = button.matched_button
|
||
for b in self.iter_buttons():
|
||
b.load_offset(button)
|
||
|
||
def clear_offset(self):
|
||
for b in self.iter_buttons():
|
||
b.clear_offset()
|
||
|
||
def load_search(self, area):
|
||
"""
|
||
Set `search` attribute.
|
||
Note that this method is irreversible.
|
||
|
||
Args:
|
||
area:
|
||
"""
|
||
for b in self.iter_buttons():
|
||
b.search = area
|
||
|
||
def set_search_offset(self, offset):
|
||
"""
|
||
Compatible with Alas’ `offset` attribute
|
||
In ALAS:
|
||
if self.appear(BUTTON, offset=(20, 20)):
|
||
pass
|
||
In SRC:
|
||
BUTTON.set_search_offset((20, 20))
|
||
if self.appear(BUTTON):
|
||
pass
|
||
Note that `search` attribute will be set, and it's irreversible.
|
||
|
||
Args:
|
||
offset (tuple): (x, y) or (left, up, right, bottom)
|
||
"""
|
||
if len(offset) == 2:
|
||
left, up, right, bottom = -offset[0], -offset[1], offset[0], offset[1]
|
||
else:
|
||
left, up, right, bottom = offset
|
||
for b in self.iter_buttons():
|
||
upper_left_x, upper_left_y, bottom_right_x, bottom_right_y = b.area
|
||
b.search = (
|
||
upper_left_x + left,
|
||
upper_left_y + up,
|
||
bottom_right_x + right,
|
||
bottom_right_y + bottom,
|
||
)
|
||
|
||
|
||
class ClickButton:
|
||
def __init__(self, area, button=None, name='CLICK_BUTTON'):
|
||
self.area = area
|
||
if button is None:
|
||
self.button = area
|
||
else:
|
||
self.button = button
|
||
self.name = name
|
||
|
||
def __str__(self):
|
||
return self.name
|
||
|
||
__repr__ = __str__
|
||
|
||
def __eq__(self, other):
|
||
return str(self) == str(other)
|
||
|
||
def __hash__(self):
|
||
return hash(self.name)
|
||
|
||
def __bool__(self):
|
||
return True
|
||
|
||
|
||
def match_template(image, template, similarity=0.85):
|
||
"""
|
||
Args:
|
||
image (np.ndarray): Screenshot
|
||
template (np.ndarray):
|
||
area (tuple): Crop area of image.
|
||
offset (int, tuple): Detection area offset.
|
||
similarity (float): 0-1. Similarity. Lower than this value will return float(0).
|
||
|
||
Returns:
|
||
bool:
|
||
"""
|
||
res = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
|
||
_, sim, _, point = cv2.minMaxLoc(res)
|
||
return sim > similarity
|