diff --git a/grid.py b/grid.py new file mode 100644 index 0000000..a1511ab --- /dev/null +++ b/grid.py @@ -0,0 +1,516 @@ +# 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": ["+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["+1"] = lambda: self.queue_color("#ff0000") + + bindings["+2"] = lambda: self.queue_color("#00ff66") + + bindings["+3"] = lambda: self.queue_color("#ffffff") + + bindings["+4"] = lambda: self.queue_color("#ffff00") + + bindings["+5"] = self.queue_custom_color + + bindings["+2"] = lambda: self.queue_mode("crosshair") + + bindings["+3"] = lambda: self.queue_mode("circle") + + bindings["+4"] = lambda: self.queue_mode("circle_dot") + + bindings["+5"] = lambda: self.queue_mode("9:16") + + bindings["+6"] = lambda: self.queue_mode("4:3") + + bindings["+7"] = lambda: self.queue_mode("1:1") + + bindings["+8"] = lambda: self.queue_mode("thirds") + + bindings["+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()