diff --git a/llama_assistant/llama_assistant_app.py b/llama_assistant/llama_assistant_app.py index 2559d5a..c2adf02 100644 --- a/llama_assistant/llama_assistant_app.py +++ b/llama_assistant/llama_assistant_app.py @@ -14,9 +14,8 @@ QVBoxLayout, QMessageBox, QSystemTrayIcon, - QRubberBand, ) -from PyQt5.QtCore import Qt, QPoint, QTimer, QSize, QRect +from PyQt5.QtCore import Qt, QTimer, QRect from PyQt5.QtGui import ( QPixmap, QPainter, @@ -24,7 +23,7 @@ QDropEvent, QBitmap, QTextCursor, - QFont, + QMouseEvent, ) from llama_assistant import config @@ -64,6 +63,24 @@ def __init__(self): self.gen_mark_down = True self.has_ocr_context = False + # Add drag-drop move support + self.setWindowFlags(Qt.FramelessWindowHint) + self.oldPos = None + + def mousePressEvent(self, event: QMouseEvent): + if event.button() == Qt.LeftButton: + self.oldPos = event.globalPos() + + def mouseMoveEvent(self, event: QMouseEvent): + if self.oldPos is not None: + delta = event.globalPos() - self.oldPos + self.move(self.x() + delta.x(), self.y() + delta.y()) + self.oldPos = event.globalPos() + + def mouseReleaseEvent(self, event: QMouseEvent): + if event.button() == Qt.LeftButton: + self.oldPos = None + def capture_screenshot(self): self.hide() self.screen_capture_widget.show() @@ -219,6 +236,10 @@ def on_ask_with_ocr_context(self): self.show() self.screen_capture_widget.hide() self.has_ocr_context = True + # Show the screenshot as reference image + if config.ocr_tmp_file.exists(): + self.dropped_image = str(config.ocr_tmp_file) + self.show_image_thumbnail(self.dropped_image) def on_submit(self): message = self.ui_manager.input_field.toPlainText() @@ -230,6 +251,7 @@ def on_submit(self): self.clear_chat() self.remove_image_thumbnail() self.dropped_image = None + self.has_ocr_context = False for file_path in self.dropped_files: self.remove_file_thumbnail(self.file_containers[file_path], file_path) @@ -239,7 +261,7 @@ def on_submit(self): self.last_response = "" self.gen_mark_down = True - if self.dropped_image: + if self.dropped_image and not self.has_ocr_context: self.process_image_with_prompt(self.dropped_image, self.dropped_files, message) self.dropped_image = None self.remove_image_thumbnail() @@ -560,6 +582,7 @@ def remove_image_thumbnail(self): self.image_label.setParent(None) self.image_label = None self.dropped_image = None + self.has_ocr_context = False self.ui_manager.input_field.setPlaceholderText("Ask me anything...") self.setFixedHeight(self.height() - 110) # Decrease height after removing image diff --git a/llama_assistant/screen_capture_widget.py b/llama_assistant/screen_capture_widget.py index ed351cf..0a95f96 100644 --- a/llama_assistant/screen_capture_widget.py +++ b/llama_assistant/screen_capture_widget.py @@ -1,7 +1,16 @@ from typing import TYPE_CHECKING -from PyQt5.QtWidgets import QApplication, QWidget, QDesktopWidget, QPushButton -from PyQt5.QtCore import Qt, QRect -from PyQt5.QtGui import QPainter, QColor, QPen +from PyQt5.QtWidgets import ( + QApplication, + QWidget, + QDesktopWidget, + QPushButton, + QLabel, + QHBoxLayout, + QVBoxLayout, + QFrame, +) +from PyQt5.QtCore import Qt, QRect, QTimer, QSize, QRectF +from PyQt5.QtGui import QPainter, QColor, QPen, QKeyEvent, QPainterPath from llama_assistant import config from llama_assistant.ocr_engine import OCREngine @@ -9,6 +18,8 @@ if TYPE_CHECKING: from llama_assistant.llama_assistant_app import LlamaAssistantApp +LEFT_BOTTOM_MARGIN = 64 + class ScreenCaptureWidget(QWidget): def __init__(self, parent: "LlamaAssistantApp"): @@ -20,7 +31,9 @@ def __init__(self, parent: "LlamaAssistantApp"): # Get screen size screen = QDesktopWidget().screenGeometry() - self.setGeometry(0, 0, screen.width(), screen.height()) + self.screen_width = screen.width() + self.screen_height = screen.height() + self.setGeometry(0, 0, self.screen_width, self.screen_height) # Set crosshairs cursor self.setCursor(Qt.CrossCursor) @@ -28,60 +41,160 @@ def __init__(self, parent: "LlamaAssistantApp"): # To store the start and end points of the mouse region self.start_point = None self.end_point = None + self.captured = False # Buttons to appear after selection self.button_widget = QWidget() + + # Create a frame for the preview with rounded corners + self.preview_frame = QFrame(self.button_widget) + self.preview_frame.setObjectName("previewFrame") + self.preview_frame.setStyleSheet( + """ + #previewFrame { + background-color: white; + border: 2px solid #808080; + border-radius: 10px; + padding: 4px; + } + """ + ) + preview_layout = QVBoxLayout(self.preview_frame) + preview_layout.setContentsMargins(4, 4, 4, 4) + + # Add close button container at the top + close_container = QWidget() + close_layout = QHBoxLayout(close_container) + close_layout.setContentsMargins(0, 0, 0, 0) + close_layout.addStretch() + + # Create close button + self.close_button = QPushButton("×", close_container) + self.close_button.setFixedSize(24, 24) + self.close_button.setCursor(Qt.PointingHandCursor) + self.close_button.setStyleSheet( + """ + QPushButton { + background-color: #ff4444; + color: white; + border: none; + border-radius: 12px; + font-size: 16px; + font-weight: bold; + } + QPushButton:hover { + background-color: #ff6666; + } + QPushButton:pressed { + background-color: #cc3333; + } + """ + ) + self.close_button.clicked.connect(self.discard_capture) + close_layout.addWidget(self.close_button) + + preview_layout.addWidget(close_container) + + self.preview_label = QLabel() + self.preview_label.setMinimumSize(QSize(400, 250)) # Bigger preview + self.preview_label.setAlignment(Qt.AlignCenter) + self.preview_label.setStyleSheet("background-color: transparent;") + self.preview_label.setScaledContents(True) # Make content scale to fit label + preview_layout.addWidget(self.preview_label) + + # Modern button styling self.ocr_button = QPushButton("OCR", self.button_widget) self.ask_button = QPushButton("Ask", self.button_widget) self.ocr_button.setCursor(Qt.PointingHandCursor) self.ask_button.setCursor(Qt.PointingHandCursor) opacity = self.parent.settings.get("transparency", 90) / 100 - base_style = f""" + base_style = """ border: none; - border-radius: 20px; + border-radius: 8px; color: white; - padding: 10px 15px; - font-size: 16px; + padding: 12px 24px; + font-size: 14px; + font-weight: 600; + min-width: 100px; """ button_style = f""" QPushButton {{ {base_style} - padding: 2.5px 5px; - border-radius: 5px; background-color: rgba{QColor(self.parent.settings["color"]).lighter(120).getRgb()[:3] + (opacity,)}; + transition: background-color 0.2s; + }} + QPushButton:hover {{ + background-color: rgba{QColor(self.parent.settings["color"]).lighter(150).getRgb()[:3] + (opacity,)}; + }} + QPushButton:pressed {{ + background-color: rgba{QColor(self.parent.settings["color"]).darker(120).getRgb()[:3] + (opacity,)}; }} """ self.ocr_button.setStyleSheet(button_style) self.ask_button.setStyleSheet(button_style) + + # Layout for buttons and preview + button_layout = QHBoxLayout() + button_layout.addStretch() + button_layout.addWidget(self.ocr_button) + button_layout.addWidget(self.ask_button) + button_layout.addStretch() + + main_layout = QVBoxLayout() + main_layout.addWidget(self.preview_frame) + main_layout.addSpacing(10) + main_layout.addLayout(button_layout) + + self.button_widget.setLayout(main_layout) self.button_widget.hide() # Connect button signals self.ocr_button.clicked.connect(self.parent.on_ocr_button_clicked) self.ask_button.clicked.connect(self.parent.on_ask_with_ocr_context) - def show(self): - # remove painting if any - self.start_point = None - self.end_point = None - self.update() + def show(self, reset=True): + if reset: + # remove painting if any + self.start_point = None + self.end_point = None + self.captured = False + self.update() - # Set window opacity to 50% - self.setWindowOpacity(0.5) - # self.setAttribute(Qt.WA_TranslucentBackground, True) + # Set window opacity to 50% + self.setWindowOpacity(0.5) + else: + self.setWindowOpacity(0.0) + self.button_widget.show() super().show() def hide(self): self.button_widget.hide() super().hide() + def discard_capture(self): + self.start_point = None + self.end_point = None + self.captured = False + self.button_widget.hide() + self.hide() + self.parent.show() + + def keyPressEvent(self, event: QKeyEvent): + if event.key() == Qt.Key_Escape: + if self.captured: + self.discard_capture() + else: + self.hide() + self.parent.show() + def mousePressEvent(self, event): - if event.button() == Qt.LeftButton: + if event.button() == Qt.LeftButton and not self.captured: self.start_point = event.pos() # Capture start position self.end_point = event.pos() # Initialize end point to start position print(f"Mouse press at {self.start_point}") def mouseReleaseEvent(self, event): - if event.button() == Qt.LeftButton: + if event.button() == Qt.LeftButton and not self.captured: self.end_point = event.pos() # Capture end position print(f"Mouse release at {self.end_point}") @@ -89,14 +202,15 @@ def mouseReleaseEvent(self, event): # Capture the region between start and end points if self.start_point and self.end_point: self.capture_region(self.start_point, self.end_point) + self.captured = True # Trigger repaint to show the red rectangle self.update() - self.show_buttons() + self.show_buttons() def mouseMoveEvent(self, event): - if self.start_point: + if self.start_point and not self.captured: # Update the end_point to the current mouse position as it moves self.end_point = event.pos() @@ -104,6 +218,16 @@ def mouseMoveEvent(self, event): self.update() def capture_region(self, start_point, end_point): + # Store current visibility state + was_visible = self.isVisible() + + # Hide the window before capturing to avoid including it in the screenshot + self.hide() + + # Small delay to ensure window is fully hidden + QTimer.singleShot(100, lambda: self._do_capture(start_point, end_point, was_visible)) + + def _do_capture(self, start_point, end_point, restore_visibility=True): # Convert local widget coordinates to global screen coordinates start_global = self.mapToGlobal(start_point) end_global = self.mapToGlobal(end_point) @@ -124,46 +248,59 @@ def capture_region(self, start_point, end_point): pixmap.save(str(config.ocr_tmp_file), "PNG") print(f"Captured region saved at '{config.ocr_tmp_file}'.") + # Update preview label with captured image + self.preview_label.setPixmap(pixmap) + + # Restore visibility if needed + if restore_visibility: + self.show(reset=False) + def paintEvent(self, event): # If the start and end points are set, draw the rectangle if self.start_point and self.end_point: # Create a painter object painter = QPainter(self) + painter.setRenderHint(QPainter.Antialiasing) - # Set the pen color to red - pen = QPen(QColor(255, 0, 0)) # Red color - pen.setWidth(3) # Set width of the border + # Set the pen color to red with a modern look + pen = QPen(QColor(255, 69, 58)) # Apple-style red + pen.setWidth(2) painter.setPen(pen) # Draw the rectangle from start_point to end_point self.region_rect = QRect(self.start_point, self.end_point) - self.region_rect = ( - self.region_rect.normalized() - ) # Normalize to ensure correct width/height + self.region_rect = self.region_rect.normalized() - painter.drawRect(self.region_rect) # Draw the rectangle + # Draw rectangle with slightly rounded corners + path = QPainterPath() + path.addRoundedRect(QRectF(self.region_rect), 4, 4) + painter.drawPath(path) - super().paintEvent(event) # Call the base class paintEvent + super().paintEvent(event) def show_buttons(self): if self.start_point and self.end_point: # Get normalized rectangle rect = QRect(self.start_point, self.end_point).normalized() - # Calculate button positions - button_y = rect.bottom() + 10 # Place buttons below the rectangle - button_width = 80 - button_height = 30 - spacing = 10 + # Calculate widget size based on preview and buttons + widget_width = 450 + widget_height = 350 - print("Showing buttons") + # Calculate position to ensure buttons stay within screen bounds + # Add LEFT_BOTTOM_MARGIN pixels offset from left to avoid macOS dock + x_pos = min( + max(LEFT_BOTTOM_MARGIN, rect.left() + (rect.width() - widget_width) // 2), + self.screen_width - widget_width, + ) - self.ocr_button.setGeometry(0, 0, button_width, button_height) - self.ask_button.setGeometry(button_width + spacing, 0, button_width, button_height) + # Check if there's enough space below the selection + y_pos = rect.bottom() + LEFT_BOTTOM_MARGIN + if y_pos + widget_height > self.screen_height: + # If not enough space below, place above the selection + y_pos = max(0, rect.top() - widget_height - LEFT_BOTTOM_MARGIN) - self.button_widget.setGeometry( - rect.left(), button_y, 2 * button_width + spacing, button_height - ) + self.button_widget.setGeometry(x_pos, y_pos, widget_width, widget_height) self.button_widget.setAttribute(Qt.WA_TranslucentBackground) self.button_widget.setWindowFlags(Qt.FramelessWindowHint) self.button_widget.show()