Gabriel20xx/VisionFlowAutomator
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix screendhot

Browse Source
This commit is contained in:
Gabriel20xx
2025-07-07 10:46:43 +02:00
parent d1190b5061
commit 5377a62ed9
2 changed files with 848 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main.py
+155 -83
View File
@@ -13,6 +13,8 @@ import numpy as np
import pyautogui
import mss
from pynput import keyboard
import tkinter as tk
from PIL import ImageGrab, Image, ImageTk
# Setup logging
logging.basicConfig(
@@ -72,46 +74,55 @@ class Scenario:
def list_all():
return [d for d in os.listdir(CONFIG_DIR) if os.path.isdir(os.path.join(CONFIG_DIR, d))]
class CropDialog(QtWidgets.QDialog):
def __init__(self, pixmap, parent=None):
super().__init__(parent)
self.setWindowTitle("Select Area")
self.pixmap = pixmap
def take_screenshot_with_tkinter():
root = tk.Tk()
root.attributes("-alpha", 0.3)
root.attributes("-fullscreen", True)
root.wait_visibility(root)
canvas = tk.Canvas(root, cursor="cross")
canvas.pack(fill="both", expand=True)
self.label = QtWidgets.QLabel(self)
self.label.setPixmap(self.pixmap)
self.label.setCursor(QtGui.QCursor(QtCore.Qt.CursorShape.CrossCursor))
rect = None
start_x = None
start_y = None
im = ImageGrab.grab()
img = ImageTk.PhotoImage(im)
canvas.create_image(0,0,image=img,anchor="nw")
selection_rect = {"x": 0, "y": 0, "width": 0, "height": 0}
def on_button_press(event):
nonlocal start_x, start_y, rect
start_x = event.x
start_y = event.y
rect = canvas.create_rectangle(start_x, start_y, start_x, start_y, outline='red', width=2)
def on_mouse_drag(event):
nonlocal rect
cur_x, cur_y = (event.x, event.y)
canvas.coords(rect, start_x, start_y, cur_x, cur_y)
def on_button_release(event):
nonlocal selection_rect
end_x, end_y = (event.x, event.y)
self.rubber_band = QtWidgets.QRubberBand(QtWidgets.QRubberBand.Shape.Rectangle, self.label)
layout = QtWidgets.QVBoxLayout(self)
layout.setContentsMargins(0, 0, 0, 0)
layout.addWidget(self.label)
self.setLayout(layout)
selection_rect["x"] = min(start_x, end_x)
selection_rect["y"] = min(start_y, end_y)
selection_rect["width"] = abs(start_x - end_x)
selection_rect["height"] = abs(start_y - end_y)
self.setWindowState(QtCore.Qt.WindowState.WindowMaximized)
self.start_pos = None
self.selection_rect = None
root.quit()
def mousePressEvent(self, event):
self.start_pos = event.position().toPoint()
self.rubber_band.setGeometry(QtCore.QRect(self.start_pos, QtCore.QSize()))
self.rubber_band.show()
canvas.bind("<ButtonPress-1>", on_button_press)
canvas.bind("<B1-Motion>", on_mouse_drag)
canvas.bind("<ButtonRelease-1>", on_button_release)
def mouseMoveEvent(self, event):
if self.start_pos:
self.rubber_band.setGeometry(QtCore.QRect(self.start_pos, event.position().toPoint()).normalized())
def mouseReleaseEvent(self, event):
if self.start_pos:
self.selection_rect = self.rubber_band.geometry()
self.accept()
def keyPressEvent(self, event):
if event.key() == QtCore.Qt.Key.Key_Escape:
self.reject()
else:
super().keyPressEvent(event)
root.mainloop()
root.destroy()
return selection_rect
class MainWindow(QtWidgets.QMainWindow):
def automation_loop(self):
@@ -241,6 +252,7 @@ class MainWindow(QtWidgets.QMainWindow):
def add_step(self):
dlg = StepDialog(self)
dlg.setWindowFlags(dlg.windowFlags() | QtCore.Qt.WindowType.WindowStaysOnTopHint)
if dlg.exec():
step = dlg.get_step()
self.current_scenario.steps.append(step)
@@ -254,6 +266,7 @@ class MainWindow(QtWidgets.QMainWindow):
return
step = self.current_scenario.steps[idx]
dlg = StepDialog(self, step)
dlg.setWindowFlags(dlg.windowFlags() | QtCore.Qt.WindowType.WindowStaysOnTopHint)
if dlg.exec():
self.current_scenario.steps[idx] = dlg.get_step()
self.current_scenario.save()
@@ -547,49 +560,90 @@ class StepDialog(QtWidgets.QDialog):
QtWidgets.QMessageBox.warning(self, "Step Name Required", "Please enter a name for the step before adding an image.")
return
logger.debug("StepDialog.add_image_to_step: Minimizing all windows for screenshot.")
all_windows = QtWidgets.QApplication.topLevelWidgets()
window_states = {}
# Store current window states and minimize all windows
for w in all_windows:
w.setWindowState(QtCore.Qt.WindowState.WindowMinimized)
if w.isVisible():
window_states[w] = w.windowState()
w.setWindowState(QtCore.Qt.WindowState.WindowMinimized)
QtWidgets.QApplication.processEvents()
time.sleep(0.5)
try:
screen = QtWidgets.QApplication.primaryScreen()
if not screen:
logger.error("StepDialog.add_image_to_step: Could not get primary screen.")
raise Exception("Could not get primary screen.")
full_screenshot_pixmap = screen.grabWindow(0)
# Use mss for more reliable screenshot
import mss.tools
with mss.mss() as sct:
monitor = sct.monitors[1]
sct_img = sct.grab(monitor)
img = np.array(sct_img)
img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
height, width, channel = img.shape
bytes_per_line = 3 * width
qimg = QtGui.QImage(img.data, width, height, bytes_per_line, QtGui.QImage.Format.Format_RGB888)
full_screenshot_pixmap = QtGui.QPixmap.fromImage(qimg)
logger.debug("StepDialog.add_image_to_step: Screenshot taken with mss.")
except Exception as e:
logger.error(f"StepDialog.add_image_to_step: Screenshot failed: {e}")
QtWidgets.QMessageBox.critical(self, 'Error', f'Failed to take screenshot: {e}')
# Restore all windows to their original states
for w, state in window_states.items():
w.setWindowState(state)
w.showNormal()
w.activateWindow()
return
finally:
for w in all_windows:
w.setWindowState(QtCore.Qt.WindowState.WindowNoState)
# Restore all windows to their original states
for w, state in window_states.items():
w.setWindowState(state)
w.showNormal()
w.activateWindow()
crop_dialog = CropDialog(full_screenshot_pixmap, self)
if crop_dialog.exec():
rect = crop_dialog.selection_rect
if rect and not rect.isEmpty():
cropped_pixmap = full_screenshot_pixmap.copy(rect)
name_dialog = self.NameImageDialog(cropped_pixmap, self)
if name_dialog.exec():
name = name_dialog.get_name()
scenario_dir = main_window.current_scenario.get_scenario_dir()
step_images_dir = os.path.join(scenario_dir, "steps", step_name)
if not os.path.exists(step_images_dir):
os.makedirs(step_images_dir)
rect_coords = take_screenshot_with_tkinter()
# Ensure windows are properly restored after area selection
QtWidgets.QApplication.processEvents()
time.sleep(0.2)
# Restore main window and step dialog visibility
main_window.showNormal()
main_window.activateWindow()
self.showNormal()
self.activateWindow()
self.raise_()
if rect_coords and rect_coords["width"] > 0 and rect_coords["height"] > 0:
# Now, use the coordinates to crop the original full screenshot
rect = QtCore.QRect(rect_coords["x"], rect_coords["y"], rect_coords["width"], rect_coords["height"])
cropped_pixmap = full_screenshot_pixmap.copy(rect)
safe_name = "".join(c for c in name if c.isalnum() or c in (' ', '_')).rstrip()
path = os.path.join(step_images_dir, f'{safe_name}.png')
if cropped_pixmap.save(path, "PNG"):
logger.info(f"Screenshot saved to {path}")
img_obj = {'path': path, 'region': [rect.x(), rect.y(), rect.width(), rect.height()], 'name': name}
self.images.append(img_obj)
self.img_list.addItem(name)
else:
logger.error(f"Failed to save screenshot to {path}")
QtWidgets.QMessageBox.critical(self, 'Error', 'Failed to save the screenshot file.')
name_dialog = self.NameImageDialog(cropped_pixmap, self)
if name_dialog.exec():
name = name_dialog.get_name()
scenario_dir = main_window.current_scenario.get_scenario_dir()
step_images_dir = os.path.join(scenario_dir, "steps", step_name)
if not os.path.exists(step_images_dir):
os.makedirs(step_images_dir)
safe_name = "".join(c for c in name if c.isalnum() or c in (' ', '_')).rstrip()
path = os.path.join(step_images_dir, f'{safe_name}.png')
if cropped_pixmap.save(path, "PNG"):
logger.info(f"Screenshot saved to {path}")
img_obj = {'path': path, 'region': [rect.x(), rect.y(), rect.width(), rect.height()], 'name': name}
self.images.append(img_obj)
self.img_list.addItem(name)
else:
logger.error(f"Failed to save screenshot to {path}")
QtWidgets.QMessageBox.critical(self, 'Error', 'Failed to save the screenshot file.')
def update_image_preview(self, current, previous):
if current:
@@ -625,8 +679,13 @@ class StepDialog(QtWidgets.QDialog):
return
all_windows = QtWidgets.QApplication.topLevelWidgets()
window_states = {}
# Store current window states and minimize all windows
for w in all_windows:
w.setWindowState(QtCore.Qt.WindowState.WindowMinimized)
if w.isVisible():
window_states[w] = w.windowState()
w.setWindowState(QtCore.Qt.WindowState.WindowMinimized)
time.sleep(0.5)
@@ -636,27 +695,40 @@ class StepDialog(QtWidgets.QDialog):
raise Exception("Could not get primary screen.")
full_screenshot_pixmap = screen.grabWindow(0)
finally:
for w in all_windows:
w.setWindowState(QtCore.Qt.WindowState.WindowNoState)
# Restore all windows to their original states
for w, state in window_states.items():
w.setWindowState(state)
w.showNormal()
w.activateWindow()
crop_dialog = CropDialog(full_screenshot_pixmap, self)
if crop_dialog.exec():
rect = crop_dialog.selection_rect
if rect and not rect.isEmpty():
cropped_pixmap = full_screenshot_pixmap.copy(rect)
rect_coords = take_screenshot_with_tkinter()
img_obj = self.images[idx]
path = img_obj['path']
# Ensure windows are properly restored after area selection
QtWidgets.QApplication.processEvents()
time.sleep(0.2)
# Restore main window and step dialog visibility
main_window = self.parent()
main_window.showNormal()
main_window.activateWindow()
self.showNormal()
self.activateWindow()
self.raise_()
if cropped_pixmap.save(path, "PNG"):
logger.info(f"Screenshot retaken and saved to {path}")
img_obj['region'] = [rect.x(), rect.y(), rect.width(), rect.height()]
self.update_image_preview(self.img_list.currentItem(), None)
else:
logger.error(f"Failed to save retaken screenshot to {path}")
QtWidgets.QMessageBox.critical(self, 'Error', 'Failed to save the retaken screenshot file.')
if rect_coords and rect_coords["width"] > 0 and rect_coords["height"] > 0:
rect = QtCore.QRect(rect_coords["x"], rect_coords["y"], rect_coords["width"], rect_coords["height"])
cropped_pixmap = full_screenshot_pixmap.copy(rect)
img_obj = self.images[idx]
path = img_obj['path']
if cropped_pixmap.save(path, "PNG"):
logger.info(f"Screenshot retaken and saved to {path}")
img_obj['region'] = [rect.x(), rect.y(), rect.width(), rect.height()]
self.update_image_preview(self.img_list.currentItem(), None)
else:
logger.error(f"Failed to save retaken screenshot to {path}")
QtWidgets.QMessageBox.critical(self, 'Error', 'Failed to save the retaken screenshot file.')
def add_action(self):
dlg = StepActionDialog(self)