Add grid.py
This commit is contained in:
admin
516
grid.py
Normal file
516
grid.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user