Add grid.py

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": ["<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()