323 lines
13 KiB
Python
323 lines
13 KiB
Python
import cv2
|
|
import threading
|
|
from screen import convert_screen_to_monitor, grab
|
|
from dataclasses import dataclass
|
|
import numpy as np
|
|
from logger import Logger
|
|
import time
|
|
import os
|
|
from config import Config
|
|
from utils.misc import cut_roi, load_template, list_files_in_folder, alpha_to_mask, roi_center, color_filter, mask_by_roi
|
|
from functools import cache
|
|
|
|
templates_lock = threading.Lock()
|
|
|
|
@dataclass
|
|
class Template:
|
|
name: str = None
|
|
img_bgra: np.ndarray = None
|
|
img_bgr: np.ndarray = None
|
|
img_gray: np.ndarray = None
|
|
alpha_mask: np.ndarray = None
|
|
|
|
@dataclass
|
|
class TemplateMatch:
|
|
name: str = None
|
|
score: float = -1.0
|
|
center: tuple[int, int] = None
|
|
center_monitor: tuple[int, int] = None
|
|
region: list[float] = None
|
|
region_monitor: list[float] = None
|
|
valid: bool = False
|
|
|
|
|
|
TEMPLATE_PATHS = [
|
|
"assets\\templates",
|
|
"assets\\npc",
|
|
"assets\\shop",
|
|
"assets\\item_properties",
|
|
"assets\\chests",
|
|
"assets\\gamble",
|
|
]
|
|
|
|
@cache
|
|
def stored_templates() -> dict[Template]:
|
|
paths = []
|
|
templates = {}
|
|
for path in TEMPLATE_PATHS:
|
|
paths += list_files_in_folder(path)
|
|
for file_path in paths:
|
|
file_name: str = os.path.basename(file_path)
|
|
if file_name.lower().endswith('.png'):
|
|
key = file_name[:-4].upper()
|
|
template_img = load_template(file_path)
|
|
templates[key] = Template(
|
|
name = key,
|
|
img_bgra = template_img,
|
|
img_bgr = cv2.cvtColor(template_img, cv2.COLOR_BGRA2BGR),
|
|
img_gray = cv2.cvtColor(template_img, cv2.COLOR_BGRA2GRAY),
|
|
alpha_mask = alpha_to_mask(template_img)
|
|
)
|
|
return templates
|
|
|
|
def get_template(key):
|
|
with templates_lock:
|
|
return stored_templates()[key].img_bgr
|
|
|
|
def _process_template_refs(ref: str | np.ndarray | list[str]) -> list[Template]:
|
|
templates = []
|
|
if type(ref) != list:
|
|
ref = [ref]
|
|
for i in ref:
|
|
# if the reference is a string, then it's a reference to a named template asset
|
|
if type(i) == str:
|
|
templates.append(stored_templates()[i.upper()])
|
|
# if the reference is an image, append new Template class object
|
|
elif type(i) == np.ndarray:
|
|
templates.append(Template(
|
|
img_bgr = i,
|
|
img_gray = cv2.cvtColor(i, cv2.COLOR_BGR2GRAY),
|
|
alpha_mask = alpha_to_mask(i)
|
|
))
|
|
return templates
|
|
|
|
def _single_template_match(template: Template, inp_img: np.ndarray = None, roi: list = None, color_match: list = None, use_grayscale: bool = False) -> TemplateMatch:
|
|
inp_img = inp_img if inp_img is not None else grab()
|
|
template_match = TemplateMatch()
|
|
|
|
# crop image to roi
|
|
if roi is None:
|
|
# if no roi is provided roi = full inp_img
|
|
roi = [0, 0, inp_img.shape[1], inp_img.shape[0]]
|
|
rx, ry, rw, rh = roi
|
|
img = inp_img[ry:ry + rh, rx:rx + rw]
|
|
|
|
# filter for desired color or make grayscale
|
|
if color_match:
|
|
template_img, = color_filter(template.img_bgr, color_match)[1],
|
|
img = color_filter(img, color_match)[1]
|
|
elif use_grayscale:
|
|
template_img = template.img_gray
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
else:
|
|
template_img = template.img_bgr
|
|
|
|
if not (img.shape[0] > template_img.shape[0] and img.shape[1] > template_img.shape[1]):
|
|
Logger.error(f"Image shape and template shape are incompatible: {template.name}. Image: {img.shape}, Template: {template_img.shape}, roi: {roi}")
|
|
else:
|
|
res = cv2.matchTemplate(img, template_img, cv2.TM_CCOEFF_NORMED, mask = template.alpha_mask)
|
|
np.nan_to_num(res, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
|
|
_, max_val, _, max_pos = cv2.minMaxLoc(res)
|
|
|
|
# save rectangle corresponding to matched region
|
|
rec_x = int((max_pos[0] + rx))
|
|
rec_y = int((max_pos[1] + ry))
|
|
rec_w = int(template_img.shape[1])
|
|
rec_h = int(template_img.shape[0])
|
|
template_match.region = [rec_x, rec_y, rec_w, rec_h]
|
|
template_match.region_monitor = [*convert_screen_to_monitor((rec_x, rec_y)), rec_w, rec_h]
|
|
template_match.center = roi_center(template_match.region)
|
|
template_match.center_monitor = convert_screen_to_monitor(template_match.center)
|
|
template_match.name = template.name
|
|
template_match.score = max_val
|
|
template_match.valid = True
|
|
|
|
return template_match
|
|
|
|
|
|
def search(
|
|
ref: str | np.ndarray | list[str],
|
|
inp_img: np.ndarray,
|
|
threshold: float = 0.68,
|
|
roi: list[float] = None,
|
|
use_grayscale: bool = False,
|
|
color_match: list = False,
|
|
best_match: bool = False
|
|
) -> TemplateMatch:
|
|
"""
|
|
Search for a template in an image
|
|
:param ref: Either key of a already loaded template, list of such keys, or a image which is used as template
|
|
:param inp_img: Image in which the template will be searched
|
|
:param threshold: Threshold which determines if a template is found or not
|
|
:param roi: Region of Interest of the inp_img to restrict search area. Format [left, top, width, height]
|
|
:param use_grayscale: Use grayscale template matching for speed up
|
|
:param color_match: Pass a color to be used by misc.color_filter to filter both image of interest and template image (format Config().colors["color"])
|
|
:param best_match: If list input, will search for list of templates by best match. Default behavior is first match.
|
|
:return: Returns a TemplateMatch object with a valid flag
|
|
"""
|
|
templates = _process_template_refs(ref)
|
|
matches = []
|
|
for template in templates:
|
|
match = _single_template_match(template, inp_img, roi, color_match, use_grayscale)
|
|
if match.score >= threshold:
|
|
if not best_match:
|
|
return match
|
|
else:
|
|
matches.append(match)
|
|
if matches:
|
|
matches = sorted(matches, key=lambda obj: obj.score, reverse=True)
|
|
return matches[0]
|
|
return TemplateMatch()
|
|
|
|
|
|
def search_and_wait(
|
|
ref: str | list[str],
|
|
roi: list[float] = None,
|
|
timeout: float = 30,
|
|
threshold: float = 0.68,
|
|
use_grayscale: bool = False,
|
|
color_match: list = False,
|
|
best_match: bool = False,
|
|
suppress_debug: bool = False,
|
|
) -> TemplateMatch:
|
|
"""
|
|
Helper function that will loop and keep searching for a template
|
|
:param timeout: After this amount of time the search will stop and it will return [False, None]
|
|
:Other params are the same as for template_finder.search()
|
|
:returns a TemplateMatch object
|
|
"""
|
|
if not suppress_debug:
|
|
Logger.debug(f"Waiting for templates: {ref}")
|
|
start = time.time()
|
|
template_match = TemplateMatch()
|
|
while (time_remains := time.time() - start < timeout):
|
|
img = grab()
|
|
is_loading_black_roi = np.average(img[:, 0:Config().ui_roi["loading_left_black"][2]]) < 1.0
|
|
if not is_loading_black_roi or "LOADING" in ref:
|
|
template_match = search(ref, img, roi=roi, threshold=threshold, use_grayscale=use_grayscale, color_match=color_match, best_match=best_match)
|
|
if template_match.valid:
|
|
break
|
|
if not time_remains:
|
|
Logger.debug(f"Could not find desired templates")
|
|
else:
|
|
Logger.debug(f"Found match: {template_match.name} ({template_match.score*100:.1f}% confidence)")
|
|
return template_match
|
|
|
|
|
|
def search_all(
|
|
ref: str | np.ndarray | list[str],
|
|
inp_img: np.ndarray,
|
|
threshold: float = 0.68,
|
|
roi: list[float] = None,
|
|
use_grayscale: bool = False,
|
|
color_match: list = False,
|
|
) -> list[TemplateMatch]:
|
|
"""
|
|
Returns a list of all templates scoring above set threshold on the screen
|
|
:Other params are the same as for template_finder.search()
|
|
:return: Returns a list of TemplateMatch objects
|
|
"""
|
|
templates = _process_template_refs(ref)
|
|
matches = []
|
|
img = inp_img.copy()
|
|
while True:
|
|
any_found = False
|
|
for template in templates:
|
|
match = _single_template_match(template, img, roi, color_match, use_grayscale)
|
|
if (ind_found := match.score >= threshold):
|
|
matches.append(match)
|
|
img = mask_by_roi(img, match.region, "inverse")
|
|
any_found |= ind_found
|
|
if not any_found:
|
|
break
|
|
return matches
|
|
|
|
|
|
# Testing: Have whatever you want to find on the screen
|
|
if __name__ == "__main__":
|
|
import keyboard
|
|
import os
|
|
from screen import start_detecting_window, stop_detecting_window
|
|
from utils.misc import wait
|
|
import template_finder
|
|
start_detecting_window()
|
|
wait(0.1)
|
|
|
|
print("\n== Hotkeys ==")
|
|
print("F11: Start")
|
|
print("F12: Exit")
|
|
print("Down arrow: decrease template matching threshold")
|
|
print("Up arrow: increase template matching threshold")
|
|
print("Left arrow: decrease template index")
|
|
print("Right arrow: increase template index")
|
|
print("F9: toggle all vs. individual template(s)")
|
|
print("F10: save visible template(s)")
|
|
|
|
keyboard.add_hotkey('f12', lambda: Logger.info('Force Exit (f12)') or stop_detecting_window() or os._exit(1))
|
|
keyboard.wait("f11")
|
|
|
|
# enter the template names you are trying to detect here
|
|
|
|
_template_list = ["SHENK_0","SHENK_1","SHENK_10","SHENK_11","SHENK_12","SHENK_13","SHENK_15","SHENK_16","SHENK_17","SHENK_18","SHENK_19","SHENK_2","SHENK_20","SHENK_3","SHENK_4","SHENK_6","SHENK_7","SHENK_8","SHENK_9","SHENK_DEATH_0","SHENK_DEATH_1","SHENK_DEATH_2","SHENK_DEATH_3","SHENK_DEATH_4","SHENK_V2_3","SHENK_V2_4","SHENK_V2_6","SHENK_V2_7","SHENK_V2_8"]
|
|
|
|
_template_list += ["ELDRITCH_0","ELDRITCH_0_V2","ELDRITCH_0_V3","ELDRITCH_1","ELDRITCH_1_V2","ELDRITCH_2","ELDRITCH_2_V2","ELDRITCH_3","ELDRITCH_4","ELDRITCH_5","ELDRITCH_6","ELDRITCH_7","ELDRITCH_7_V2","ELDRITCH_8","ELDRITCH_8_V2","ELDRITCH_9","ELDRITCH_START","ELDRITCH_START_V2"]
|
|
|
|
_current_template_idx = -1
|
|
_last_stored_idx = 0
|
|
_current_threshold = 0.6
|
|
_visible_templates = []
|
|
|
|
def _save_visible_templates():
|
|
os.makedirs("log/screenshots/info", exist_ok=True)
|
|
for match in _visible_templates:
|
|
cv2.imwrite(match['filename'], match['img'])
|
|
Logger.info(f"{match['filename']} saved")
|
|
|
|
def _toggle_templates():
|
|
global _current_template_idx
|
|
_current_template_idx = -1 if _current_template_idx != -1 else _last_stored_idx
|
|
if _current_template_idx == -1:
|
|
Logger.info(f"Searching for templates: {_template_list}")
|
|
else:
|
|
Logger.info(f"Searching for template: {_template_list[_current_template_idx]}")
|
|
|
|
def _incr_template_idx(direction: int = 1):
|
|
global _current_template_idx, _last_stored_idx
|
|
if \
|
|
(-1 < _current_template_idx < len(_template_list) - 1) or \
|
|
(_current_template_idx == -1 and direction > 0) or \
|
|
(_current_template_idx == len(_template_list) - 1 and direction < 0):
|
|
_current_template_idx += direction
|
|
_last_stored_idx = _current_template_idx
|
|
if _current_template_idx == -1:
|
|
Logger.info(f"Searching for templates: {_template_list}")
|
|
else:
|
|
Logger.info(f"Searching for template: {_template_list[_current_template_idx]}")
|
|
|
|
def _incr_threshold(direction: int = 1):
|
|
global _current_threshold
|
|
if (_current_threshold < 1 and direction > 0) or (_current_threshold > 0 and direction < 0):
|
|
_current_threshold = round(_current_threshold + direction, 2)
|
|
Logger.info(f"_current_threshold = {_current_threshold}")
|
|
|
|
keyboard.add_hotkey('down', lambda: _incr_threshold(-0.05))
|
|
keyboard.add_hotkey('up', lambda: _incr_threshold(0.05))
|
|
keyboard.add_hotkey('left', lambda: _incr_template_idx(-1))
|
|
keyboard.add_hotkey('right', lambda: _incr_template_idx(1))
|
|
keyboard.add_hotkey('f9', lambda: _toggle_templates())
|
|
keyboard.add_hotkey('f10', lambda: _save_visible_templates())
|
|
|
|
while 1:
|
|
_visible_templates = []
|
|
img = grab()
|
|
display_img = img.copy()
|
|
if _current_template_idx < 0:
|
|
templates = _template_list
|
|
else:
|
|
templates = [_template_list[_current_template_idx]]
|
|
for key in templates:
|
|
template_match = template_finder.search(key, img, threshold=_current_threshold)
|
|
if template_match.valid:
|
|
x, y = template_match.center
|
|
cv2.putText(display_img, str(template_match.name), template_match.center, cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
|
|
cv2.circle(display_img, template_match.center, 7, (255, 0, 0), thickness=5)
|
|
cv2.rectangle(display_img, template_match.region[:2], (template_match.region[0] + template_match.region[2], template_match.region[1] + template_match.region[3]), (0, 0, 255), 1)
|
|
print(f"Name: {template_match.name} Pos: {template_match.center}, Dist: {625-x, 360-y}, Score: {template_match.score}")
|
|
match = {
|
|
'filename': f"./log/screenshots/info/{key.lower()}.png",
|
|
'img': cut_roi(img, template_match.region)
|
|
}
|
|
_visible_templates.append(match)
|
|
cv2.imshow('test', display_img)
|
|
key = cv2.waitKey(3000)
|