159 lines
7.6 KiB
Python
159 lines
7.6 KiB
Python
from screen import grab, convert_monitor_to_screen, convert_abs_to_screen, convert_screen_to_abs
|
|
import cv2
|
|
from config import Config
|
|
import template_finder
|
|
from utils.misc import load_template, cut_roi
|
|
import mouse
|
|
from input_layer import keyboard
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
import math
|
|
|
|
|
|
class NodeRecorder:
|
|
def __init__(self, run_name):
|
|
if os.path.exists("log/screenshots/generated"):
|
|
for path in Path("log/screenshots/generated").glob("**/*"):
|
|
if path.is_file():
|
|
os.remove(path)
|
|
elif path.is_dir():
|
|
shutil.rmtree(path)
|
|
shutil.rmtree("log/screenshots/generated")
|
|
os.makedirs("log/screenshots/generated", exist_ok=True)
|
|
os.system(f"cd log/screenshots/generated && mkdir templates && cd templates && mkdir {run_name} && cd {run_name} && mkdir nodes")
|
|
self._run_name = run_name
|
|
self._offset = 100
|
|
self._template_counter = 0
|
|
self._half_width = Config().ui_pos["screen_width"] // 2
|
|
self._half_height = Config().ui_pos["screen_height"] // 2
|
|
self._curr_state = 0
|
|
self._upper_left = None
|
|
template_finder._templates = {}
|
|
self._pather_code_file = "log/screenshots/generated/pather_generated.py"
|
|
self.ref_points = {}
|
|
self.nodes = {}
|
|
self.debug_node_pos = {}
|
|
# Starting with template recording:
|
|
print("1 - Select top-left corner and press f8")
|
|
|
|
@staticmethod
|
|
def _convert_rel_to_abs(rel_loc: tuple[float, float], pos_abs: tuple[float, float]) -> tuple[float, float]:
|
|
return (rel_loc[0] + pos_abs[0], rel_loc[1] + pos_abs[1])
|
|
|
|
def find_templates(self, img):
|
|
ref_points = {}
|
|
for key in template_finder._templates:
|
|
found = template_finder.search(key, img, use_grayscale=False, threshold=0.77)
|
|
if found.valid:
|
|
ref_points[key] = found.center
|
|
return ref_points
|
|
|
|
def hook(self, e):
|
|
if e.event_type == "down":
|
|
if e.name == "f12":
|
|
os._exit(1)
|
|
self.ref_points = {}
|
|
self.debug_node_pos = {}
|
|
img = grab()
|
|
loc_monitor = mouse.get_position()
|
|
loc_screen = convert_monitor_to_screen(loc_monitor)
|
|
if e.name == "f8" and self._curr_state == 0:
|
|
# create a tempalte
|
|
if self._upper_left is None:
|
|
self._upper_left = loc_screen
|
|
print("-- Select bottom-right corner and press f8")
|
|
return
|
|
else:
|
|
bottom_right = loc_screen
|
|
width = (bottom_right[0] - self._upper_left[0])
|
|
height = (bottom_right[1] - self._upper_left[1])
|
|
ref_point_name = f"{self._run_name}_{self._template_counter}"
|
|
self._template_counter += 1
|
|
# save as png
|
|
template_img = cut_roi(img, [*self._upper_left, width, height])
|
|
template_path = f"log/screenshots/generated/templates/{self._run_name}/{ref_point_name}.png"
|
|
cv2.imwrite(template_path, template_img)
|
|
self._upper_left = None
|
|
template_img = load_template(template_path)
|
|
template_finder._templates[ref_point_name] = [template_img, cv2.cvtColor(template_img, cv2.COLOR_BGRA2GRAY), 1.0, None]
|
|
elif e.name == "f7":
|
|
self.ref_points = {}
|
|
else:
|
|
self.ref_points = self.find_templates(img)
|
|
|
|
if e.name == "f9":
|
|
# add current loc screen as node
|
|
new_node_idx = self._offset + len(self.nodes)
|
|
self.nodes[new_node_idx] = {}
|
|
for key in self.ref_points:
|
|
rel_loc = (loc_screen[0] - self.ref_points[key][0], loc_screen[1] - self.ref_points[key][1])
|
|
self.nodes[new_node_idx][key] = rel_loc
|
|
|
|
if e.name == "f10" or e.name == "f9" and self.ref_points is not None:
|
|
# find all nodes:
|
|
for node_idx in self.nodes:
|
|
# try to find the screen coordinate of the node
|
|
node_screen_pos = None
|
|
for template_key in self.nodes[node_idx]:
|
|
if template_key in self.ref_points:
|
|
ref_pos_screen = self.ref_points[template_key]
|
|
# Get reference position of template in abs coordinates
|
|
ref_pos_abs = convert_screen_to_abs(ref_pos_screen)
|
|
# Calc the abs node position with the relative coordinates (relative to ref)
|
|
node_pos_rel = self.nodes[node_idx][template_key]
|
|
node_pos_abs = self._convert_rel_to_abs(node_pos_rel, ref_pos_abs)
|
|
node_screen_pos = convert_abs_to_screen(node_pos_abs)
|
|
self.debug_node_pos[node_idx] = node_screen_pos
|
|
break
|
|
# if it was found try to add all other visible templates to it that are not already included
|
|
if node_screen_pos is not None:
|
|
for template_key in self.ref_points:
|
|
if template_key not in self.nodes[node_idx]:
|
|
rel_loc = (node_screen_pos[0] - self.ref_points[template_key][0], node_screen_pos[1] - self.ref_points[template_key][1])
|
|
self.nodes[node_idx][template_key] = rel_loc
|
|
# print info to console and write code
|
|
f = open(self._pather_code_file, 'w')
|
|
print("---- Current Recorded Nodes: ----")
|
|
for k in self.nodes:
|
|
print(self.nodes[k])
|
|
new_path = []
|
|
for template in self.nodes[k]:
|
|
dist = math.dist((0, 0), self.nodes[k][template])
|
|
new_path.append({"key": template, "pos": self.nodes[k][template], "dist": dist})
|
|
results = sorted(new_path, key=lambda r: r["dist"])
|
|
code = f"{k}: " + "{"
|
|
for i, res in enumerate(results):
|
|
if res["dist"] < 1100 and i < 8:
|
|
code += f'"{res["key"].upper()}": {res["pos"]}, '
|
|
f.write(code + "}\n")
|
|
f.close()
|
|
print("")
|
|
print("f8: Create Template | f9: New node at cursor | f10: update nodes with visible templates")
|
|
|
|
if __name__ == "__main__":
|
|
keyboard.add_hotkey('f12', lambda: print('Force Exit (f12)') or os._exit(1))
|
|
print("Enter run name...")
|
|
run_name = input()
|
|
from screen import stop_detecting_window, start_detecting_window
|
|
start_detecting_window()
|
|
|
|
recorder = NodeRecorder(run_name)
|
|
keyboard.hook(recorder.hook, suppress=True)
|
|
|
|
while 1:
|
|
img = grab().copy()
|
|
try:
|
|
for key in recorder.debug_node_pos:
|
|
cv2.circle(img, recorder.debug_node_pos[key], 8, (0, 0, 255), 4)
|
|
cv2.putText(img, str(key), recorder.debug_node_pos[key], cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
|
|
for key in recorder.ref_points:
|
|
cv2.circle(img, recorder.ref_points[key], 8, (0, 255, 0), 4)
|
|
cv2.putText(img, key, recorder.ref_points[key], cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
|
|
except Exception:
|
|
pass
|
|
img = cv2.resize(img, None, fx=0.5, fy=0.5)
|
|
cv2.imshow("vis", img)
|
|
cv2.waitKey(1)
|
|
|
|
stop_detecting_window() |