mirror of
https://github.com/LmeSzinc/StarRailCopilot.git
synced 2024-11-16 14:31:16 +00:00
476 lines
16 KiB
Python
476 lines
16 KiB
Python
import os
|
|
import traceback
|
|
|
|
import imageio
|
|
from PIL import ImageDraw
|
|
|
|
from module.base.decorator import cached_property
|
|
from module.base.resource import Resource
|
|
from module.base.utils import *
|
|
from module.config.server import VALID_SERVER
|
|
|
|
|
|
class Button(Resource):
|
|
def __init__(self, area, color, button, file=None, name=None):
|
|
"""Initialize a Button instance.
|
|
|
|
Args:
|
|
area (dict[tuple], tuple): Area that the button would appear on the image.
|
|
(upper_left_x, upper_left_y, bottom_right_x, bottom_right_y)
|
|
color (dict[tuple], tuple): Color we expect the area would be.
|
|
(r, g, b)
|
|
button (dict[tuple], tuple): Area to be click if button appears on the image.
|
|
(upper_left_x, upper_left_y, bottom_right_x, bottom_right_y)
|
|
If tuple is empty, this object can be use as a checker.
|
|
Examples:
|
|
BATTLE_PREPARATION = Button(
|
|
area=(1562, 908, 1864, 1003),
|
|
color=(231, 181, 90),
|
|
button=(1562, 908, 1864, 1003)
|
|
)
|
|
"""
|
|
self.raw_area = area
|
|
self.raw_color = color
|
|
self.raw_button = button
|
|
self.raw_file = file
|
|
self.raw_name = name
|
|
|
|
self._button_offset = None
|
|
self._match_init = False
|
|
self._match_binary_init = False
|
|
self._match_luma_init = False
|
|
self.image = None
|
|
self.image_binary = None
|
|
self.image_luma = None
|
|
|
|
if self.file:
|
|
self.resource_add(key=self.file)
|
|
|
|
cached = ['area', 'color', '_button', 'file', 'name', 'is_gif']
|
|
|
|
@cached_property
|
|
def area(self):
|
|
return self.parse_property(self.raw_area)
|
|
|
|
@cached_property
|
|
def color(self):
|
|
return self.parse_property(self.raw_color)
|
|
|
|
@cached_property
|
|
def _button(self):
|
|
return self.parse_property(self.raw_button)
|
|
|
|
@cached_property
|
|
def file(self):
|
|
return self.parse_property(self.raw_file)
|
|
|
|
@cached_property
|
|
def name(self):
|
|
if self.raw_name:
|
|
return self.raw_name
|
|
elif self.file:
|
|
return os.path.splitext(os.path.split(self.file)[1])[0]
|
|
else:
|
|
return 'BUTTON'
|
|
|
|
@cached_property
|
|
def is_gif(self):
|
|
if self.file:
|
|
return os.path.splitext(self.file)[1] == '.gif'
|
|
else:
|
|
return False
|
|
|
|
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
|
|
|
|
@property
|
|
def button(self):
|
|
if self._button_offset is None:
|
|
return self._button
|
|
else:
|
|
return self._button_offset
|
|
|
|
def appear_on(self, image, threshold=10):
|
|
"""Check if the button appears on the image.
|
|
|
|
Args:
|
|
image (np.ndarray): Screenshot.
|
|
threshold (int): Default to 10.
|
|
|
|
Returns:
|
|
bool: True if button appears on screenshot.
|
|
"""
|
|
return color_similar(
|
|
color1=get_color(image, self.area),
|
|
color2=self.color,
|
|
threshold=threshold
|
|
)
|
|
|
|
def load_color(self, image):
|
|
"""Load color from the specific area of the given image.
|
|
This method is irreversible, this would be only use in some special occasion.
|
|
|
|
Args:
|
|
image: Another screenshot.
|
|
|
|
Returns:
|
|
tuple: Color (r, g, b).
|
|
"""
|
|
self.__dict__['color'] = get_color(image, self.area)
|
|
self.image = crop(image, self.area)
|
|
self.__dict__['is_gif'] = False
|
|
return self.color
|
|
|
|
def load_offset(self, button):
|
|
"""
|
|
Load offset from another button.
|
|
|
|
Args:
|
|
button (Button):
|
|
"""
|
|
offset = np.subtract(button.button, button._button)[:2]
|
|
self._button_offset = area_offset(self._button, offset=offset)
|
|
|
|
def clear_offset(self):
|
|
self._button_offset = None
|
|
|
|
def ensure_template(self):
|
|
"""
|
|
Load asset image.
|
|
If needs to call self.match, call this first.
|
|
"""
|
|
if not self._match_init:
|
|
if self.is_gif:
|
|
self.image = []
|
|
for image in imageio.mimread(self.file):
|
|
image = image[:, :, :3].copy() if len(image.shape) == 3 else image
|
|
image = crop(image, self.area)
|
|
self.image.append(image)
|
|
else:
|
|
self.image = load_image(self.file, self.area)
|
|
self._match_init = True
|
|
|
|
def ensure_binary_template(self):
|
|
"""
|
|
Load asset image.
|
|
If needs to call self.match, call this first.
|
|
"""
|
|
if not self._match_binary_init:
|
|
if self.is_gif:
|
|
self.image_binary = []
|
|
for image in self.image:
|
|
image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
_, image_binary = cv2.threshold(image_gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
|
|
self.image_binary.append(image_binary)
|
|
else:
|
|
image_gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)
|
|
_, self.image_binary = cv2.threshold(image_gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
|
|
self._match_binary_init = True
|
|
|
|
def ensure_luma_template(self):
|
|
if not self._match_luma_init:
|
|
if self.is_gif:
|
|
self.image_luma = []
|
|
for image in self.image:
|
|
luma = rgb2luma(image)
|
|
self.image_luma.append(luma)
|
|
else:
|
|
self.image_luma = rgb2luma(self.image)
|
|
self._match_luma_init = True
|
|
|
|
def resource_release(self):
|
|
super().resource_release()
|
|
self.image = None
|
|
self.image_binary = None
|
|
self.image_luma = None
|
|
self._match_init = False
|
|
self._match_binary_init = False
|
|
self._match_luma_init = False
|
|
|
|
def match(self, image, offset=30, threshold=0.85):
|
|
"""Detects button by template matching. To Some button, its location may not be static.
|
|
|
|
Args:
|
|
image: Screenshot.
|
|
offset (int, tuple): Detection area offset.
|
|
threshold (float): 0-1. Similarity.
|
|
|
|
Returns:
|
|
bool.
|
|
"""
|
|
self.ensure_template()
|
|
|
|
if isinstance(offset, tuple):
|
|
if len(offset) == 2:
|
|
offset = np.array((-offset[0], -offset[1], offset[0], offset[1]))
|
|
else:
|
|
offset = np.array(offset)
|
|
else:
|
|
offset = np.array((-3, -offset, 3, offset))
|
|
image = crop(image, offset + self.area)
|
|
|
|
if self.is_gif:
|
|
for template in self.image:
|
|
res = cv2.matchTemplate(template, image, cv2.TM_CCOEFF_NORMED)
|
|
_, similarity, _, point = cv2.minMaxLoc(res)
|
|
self._button_offset = area_offset(self._button, offset[:2] + np.array(point))
|
|
if similarity > threshold:
|
|
return True
|
|
return False
|
|
else:
|
|
res = cv2.matchTemplate(self.image, image, cv2.TM_CCOEFF_NORMED)
|
|
_, similarity, _, point = cv2.minMaxLoc(res)
|
|
self._button_offset = area_offset(self._button, offset[:2] + np.array(point))
|
|
return similarity > threshold
|
|
|
|
def match_binary(self, image, offset=30, threshold=0.85):
|
|
"""Detects button by template matching. To Some button, its location may not be static.
|
|
This method will apply template matching under binarization.
|
|
|
|
Args:
|
|
image: Screenshot.
|
|
offset (int, tuple): Detection area offset.
|
|
threshold (float): 0-1. Similarity.
|
|
|
|
Returns:
|
|
bool.
|
|
"""
|
|
self.ensure_template()
|
|
self.ensure_binary_template()
|
|
|
|
if isinstance(offset, tuple):
|
|
if len(offset) == 2:
|
|
offset = np.array((-offset[0], -offset[1], offset[0], offset[1]))
|
|
else:
|
|
offset = np.array(offset)
|
|
else:
|
|
offset = np.array((-3, -offset, 3, offset))
|
|
image = crop(image, offset + self.area)
|
|
|
|
if self.is_gif:
|
|
for template in self.image_binary:
|
|
# graying
|
|
image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
# binarization
|
|
_, image_binary = cv2.threshold(image_gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
|
|
# template matching
|
|
res = cv2.matchTemplate(template, image_binary, cv2.TM_CCOEFF_NORMED)
|
|
_, similarity, _, point = cv2.minMaxLoc(res)
|
|
self._button_offset = area_offset(self._button, offset[:2] + np.array(point))
|
|
if similarity > threshold:
|
|
return True
|
|
return False
|
|
else:
|
|
# graying
|
|
image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
# binarization
|
|
_, image_binary = cv2.threshold(image_gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
|
|
# template matching
|
|
res = cv2.matchTemplate(self.image_binary, image_binary, cv2.TM_CCOEFF_NORMED)
|
|
_, similarity, _, point = cv2.minMaxLoc(res)
|
|
self._button_offset = area_offset(self._button, offset[:2] + np.array(point))
|
|
return similarity > threshold
|
|
|
|
def match_luma(self, image, offset=30, threshold=0.85):
|
|
"""
|
|
Detects button by template matching under Y channel (Luminance)
|
|
|
|
Args:
|
|
image: Screenshot.
|
|
offset (int, tuple): Detection area offset.
|
|
threshold (float): 0-1. Similarity.
|
|
|
|
Returns:
|
|
bool.
|
|
"""
|
|
self.ensure_template()
|
|
self.ensure_luma_template()
|
|
|
|
if isinstance(offset, tuple):
|
|
if len(offset) == 2:
|
|
offset = np.array((-offset[0], -offset[1], offset[0], offset[1]))
|
|
else:
|
|
offset = np.array(offset)
|
|
else:
|
|
offset = np.array((-3, -offset, 3, offset))
|
|
image = crop(image, offset + self.area)
|
|
|
|
if self.is_gif:
|
|
image_luma = rgb2luma(image)
|
|
for template in self.image_luma:
|
|
res = cv2.matchTemplate(template, image_luma, cv2.TM_CCOEFF_NORMED)
|
|
_, similarity, _, point = cv2.minMaxLoc(res)
|
|
self._button_offset = area_offset(self._button, offset[:2] + np.array(point))
|
|
if similarity > threshold:
|
|
return True
|
|
else:
|
|
image_luma = rgb2luma(image)
|
|
res = cv2.matchTemplate(self.image_luma, image_luma, cv2.TM_CCOEFF_NORMED)
|
|
_, similarity, _, point = cv2.minMaxLoc(res)
|
|
self._button_offset = area_offset(self._button, offset[:2] + np.array(point))
|
|
return similarity > threshold
|
|
|
|
def match_appear_on(self, image, threshold=30):
|
|
"""
|
|
Args:
|
|
image: Screenshot.
|
|
threshold: Default to 10.
|
|
|
|
Returns:
|
|
bool:
|
|
"""
|
|
diff = np.subtract(self.button, self._button)[:2]
|
|
area = area_offset(self.area, offset=diff)
|
|
return color_similar(color1=get_color(image, area), color2=self.color, threshold=threshold)
|
|
|
|
def crop(self, area, image=None, name=None):
|
|
"""
|
|
Get a new button by relative coordinates.
|
|
|
|
Args:
|
|
area (tuple):
|
|
image (np.ndarray): Screenshot. If provided, load color and image from it.
|
|
name (str):
|
|
|
|
Returns:
|
|
Button:
|
|
"""
|
|
if name is None:
|
|
name = self.name
|
|
new_area = area_offset(area, offset=self.area[:2])
|
|
new_button = area_offset(area, offset=self.button[:2])
|
|
button = Button(area=new_area, color=self.color, button=new_button, file=self.file, name=name)
|
|
if image is not None:
|
|
button.load_color(image)
|
|
return button
|
|
|
|
def move(self, vector, image=None, name=None):
|
|
"""
|
|
Move button.
|
|
|
|
Args:
|
|
vector (tuple):
|
|
image (np.ndarray): Screenshot. If provided, load color and image from it.
|
|
name (str):
|
|
|
|
Returns:
|
|
Button:
|
|
"""
|
|
if name is None:
|
|
name = self.name
|
|
new_area = area_offset(self.area, offset=vector)
|
|
new_button = area_offset(self.button, offset=vector)
|
|
button = Button(area=new_area, color=self.color, button=new_button, file=self.file, name=name)
|
|
if image is not None:
|
|
button.load_color(image)
|
|
return button
|
|
|
|
def split_server(self):
|
|
"""
|
|
Split into 4 server specific buttons.
|
|
|
|
Returns:
|
|
dict[str, Button]:
|
|
"""
|
|
out = {}
|
|
for s in VALID_SERVER:
|
|
out[s] = Button(
|
|
area=self.parse_property(self.raw_area, s),
|
|
color=self.parse_property(self.raw_color, s),
|
|
button=self.parse_property(self.raw_button, s),
|
|
file=self.parse_property(self.raw_file, s),
|
|
name=self.name
|
|
)
|
|
return out
|
|
|
|
|
|
class ButtonGrid:
|
|
def __init__(self, origin, delta, button_shape, grid_shape, name=None):
|
|
self.origin = np.array(origin)
|
|
self.delta = np.array(delta)
|
|
self.button_shape = np.array(button_shape)
|
|
self.grid_shape = np.array(grid_shape)
|
|
if name:
|
|
self._name = name
|
|
else:
|
|
(filename, line_number, function_name, text) = traceback.extract_stack()[-2]
|
|
self._name = text[:text.find('=')].strip()
|
|
|
|
def __getitem__(self, item):
|
|
base = np.round(np.array(item) * self.delta + self.origin).astype(int)
|
|
area = tuple(np.append(base, base + self.button_shape))
|
|
return Button(area=area, color=(), button=area, name='%s_%s_%s' % (self._name, item[0], item[1]))
|
|
|
|
def generate(self):
|
|
for y in range(self.grid_shape[1]):
|
|
for x in range(self.grid_shape[0]):
|
|
yield x, y, self[x, y]
|
|
|
|
@cached_property
|
|
def buttons(self):
|
|
return list([button for _, _, button in self.generate()])
|
|
|
|
def crop(self, area, name=None):
|
|
"""
|
|
Args:
|
|
area (tuple): Area related to self.origin
|
|
name (str): Name of the new ButtonGrid instance.
|
|
|
|
Returns:
|
|
ButtonGrid:
|
|
"""
|
|
if name is None:
|
|
name = self._name
|
|
origin = self.origin + area[:2]
|
|
button_shape = np.subtract(area[2:], area[:2])
|
|
return ButtonGrid(
|
|
origin=origin, delta=self.delta, button_shape=button_shape, grid_shape=self.grid_shape, name=name)
|
|
|
|
def move(self, vector, name=None):
|
|
"""
|
|
Args:
|
|
vector (tuple): Move vector.
|
|
name (str): Name of the new ButtonGrid instance.
|
|
|
|
Returns:
|
|
ButtonGrid:
|
|
"""
|
|
if name is None:
|
|
name = self._name
|
|
origin = self.origin + vector
|
|
return ButtonGrid(
|
|
origin=origin, delta=self.delta, button_shape=self.button_shape, grid_shape=self.grid_shape, name=name)
|
|
|
|
def gen_mask(self):
|
|
"""
|
|
Generate a mask image to display this ButtonGrid object for debugging.
|
|
|
|
Returns:
|
|
PIL.Image.Image: Area in white, background in black.
|
|
"""
|
|
image = Image.new("RGB", (1280, 720), (0, 0, 0))
|
|
draw = ImageDraw.Draw(image)
|
|
for button in self.buttons:
|
|
draw.rectangle((button.area[:2], button.button[2:]), fill=(255, 255, 255), outline=None)
|
|
return image
|
|
|
|
def show_mask(self):
|
|
self.gen_mask().show()
|
|
|
|
def save_mask(self):
|
|
"""
|
|
Save mask to {name}.png
|
|
"""
|
|
self.gen_mask().save(f'{self._name}.png')
|