501 lines
20 KiB
Python
501 lines
20 KiB
Python
import cv2
|
|
import threading
|
|
import concurrent.futures
|
|
from screen import convert_screen_to_monitor, grab
|
|
from dataclasses import dataclass
|
|
import numpy as np
|
|
from logger import Logger
|
|
from utils.log_rotation import safe_imwrite
|
|
import time
|
|
import os
|
|
import re
|
|
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
|
|
|
|
|
|
from pathlib import Path
|
|
|
|
TEMPLATE_PATHS = [
|
|
str(Path(__file__).parent.parent / "assets" / "templates"),
|
|
str(Path(__file__).parent.parent / "assets" / "npc"),
|
|
str(Path(__file__).parent.parent / "assets" / "shop"),
|
|
str(Path(__file__).parent.parent / "assets" / "item_properties"),
|
|
str(Path(__file__).parent.parent / "assets" / "chests"),
|
|
str(Path(__file__).parent.parent / "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:
|
|
key = i.upper()
|
|
if key in stored_templates():
|
|
templates.append(stored_templates()[key])
|
|
else:
|
|
Logger.warning(f"Template reference does not exist: {key}")
|
|
# 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]):
|
|
from logger import Logger
|
|
Logger.debug(f"Incompatible shapes for {template.name}: Image {img.shape}, Template {template_img.shape}, ROI {roi}")
|
|
else:
|
|
try:
|
|
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
|
|
except Exception as e:
|
|
from logger import Logger
|
|
Logger.error(f"Template match failed for {template.name}: {e}")
|
|
import traceback
|
|
Logger.debug(traceback.format_exc())
|
|
|
|
return template_match
|
|
|
|
|
|
def _match_template_worker(template, inp_img, roi, color_match, use_grayscale):
|
|
"""Worker for parallel template matching (runs in thread pool)."""
|
|
return _single_template_match(template, inp_img, roi, color_match, use_grayscale)
|
|
|
|
|
|
def search(
|
|
ref,
|
|
inp_img,
|
|
threshold=0.68,
|
|
roi=None,
|
|
use_grayscale=False,
|
|
color_match=False,
|
|
best_match=False,
|
|
max_workers=4
|
|
):
|
|
"""
|
|
Search for a template in an image. Uses parallel matching for list inputs.
|
|
|
|
: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
|
|
:param best_match: If list input, will search for list of templates by best match. Default behavior is first match.
|
|
:param max_workers: Max parallel threads for template matching (default 4)
|
|
:return: Returns a TemplateMatch object with a valid flag
|
|
"""
|
|
templates = _process_template_refs(ref)
|
|
if not templates:
|
|
return TemplateMatch()
|
|
|
|
# Single template — no benefit to parallelize
|
|
if len(templates) == 1:
|
|
match = _single_template_match(templates[0], inp_img, roi, color_match, use_grayscale)
|
|
match.valid = match.score >= threshold
|
|
return match
|
|
|
|
# Multiple templates default to first valid match in input order.
|
|
if not best_match:
|
|
best = TemplateMatch()
|
|
for template in templates:
|
|
match = _single_template_match(template, inp_img, roi, color_match, use_grayscale)
|
|
match.valid = match.score >= threshold
|
|
if match.score > best.score:
|
|
best = match
|
|
if match.valid:
|
|
return match
|
|
return best
|
|
|
|
# Multiple templates with best_match=True return the highest score.
|
|
matches = []
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
futures = [
|
|
executor.submit(_match_template_worker, t, inp_img, roi, color_match, use_grayscale)
|
|
for t in templates
|
|
]
|
|
for future in concurrent.futures.as_completed(futures):
|
|
match = future.result()
|
|
match.valid = match.score >= threshold
|
|
matches.append(match)
|
|
|
|
if matches:
|
|
matches = sorted(matches, key=lambda obj: obj.score, reverse=True)
|
|
return matches[0]
|
|
return TemplateMatch()
|
|
|
|
|
|
def _should_save_missing_debug(timeout: float, suppress_debug: bool, save_debug: bool | None) -> bool:
|
|
if save_debug is not None:
|
|
return bool(save_debug)
|
|
return not suppress_debug and timeout >= 1.0 and Config().general.get("info_screenshots", 0)
|
|
|
|
|
|
def _safe_ref_label(ref) -> str:
|
|
refs = ref if isinstance(ref, list) else [ref]
|
|
label = "_".join(str(r) for r in refs[:4])
|
|
if len(refs) > 4:
|
|
label += f"_plus{len(refs) - 4}"
|
|
return re.sub(r"[^A-Za-z0-9_-]+", "_", label).strip("_")[:90] or "unknown"
|
|
|
|
|
|
def _save_missing_template_debug(ref, img, best_match: TemplateMatch, threshold: float, roi: list = None, context: str = "search"):
|
|
if img is None:
|
|
return
|
|
try:
|
|
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
|
label = _safe_ref_label(ref)
|
|
best = "none"
|
|
if best_match and best_match.name:
|
|
best = f"{best_match.name}_{best_match.score:.2f}".replace(".", "p")
|
|
base = f"./log/screenshots/info/missing_{context}_{label}_thr{threshold:.2f}_best_{best}_{timestamp}"
|
|
safe_imwrite(base + ".png", img)
|
|
if roi:
|
|
safe_imwrite(base + "_roi.png", cut_roi(img, roi))
|
|
Logger.info(f"Saved missing-template debug screenshot: {base}.png")
|
|
except Exception as e:
|
|
Logger.warning(f"Failed to save missing-template debug screenshot for {ref}: {e}")
|
|
|
|
|
|
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,
|
|
save_debug: bool = None,
|
|
) -> 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
|
|
"""
|
|
from utils.misc import should_stop
|
|
if not suppress_debug:
|
|
Logger.debug(f"Waiting for templates: {ref}")
|
|
start = time.time()
|
|
template_match = TemplateMatch()
|
|
last_img = None
|
|
while (time_remains := time.time() - start < timeout):
|
|
if should_stop():
|
|
raise SystemExit()
|
|
img = grab()
|
|
last_img = img
|
|
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")
|
|
if _should_save_missing_debug(timeout, suppress_debug, save_debug):
|
|
_save_missing_template_debug(ref, last_img, template_match, threshold, roi, "wait")
|
|
else:
|
|
Logger.debug(f"Found match: {template_match.name} ({template_match.score*100:.1f}% confidence)")
|
|
return template_match
|
|
|
|
|
|
def search_and_wait_stable(
|
|
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 = True,
|
|
confirmations: int = 2,
|
|
interval: float = 0.08,
|
|
suppress_debug: bool = False,
|
|
save_debug: bool = None,
|
|
) -> TemplateMatch:
|
|
"""
|
|
Wait until a template is visible across consecutive fresh frames.
|
|
|
|
Use this for state/location verification where a single-frame false positive
|
|
is more expensive than a short confirmation delay.
|
|
"""
|
|
from utils.misc import should_stop, wait as _wait
|
|
|
|
confirmations = max(1, int(confirmations))
|
|
if not suppress_debug:
|
|
Logger.debug(f"Waiting for stable templates: {ref} ({confirmations} confirmations)")
|
|
|
|
start = time.time()
|
|
last_name = None
|
|
stable_count = 0
|
|
best_seen = TemplateMatch()
|
|
template_match = TemplateMatch()
|
|
last_img = None
|
|
|
|
while time.time() - start < timeout:
|
|
if should_stop():
|
|
raise SystemExit()
|
|
|
|
img = grab(force_new=True)
|
|
last_img = img
|
|
refs = ref if isinstance(ref, list) else [ref]
|
|
is_loading_black_roi = np.average(img[:, 0:Config().ui_roi["loading_left_black"][2]]) < 1.0
|
|
if is_loading_black_roi and "LOADING" not in refs:
|
|
_wait(interval)
|
|
continue
|
|
|
|
template_match = search(
|
|
ref,
|
|
img,
|
|
roi=roi,
|
|
threshold=threshold,
|
|
use_grayscale=use_grayscale,
|
|
color_match=color_match,
|
|
best_match=best_match,
|
|
)
|
|
if template_match.score > best_seen.score:
|
|
best_seen = template_match
|
|
|
|
if template_match.valid:
|
|
if template_match.name == last_name:
|
|
stable_count += 1
|
|
else:
|
|
last_name = template_match.name
|
|
stable_count = 1
|
|
if stable_count >= confirmations:
|
|
if not suppress_debug:
|
|
Logger.debug(
|
|
f"Found stable match: {template_match.name} "
|
|
f"({template_match.score*100:.1f}% confidence)"
|
|
)
|
|
return template_match
|
|
else:
|
|
last_name = None
|
|
stable_count = 0
|
|
|
|
_wait(interval)
|
|
|
|
if not suppress_debug:
|
|
if best_seen.name:
|
|
Logger.debug(
|
|
f"Could not find stable templates; best was {best_seen.name} "
|
|
f"({best_seen.score*100:.1f}% confidence)"
|
|
)
|
|
else:
|
|
Logger.debug("Could not find stable templates")
|
|
if _should_save_missing_debug(timeout, suppress_debug, save_debug):
|
|
_save_missing_template_debug(ref, last_img, best_seen, threshold, roi, "stable")
|
|
best_seen.valid = False
|
|
return best_seen
|
|
|
|
|
|
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__":
|
|
from input_layer 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:
|
|
safe_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)
|