onscreen-translator/create_overlay.py

305 lines
11 KiB
Python

from PySide6.QtCore import Qt, QPoint, QRect, QTimer, QBuffer
from PySide6.QtGui import (QKeySequence, QShortcut, QAction, QPainter, QFont, QScreen, QIcon)
from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
QLabel, QSystemTrayIcon, QMenu)
import sys, io, os, signal, time, platform
from PIL import Image
from dataclasses import dataclass
from typing import List, Optional
from config import ADD_OVERLAY, FONT_FILE, FONT_SIZE
from logging_config import logger
def qpixmap_to_bytes(qpixmap):
qimage = qpixmap.toImage()
buffer = QBuffer()
buffer.open(QBuffer.ReadWrite)
qimage.save(buffer, "PNG")
return qimage
@dataclass
class TextEntry:
text: str
x: int
y: int
font: QFont = QFont('Arial', FONT_SIZE)
visible: bool = True
text_color: Qt.GlobalColor = Qt.GlobalColor.red
background_color: Optional[Qt.GlobalColor] = None
padding: int = 1
class TranslationOverlay(QMainWindow):
def __init__(self):
super().__init__()
self.setWindowTitle("Translation Overlay")
self.is_passthrough = True
self.text_entries: List[TextEntry] = []
self.setup_window_attributes()
self.setup_shortcuts()
self.closeEvent = lambda event: QApplication.quit()
self.default_font = QFont('Arial', FONT_SIZE)
self.text_entries_copy: List[TextEntry] = []
self.next_text_entries: List[TextEntry] = []
#self.show_background = True
self.background_opacity = 0.5
# self.setup_tray()
def prepare_for_capture(self):
"""Preserve current state and clear overlay"""
if ADD_OVERLAY:
self.text_entries_copy = self.text_entries.copy()
self.clear_all_text()
self.update()
def restore_after_capture(self):
"""Restore overlay state after capture"""
if ADD_OVERLAY:
logger.debug(f'Text entries copy during initial phase of restore_after_capture: {self.text_entries_copy}')
self.text_entries = self.text_entries_copy.copy()
logger.debug(f"Restored text entries: {self.text_entries}")
self.update()
def add_next_text_at_position_no_update(self, x: int, y: int, text: str,
font: Optional[QFont] = None, text_color: Qt.GlobalColor = Qt.GlobalColor.red):
"""Add new text without triggering update"""
entry = TextEntry(
text=text,
x=x,
y=y,
font=font or self.default_font,
text_color=text_color
)
self.next_text_entries.append(entry)
def update_translation(self, ocr_output, translation):
# Update your overlay with new translations here
# You'll need to implement the logic to display the translations
self.clear_all_text()
self.text_entries = self.next_text_entries.copy()
self.next_text_entries.clear()
self.update()
def capture_behind(self, x=None, y=None, width=None, height=None):
"""
Capture the screen area behind the overlay.
If no coordinates provided, captures the area under the window.
"""
# Temporarily hide the window
self.hide()
# Get screen
screen = QScreen.grabWindow(
self.screen(),
0,
x if x is not None else self.x(),
y if y is not None else self.y(),
width if width is not None else self.width(),
height if height is not None else self.height()
)
# Show the window again
self.show()
screen_bytes = qpixmap_to_bytes(screen)
return screen_bytes
def clear_all_text(self):
"""Clear all text entries"""
self.text_entries.clear()
self.update()
def setup_window_attributes(self):
# Set window flags for overlay behavior
self.setWindowFlags(
Qt.WindowType.FramelessWindowHint |
Qt.WindowType.WindowStaysOnTopHint |
Qt.WindowType.Tool
)
# Set attributes for transparency
self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)
# Make the window cover the entire screen
self.setGeometry(QApplication.primaryScreen().geometry())
# Special handling for Wayland
if platform.system() == "Linux":
if "WAYLAND_DISPLAY" in os.environ:
self.setAttribute(Qt.WidgetAttribute.WA_X11NetWmWindowTypeCombo)
self.setAttribute(Qt.WidgetAttribute.WA_DontCreateNativeAncestors)
def setup_shortcuts(self):
# Toggle visibility (Alt+Shift+T)
self.toggle_visibility_shortcut = QShortcut(QKeySequence("Alt+Shift+T"), self)
self.toggle_visibility_shortcut.activated.connect(self.toggle_visibility)
# Toggle passthrough mode (Alt+Shift+P)
self.toggle_passthrough_shortcut = QShortcut(QKeySequence("Alt+Shift+P"), self)
self.toggle_passthrough_shortcut.activated.connect(self.toggle_passthrough)
# Quick hide (Escape)
self.hide_shortcut = QShortcut(QKeySequence("Esc"), self)
self.hide_shortcut.activated.connect(self.hide)
# Clear all text (Alt+Shift+C)
self.clear_shortcut = QShortcut(QKeySequence("Alt+Shift+C"), self)
self.clear_shortcut.activated.connect(self.clear_all_text)
# Toggle background
self.toggle_background_shortcut = QShortcut(QKeySequence("Alt+Shift+B"), self)
self.toggle_background_shortcut.activated.connect(self.toggle_background)
def toggle_visibility(self):
if self.isVisible():
self.hide()
else:
self.show()
def toggle_passthrough(self):
self.is_passthrough = not self.is_passthrough
if self.is_passthrough:
self.setAttribute(Qt.WidgetAttribute.WA_TransparentForMouseEvents)
if platform.system() == "Linux" and "WAYLAND_DISPLAY" not in os.environ:
self.setWindowFlags(self.windowFlags() | Qt.WindowType.X11BypassWindowManagerHint)
else:
self.setAttribute(Qt.WidgetAttribute.WA_TransparentForMouseEvents, False)
if platform.system() == "Linux" and "WAYLAND_DISPLAY" not in os.environ:
self.setWindowFlags(self.windowFlags() & ~Qt.WindowType.X11BypassWindowManagerHint)
self.hide()
self.show()
def toggle_background(self):
"""Toggle background visibility"""
self.show_background = not self.show_background
self.update()
def set_background_opacity(self, opacity: float):
"""Set background opacity (0.0 to 1.0)"""
self.background_opacity = max(0.0, min(1.0, opacity))
self.update()
def add_text_at_position(self, x: int, y: int, text: str):
"""Add new text at specific coordinates"""
entry = TextEntry(text, x, y)
self.text_entries.append(entry)
self.update()
def update_text_at_position(self, x: int, y: int, text: str):
"""Update text at specific coordinates, or add if none exists"""
# Look for existing text entry near these coordinates (within 5 pixels)
for entry in self.text_entries:
if abs(entry.x - x) <= 1 and abs(entry.y - y) <= 1:
entry.text = text
self.update()
return
# If no existing entry found, add new one
self.add_text_at_position(x, y, text)
def setup_tray(self):
self.tray_icon = QSystemTrayIcon(self)
self.tray_icon.setIcon(QIcon.fromTheme("applications-system"))
tray_menu = QMenu()
toggle_action = tray_menu.addAction("Show/Hide Overlay")
toggle_action.triggered.connect(self.toggle_visibility)
toggle_passthrough = tray_menu.addAction("Toggle Passthrough")
toggle_passthrough.triggered.connect(self.toggle_passthrough)
# Add background toggle to tray menu
toggle_background = tray_menu.addAction("Toggle Background")
toggle_background.triggered.connect(self.toggle_background)
clear_action = tray_menu.addAction("Clear All Text")
clear_action.triggered.connect(self.clear_all_text)
tray_menu.addSeparator()
quit_action = tray_menu.addAction("Quit")
quit_action.triggered.connect(self.clean_exit)
self.tray_icon.setToolTip("Translation Overlay")
self.tray_icon.setContextMenu(tray_menu)
self.tray_icon.show()
self.tray_icon.activated.connect(self.tray_activated)
def remove_text_at_position(self, x: int, y: int):
"""Remove text entry near specified coordinates"""
self.text_entries = [
entry for entry in self.text_entries
if abs(entry.x - x) > 1 or abs(entry.y - y) > 1
]
self.update()
def paintEvent(self, event):
painter = QPainter(self)
painter.setRenderHint(QPainter.RenderHint.Antialiasing)
# Draw each text entry
for entry in self.text_entries:
if not entry.visible:
continue
# Set the font for this specific entry
painter.setFont(entry.font)
text_metrics = painter.fontMetrics()
# Get the bounding rectangles for text
text_bounds = text_metrics.boundingRect(
entry.text
)
total_width = text_bounds.width()
total_height = text_bounds.height()
# Create rectangles for text placement
text_rect = QRect(entry.x, entry.y, total_width, total_height)
# Calculate background rectangle that encompasses both texts
if entry.background_color is not None:
bg_rect = QRect(entry.x - entry.padding,
entry.y - entry.padding,
total_width + (2 * entry.padding),
total_height + (2 * entry.padding))
painter.setPen(Qt.PenStyle.NoPen)
painter.setBrush(entry.background_color)
painter.drawRect(bg_rect)
# Draw the texts
painter.setPen(entry.text_color)
painter.drawText(text_rect, Qt.AlignmentFlag.AlignLeft, entry.text)
def handle_exit(signum, frame):
QApplication.quit()
def start_overlay():
app = QApplication(sys.argv)
# Enable Wayland support if available
if platform.system() == "Linux" and "WAYLAND_DISPLAY" in os.environ:
app.setProperty("platform", "wayland")
overlay = TranslationOverlay()
overlay.show()
signal.signal(signal.SIGINT, handle_exit) # Handle Ctrl+C (KeyboardInterrupt)
signal.signal(signal.SIGTERM, handle_exit)
return (app, overlay)
# sys.exit(app.exec())
if ADD_OVERLAY:
app, overlay = start_overlay()
if __name__ == "__main__":
ADD_OVERLAY = True
if not ADD_OVERLAY:
app, overlay = start_overlay()
overlay.add_text_at_position(600, 100, "Hello World I AM A BIG FAAT FOROGGGGGGGGGG")
capture = overlay.capture_behind()
capture.save("capture.png")
sys.exit(app.exec())