517 lines
11 KiB
Python
517 lines
11 KiB
Python
# v4.1
|
|
import json
|
|
import signal
|
|
import threading
|
|
import time
|
|
import tkinter.colorchooser as colorchooser
|
|
|
|
import customtkinter as ctk
|
|
import cv2
|
|
import dxcam
|
|
import numpy as np
|
|
|
|
from pathlib import Path
|
|
from queue import Queue, Empty
|
|
|
|
from pynput import keyboard, mouse
|
|
from pystray import Icon, Menu, MenuItem
|
|
from PIL import Image, ImageDraw
|
|
|
|
from analyze import CompositionAnalyzer
|
|
from helper import GridOverlay
|
|
|
|
CONFIG_PATH = Path("config.json")
|
|
|
|
DEFAULT_CONFIG = {
|
|
"grid_color": "#00ff66",
|
|
"line_scale": 1.0,
|
|
"mouse_toggle": True,
|
|
"grid_hotkeys": ["<ctrl>+9"],
|
|
"mode": "thirds",
|
|
"analysis_enabled": True,
|
|
"feedback_enabled": True,
|
|
"target_fps": 144,
|
|
}
|
|
|
|
ctk.set_appearance_mode("dark")
|
|
ctk.set_default_color_theme("dark-blue")
|
|
|
|
|
|
class ConfigManager:
|
|
def __init__(self):
|
|
self.config = self.load()
|
|
|
|
def load(self):
|
|
if not CONFIG_PATH.exists():
|
|
self.save(DEFAULT_CONFIG)
|
|
return DEFAULT_CONFIG.copy()
|
|
|
|
try:
|
|
data = json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
|
|
|
|
merged = DEFAULT_CONFIG.copy()
|
|
merged.update(data)
|
|
|
|
return merged
|
|
|
|
except Exception:
|
|
return DEFAULT_CONFIG.copy()
|
|
|
|
def save(self, data):
|
|
CONFIG_PATH.write_text(
|
|
json.dumps(data, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
class HotkeyManager:
|
|
def __init__(self, bindings):
|
|
self.bindings = bindings
|
|
self.listener = None
|
|
|
|
def start(self):
|
|
self.stop()
|
|
|
|
self.listener = keyboard.GlobalHotKeys(self.bindings)
|
|
|
|
self.listener.start()
|
|
|
|
def stop(self):
|
|
if self.listener:
|
|
try:
|
|
self.listener.stop()
|
|
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
class CaptureWorker:
|
|
def __init__(self):
|
|
self.camera = dxcam.create(output_idx=0)
|
|
|
|
self.frame_queue = Queue(maxsize=1)
|
|
|
|
self.running = False
|
|
|
|
def start(self):
|
|
self.running = True
|
|
|
|
threading.Thread(
|
|
target=self.capture_loop,
|
|
daemon=True,
|
|
).start()
|
|
|
|
def stop(self):
|
|
self.running = False
|
|
|
|
def capture_loop(self):
|
|
self.camera.start(
|
|
target_fps=144,
|
|
video_mode=True,
|
|
)
|
|
|
|
while self.running:
|
|
frame = self.camera.get_latest_frame()
|
|
|
|
if frame is None:
|
|
continue
|
|
|
|
try:
|
|
if self.frame_queue.full():
|
|
self.frame_queue.get_nowait()
|
|
|
|
self.frame_queue.put_nowait(frame)
|
|
|
|
except Exception:
|
|
pass
|
|
|
|
time.sleep(0.001)
|
|
|
|
|
|
class ValidationWorker:
|
|
def __init__(
|
|
self,
|
|
overlay,
|
|
capture,
|
|
):
|
|
self.overlay = overlay
|
|
self.capture = capture
|
|
|
|
self.analyzer = CompositionAnalyzer()
|
|
|
|
self.running = False
|
|
|
|
def start(self):
|
|
self.running = True
|
|
|
|
threading.Thread(
|
|
target=self.validation_loop,
|
|
daemon=True,
|
|
).start()
|
|
|
|
def stop(self):
|
|
self.running = False
|
|
|
|
def validation_loop(self):
|
|
while self.running:
|
|
try:
|
|
frame = self.capture.frame_queue.get(timeout=0.1)
|
|
|
|
except Empty:
|
|
continue
|
|
|
|
result = self.analyzer.analyze(frame)
|
|
|
|
self.overlay.update_feedback(result)
|
|
|
|
|
|
class GridApplication:
|
|
def __init__(self):
|
|
self.running = True
|
|
|
|
self.queue = Queue(maxsize=256)
|
|
|
|
self.config_manager = ConfigManager()
|
|
|
|
self.config = self.config_manager.config
|
|
|
|
self.overlay = GridOverlay(
|
|
color=self.config["grid_color"],
|
|
line_width=max(
|
|
1,
|
|
round(self.config["line_scale"] * 2),
|
|
),
|
|
mode=self.config["mode"],
|
|
)
|
|
|
|
self.capture = CaptureWorker()
|
|
|
|
self.validator = ValidationWorker(
|
|
self.overlay,
|
|
self.capture,
|
|
)
|
|
|
|
self.mouse_listener = mouse.Listener(on_click=self.on_click)
|
|
|
|
self.hotkeys = HotkeyManager(self.build_hotkeys())
|
|
|
|
self.mouse_listener.start()
|
|
|
|
self.hotkeys.start()
|
|
|
|
self.capture.start()
|
|
|
|
self.validator.start()
|
|
|
|
self.create_tray()
|
|
|
|
self.overlay.schedule(
|
|
self.process_queue,
|
|
8,
|
|
)
|
|
|
|
self.overlay.schedule(
|
|
self.overlay.schedule_monitor_refresh,
|
|
1000,
|
|
)
|
|
|
|
signal.signal(
|
|
signal.SIGINT,
|
|
self.shutdown,
|
|
)
|
|
|
|
def build_hotkeys(self):
|
|
bindings = {}
|
|
|
|
for hotkey in self.config["grid_hotkeys"]:
|
|
bindings[hotkey] = self.queue_toggle
|
|
|
|
bindings["<shift>+1"] = lambda: self.queue_color("#ff0000")
|
|
|
|
bindings["<shift>+2"] = lambda: self.queue_color("#00ff66")
|
|
|
|
bindings["<shift>+3"] = lambda: self.queue_color("#ffffff")
|
|
|
|
bindings["<shift>+4"] = lambda: self.queue_color("#ffff00")
|
|
|
|
bindings["<shift>+5"] = self.queue_custom_color
|
|
|
|
bindings["<ctrl>+2"] = lambda: self.queue_mode("crosshair")
|
|
|
|
bindings["<ctrl>+3"] = lambda: self.queue_mode("circle")
|
|
|
|
bindings["<ctrl>+4"] = lambda: self.queue_mode("circle_dot")
|
|
|
|
bindings["<ctrl>+5"] = lambda: self.queue_mode("9:16")
|
|
|
|
bindings["<ctrl>+6"] = lambda: self.queue_mode("4:3")
|
|
|
|
bindings["<ctrl>+7"] = lambda: self.queue_mode("1:1")
|
|
|
|
bindings["<ctrl>+8"] = lambda: self.queue_mode("thirds")
|
|
|
|
bindings["<ctrl>+1"] = self.quit
|
|
|
|
return bindings
|
|
|
|
def create_tray_icon(self):
|
|
image = Image.new(
|
|
"RGBA",
|
|
(64, 64),
|
|
(0, 0, 0, 0),
|
|
)
|
|
|
|
draw = ImageDraw.Draw(image)
|
|
|
|
draw.rectangle(
|
|
(10, 10, 54, 54),
|
|
outline=(
|
|
0,
|
|
255,
|
|
100,
|
|
255,
|
|
),
|
|
width=5,
|
|
)
|
|
|
|
return image
|
|
|
|
def create_tray(self):
|
|
self.tray = Icon(
|
|
"overlay-grid",
|
|
self.create_tray_icon(),
|
|
"Overlay Grid",
|
|
menu=Menu(
|
|
MenuItem(
|
|
"Toggle",
|
|
self.tray_toggle,
|
|
),
|
|
MenuItem(
|
|
"Settings",
|
|
self.open_menu,
|
|
),
|
|
MenuItem(
|
|
"Quit",
|
|
self.quit,
|
|
),
|
|
),
|
|
)
|
|
|
|
threading.Thread(
|
|
target=self.tray.run,
|
|
daemon=True,
|
|
).start()
|
|
|
|
def tray_toggle(
|
|
self,
|
|
icon=None,
|
|
item=None,
|
|
):
|
|
self.queue_toggle()
|
|
|
|
def open_menu(
|
|
self,
|
|
icon=None,
|
|
item=None,
|
|
):
|
|
if hasattr(
|
|
self,
|
|
"window",
|
|
):
|
|
try:
|
|
self.window.focus()
|
|
|
|
return
|
|
|
|
except Exception:
|
|
pass
|
|
|
|
self.window = ctk.CTkToplevel()
|
|
|
|
self.window.geometry("420x320")
|
|
|
|
self.window.title("Overlay Grid")
|
|
|
|
self.window.attributes(
|
|
"-topmost",
|
|
True,
|
|
)
|
|
|
|
frame = ctk.CTkFrame(self.window)
|
|
|
|
frame.pack(
|
|
fill="both",
|
|
expand=True,
|
|
padx=12,
|
|
pady=12,
|
|
)
|
|
|
|
self.hotkey_entry = ctk.CTkEntry(
|
|
frame,
|
|
width=360,
|
|
)
|
|
|
|
self.hotkey_entry.insert(
|
|
0,
|
|
", ".join(self.config["grid_hotkeys"]),
|
|
)
|
|
|
|
self.hotkey_entry.pack(pady=12)
|
|
|
|
ctk.CTkButton(
|
|
frame,
|
|
text="Save Hotkeys",
|
|
command=self.save_hotkeys,
|
|
).pack(
|
|
fill="x",
|
|
pady=6,
|
|
)
|
|
|
|
ctk.CTkButton(
|
|
frame,
|
|
text="Pick Color",
|
|
command=self.queue_custom_color,
|
|
).pack(
|
|
fill="x",
|
|
pady=6,
|
|
)
|
|
|
|
def save_hotkeys(self):
|
|
keys = [x.strip() for x in (self.hotkey_entry.get().split(",")) if x.strip()]
|
|
|
|
if not keys:
|
|
return
|
|
|
|
self.config["grid_hotkeys"] = keys
|
|
|
|
self.config_manager.save(self.config)
|
|
|
|
self.hotkeys.bindings = self.build_hotkeys()
|
|
|
|
self.hotkeys.start()
|
|
|
|
def on_click(
|
|
self,
|
|
x,
|
|
y,
|
|
button,
|
|
pressed,
|
|
):
|
|
if pressed and button == mouse.Button.middle and self.config["mouse_toggle"]:
|
|
self.queue_toggle()
|
|
|
|
def safe_queue_put(
|
|
self,
|
|
event,
|
|
):
|
|
try:
|
|
if self.queue.full():
|
|
self.queue.get_nowait()
|
|
|
|
self.queue.put_nowait(event)
|
|
|
|
except Exception:
|
|
pass
|
|
|
|
def queue_toggle(self):
|
|
self.safe_queue_put(("toggle", None))
|
|
|
|
def queue_color(
|
|
self,
|
|
color,
|
|
):
|
|
self.safe_queue_put(("color", color))
|
|
|
|
def queue_mode(
|
|
self,
|
|
mode,
|
|
):
|
|
self.safe_queue_put(("mode", mode))
|
|
|
|
def queue_custom_color(
|
|
self,
|
|
):
|
|
self.safe_queue_put(("custom", None))
|
|
|
|
def process_queue(self):
|
|
processed = 0
|
|
|
|
try:
|
|
while processed < 32:
|
|
action, value = self.queue.get_nowait()
|
|
|
|
if action == "toggle":
|
|
self.overlay.toggle()
|
|
|
|
elif action == "color":
|
|
self.overlay.set_color(value)
|
|
|
|
self.config["grid_color"] = value
|
|
|
|
self.config_manager.save(self.config)
|
|
|
|
elif action == "mode":
|
|
self.overlay.set_mode(value)
|
|
|
|
self.config["mode"] = value
|
|
|
|
self.config_manager.save(self.config)
|
|
|
|
elif action == "custom":
|
|
color = colorchooser.askcolor(
|
|
initialcolor=(self.overlay.color),
|
|
)[1]
|
|
|
|
if color:
|
|
self.overlay.set_color(color)
|
|
|
|
self.config["grid_color"] = color
|
|
|
|
self.config_manager.save(self.config)
|
|
|
|
processed += 1
|
|
|
|
except Empty:
|
|
pass
|
|
|
|
if self.running:
|
|
self.overlay.schedule(
|
|
self.process_queue,
|
|
8,
|
|
)
|
|
|
|
def shutdown(
|
|
self,
|
|
*args,
|
|
):
|
|
self.quit()
|
|
|
|
def quit(
|
|
self,
|
|
*args,
|
|
):
|
|
self.running = False
|
|
|
|
self.hotkeys.stop()
|
|
|
|
self.capture.stop()
|
|
|
|
self.validator.stop()
|
|
|
|
try:
|
|
self.mouse_listener.stop()
|
|
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
self.tray.stop()
|
|
|
|
except Exception:
|
|
pass
|
|
|
|
self.overlay.destroy()
|
|
|
|
def run(self):
|
|
self.overlay.run()
|