基于python与YOLO的GUI元素检测模型

ultralytics>=8.0.0torch>=2.0.0torchvision>=0.15.0opencv-python>=4.5.0Pillow>=8.0.0pyautogui>=0.9.0numpy>=1.19.0

python run_detector.py

python gui_detector.py

import tkinter as tkfrom tkinter import ttk, filedialog, messageboximport cv2import numpy as npfrom PIL import Image, ImageTkimport pyautoguifrom ultralytics import YOLOimport threadingimport osclass GUIDetectorApp:    def __init__(self, root):        self.root = root        self.root.title("GPA GUI 元素检测器")        self.root.geometry("1200x800")        # 加载模型        self.model = None        self.load_model()        # 当前图像        self.current_image = None        self.photo = None        self.detection_results = None        # 创建界面        self.create_widgets()    def load_model(self):        """加载YOLO模型"""        try:            self.model = YOLO("model.pt")            print("模型加载成功")        except Exception as e:            messagebox.showerror("错误", f"模型加载失败: {str(e)}")    def create_widgets(self):        """创建界面组件"""        # 主框架        main_frame = ttk.Frame(self.root)        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)        # 控制面板        control_frame = ttk.LabelFrame(main_frame, text="控制面板", padding=10)        control_frame.pack(fill=tk.X, pady=(0, 10))        # 截图按钮        ttk.Button(control_frame, text="截图屏幕", command=self.capture_screen).pack(side=tk.LEFT, padx=5)        ttk.Button(control_frame, text="选择图像", command=self.select_image).pack(side=tk.LEFT, padx=5)        # 参数设置        param_frame = ttk.Frame(control_frame)        param_frame.pack(side=tk.LEFT, padx=20)        ttk.Label(param_frame, text="置信度阈值:").grid(row=0, column=0, sticky=tk.W)        self.conf_var = tk.DoubleVar(value=0.05)        ttk.Scale(param_frame, from_=0.01, to=1.0, variable=self.conf_var,                  orient=tk.HORIZONTAL, length=150).grid(row=0, column=1, padx=5)        self.conf_label = ttk.Label(param_frame, text="0.05")        self.conf_label.grid(row=0, column=2)        # 缩放控制        ttk.Label(param_frame, text="图像缩放:").grid(row=1, column=0, sticky=tk.W)        self.zoom_var = tk.DoubleVar(value=1.0)        ttk.Scale(param_frame, from_=0.1, to=3.0, variable=self.zoom_var,                  orient=tk.HORIZONTAL, length=150).grid(row=1, column=1, padx=5)        self.zoom_label = ttk.Label(param_frame, text="100%")        self.zoom_label.grid(row=1, column=2)        # 操作按钮        ttk.Button(control_frame, text="保存图像", command=self.save_image).pack(side=tk.RIGHT, padx=5)        ttk.Button(control_frame, text="开始检测", command=self.start_detection,                   style="Accent.TButton").pack(side=tk.RIGHT, padx=5)        # 绑定事件        self.conf_var.trace('w', self.update_conf_label)        self.zoom_var.trace('w', self.update_zoom)        # 图像显示区域        self.image_frame = ttk.LabelFrame(main_frame, text="图像预览", padding=10)        self.image_frame.pack(fill=tk.BOTH, expand=True)        self.canvas = tk.Canvas(self.image_frame, bg='white')        self.canvas.pack(fill=tk.BOTH, expand=True)        # 结果区域        self.result_frame = ttk.LabelFrame(main_frame, text="检测结果", padding=10)        self.result_frame.pack(fill=tk.X, pady=(10, 0))        # 结果文本框        self.result_text = tk.Text(self.result_frame, height=8, width=80)        scrollbar = ttk.Scrollbar(self.result_frame, orient=tk.VERTICAL, command=self.result_text.yview)        self.result_text.configure(yscrollcommand=scrollbar.set)        self.result_text.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)    def update_conf_label(self, *args):        """更新置信度阈值显示"""        self.conf_label.config(text=f"{self.conf_var.get():.2f}")    def update_zoom(self, *args):        """更新缩放显示并重新显示图像"""        zoom = self.zoom_var.get()        self.zoom_label.config(text=f"{zoom*100:.0f}%")        if self.current_image:            self.display_image(self.current_image)    def capture_screen(self):        """截取屏幕"""        try:            # 截图            screenshot = pyautogui.screenshot()            self.current_image = screenshot            # 转换为适合显示的尺寸            display_image = self.resize_image_for_display(screenshot)            self.display_image(display_image)            self.result_text.delete(1.0, tk.END)            self.result_text.insert(tk.END, f"屏幕截图已捕获: {screenshot.size}\n")        except Exception as e:            messagebox.showerror("错误", f"截图失败: {str(e)}")    def select_image(self):        """选择图像文件"""        file_path = filedialog.askopenfilename(            title="选择图像文件",            filetypes=[("图像文件", "*.png *.jpg *.jpeg *.bmp"), ("所有文件", "*.*")]        )        if file_path:            try:                image = Image.open(file_path)                self.current_image = image                # 转换为适合显示的尺寸                display_image = self.resize_image_for_display(image)                self.display_image(display_image)                self.result_text.delete(1.0, tk.END)                self.result_text.insert(tk.END, f"图像已加载: {file_path}\n尺寸: {image.size}\n")            except Exception as e:                messagebox.showerror("错误", f"图像加载失败: {str(e)}")    def resize_image_for_display(self, image):        """根据缩放比例调整图像尺寸"""        if not self.current_image:            return image        zoom = self.zoom_var.get()        width, height = image.size        # 应用缩放比例        new_width = int(width * zoom)        new_height = int(height * zoom)        return image.resize((new_width, new_height), Image.Resampling.LANCZOS)    def display_image(self, image):        """在画布上显示图像"""        # 清除画布        self.canvas.delete("all")        # 调整图像尺寸        display_image = self.resize_image_for_display(image)        # 转换为PhotoImage        self.photo = ImageTk.PhotoImage(display_image)        # 在画布中心显示图像        canvas_width = self.canvas.winfo_width()        canvas_height = self.canvas.winfo_height()        if canvas_width > 1 and canvas_height > 1:            x = canvas_width // 2            y = canvas_height // 2        else:            x = self.photo.width() // 2            y = self.photo.height() // 2        self.canvas.create_image(x, y, image=self.photo, anchor=tk.CENTER)        # 如果存在检测结果，重新绘制边界框        if self.detection_results:            self.draw_detection_results()    def start_detection(self):        """开始检测（在新线程中运行）"""        if self.current_image is None:            messagebox.showwarning("警告", "请先截图或选择图像")            return        if self.model is None:            messagebox.showerror("错误", "模型未加载成功")            return        # 在新线程中运行检测，避免界面冻结        thread = threading.Thread(target=self.run_detection)        thread.daemon = True        thread.start()    def run_detection(self):        """运行检测算法"""        try:            # 更新界面状态            self.root.after(0, lambda: self.result_text.insert(tk.END, "开始检测...\n"))            # 获取参数            conf = self.conf_var.get()            # 自动确定图像尺寸（使用图像宽度，不超过1280）            width, height = self.current_image.size            imgsz = min(width, 1280)            # 运行YOLO检测            results = self.model.predict(                source=self.current_image,                conf=conf,                imgsz=imgsz,                iou=0.7            )            # 解析结果            boxes = results[0].boxes.xyxy.cpu().numpy()            scores = results[0].boxes.conf.cpu().numpy()            self.detection_results = {                'boxes': boxes,                'scores': scores,                'original_size': self.current_image.size            }            # 更新结果显示            self.root.after(0, self.update_results_display)            # 绘制边界框            self.root.after(0, self.draw_detection_results)        except Exception as e:            self.root.after(0, lambda: messagebox.showerror("错误", f"检测失败: {str(e)}"))    def save_image(self):        """保存当前图像（带检测结果）"""        if self.current_image is None:            messagebox.showwarning("警告", "没有可保存的图像")            return        file_path = filedialog.asksaveasfilename(            title="保存图像",            defaultextension=".png",            filetypes=[("PNG图像", "*.png"), ("JPEG图像", "*.jpg"), ("所有文件", "*.*")]        )        if file_path:            try:                # 如果存在检测结果，创建带标注的图像                if self.detection_results:                    # 使用PIL绘制边界框                    from PIL import ImageDraw                    annotated_image = self.current_image.copy()                    draw = ImageDraw.Draw(annotated_image)                    boxes = self.detection_results['boxes']                    scores = self.detection_results['scores']                    for i, (box, score) in enumerate(zip(boxes, scores)):                        x1, y1, x2, y2 = box                        # 绘制紫色边界框                        draw.rectangle([x1, y1, x2, y2], outline="purple", width=3)                        # 绘制标签                        label = f"{i+1}: {score:.2f}"                        draw.text((x1, y1 - 15), label, fill="purple")                    annotated_image.save(file_path)                else:                    # 保存原始图像                    self.current_image.save(file_path)                messagebox.showinfo("成功", f"图像已保存到: {file_path}")            except Exception as e:                messagebox.showerror("错误", f"保存失败: {str(e)}")    def update_results_display(self):        """更新结果文本框"""        if self.detection_results:            boxes = self.detection_results['boxes']            scores = self.detection_results['scores']            self.result_text.delete(1.0, tk.END)            self.result_text.insert(tk.END, f"检测完成! 发现 {len(boxes)} 个UI元素\n\n")            for i, (box, score) in enumerate(zip(boxes, scores)):                x1, y1, x2, y2 = box                self.result_text.insert(tk.END,                     f"元素 {i+1}: 位置 [{x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f}] 置信度: {score:.3f}\n")    def draw_detection_results(self):        """在图像上绘制检测结果"""        if not self.detection_results or not self.photo:            return        # 获取原始图像和显示图像的尺寸比例        orig_width, orig_height = self.detection_results['original_size']        display_width = self.photo.width()        display_height = self.photo.height()        scale_x = display_width / orig_width        scale_y = display_height / orig_height        # 绘制边界框        boxes = self.detection_results['boxes']        scores = self.detection_results['scores']        for i, (box, score) in enumerate(zip(boxes, scores)):            x1, y1, x2, y2 = box            # 缩放坐标到显示尺寸            x1_disp = x1 * scale_x            y1_disp = y1 * scale_y            x2_disp = x2 * scale_x            y2_disp = y2 * scale_y            # 计算画布中心偏移            canvas_width = self.canvas.winfo_width()            canvas_height = self.canvas.winfo_height()            if canvas_width > 1 and canvas_height > 1:                offset_x = (canvas_width - display_width) // 2                offset_y = (canvas_height - display_height) // 2            else:                offset_x = 0                offset_y = 0            # 绘制边界框（紫色）            self.canvas.create_rectangle(                x1_disp + offset_x, y1_disp + offset_y,                x2_disp + offset_x, y2_disp + offset_y,                outline="purple", width=3            )            # 绘制标签（紫色）            label_text = f"{i+1}: {score:.2f}"            self.canvas.create_text(                x1_disp + offset_x, y1_disp + offset_y - 10,                text=label_text, fill="purple", anchor=tk.SW, font=("Arial", 10, "bold")            )def main():    """主函数"""    root = tk.Tk()    app = GUIDetectorApp(root)    root.mainloop()if __name__ == "__main__":    main()

#!/usr/bin/env python3"""修复torchvision版本兼容性问题问题描述：operator torchvision::nms does not exist原因：PyTorch和torchvision版本不兼容解决方案：重新安装兼容版本"""import sysimport subprocessimport importlib.utildef check_torch_versions():    """检查当前安装的PyTorch和torchvision版本"""    try:        import torch        import torchvision        print("当前安装版本:")        print(f"PyTorch: {torch.__version__}")        print(f"torchvision: {torchvision.__version__}")        print(f"CUDA可用: {torch.cuda.is_available()}")        # 检查nms操作符是否存在        try:            from torchvision.ops import nms            print("✓ torchvision::nms 操作符可用")            return True        except ImportError as e:            print(f"✗ torchvision::nms 操作符不可用: {e}")            return False    except ImportError as e:        print(f"导入失败: {e}")        return Falsedef install_compatible_versions():    """安装兼容的PyTorch和torchvision版本"""    print("\n正在安装兼容版本...")    # 根据Python版本和系统推荐兼容版本    if sys.platform.startswith('win'):        # Windows系统推荐版本        torch_cmd = [sys.executable, "-m", "pip", "install", "torch==2.0.1", "torchvision==0.15.2", "torchaudio==2.0.2", "--index-url", "https://download.pytorch.org/whl/cu118"]    else:        # Linux/Mac系统        torch_cmd = [sys.executable, "-m", "pip", "install", "torch==2.0.1", "torchvision==0.15.2", "torchaudio==2.0.2"]    try:        print("安装PyTorch和torchvision...")        subprocess.check_call(torch_cmd)        print("✓ PyTorch安装成功")        # 重新安装ultralytics以确保兼容性        print("重新安装ultralytics...")        subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "ultralytics"])        print("✓ ultralytics安装成功")        return True    except subprocess.CalledProcessError as e:        print(f"✗ 安装失败: {e}")        return Falsedef main():    """主函数"""    print("=" * 60)    print("GPA GUI检测器 - torchvision兼容性修复工具")    print("=" * 60)    # 检查当前版本    if check_torch_versions():        print("\n✓ 当前版本兼容，无需修复")        return    print("\n检测到版本兼容性问题，开始修复...")    # 卸载现有版本    print("卸载现有版本...")    try:        subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "-y", "torch", "torchvision", "torchaudio"])    except:        pass  # 忽略卸载错误    # 安装兼容版本    if install_compatible_versions():        print("\n✓ 修复完成！请重新运行检测器")        # 验证修复结果        print("\n验证修复结果:")        check_torch_versions()    else:        print("\n✗ 修复失败，请手动安装兼容版本")        print("推荐命令:")        print("pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2")        print("pip install ultralytics")if __name__ == "__main__":    main()

#!/usr/bin/env python3"""GPA GUI 检测器 - 启动脚本使用方法:1. 确保已安装依赖: pip install -r requirements.txt2. 运行此脚本: python run_detector.py"""import sysimport osdef check_dependencies():    """检查必要的依赖是否已安装"""    required_packages = [        ('ultralytics', 'ultralytics'),        ('torch', 'torch'),        ('torchvision', 'torchvision'),        ('cv2', 'opencv-python'),        ('PIL', 'Pillow'),        ('pyautogui', 'pyautogui'),        ('tkinter', 'tkinter')  # 通常是Python自带的    ]    missing_packages = []    for import_name, package_name in required_packages:        try:            __import__(import_name)        except ImportError:            missing_packages.append(package_name)    return missing_packagesdef install_dependencies():    """安装缺失的依赖"""    missing = check_dependencies()    if missing:        print("检测到缺失的依赖包:")        for package in missing:            print(f"  - {package}")        print("\n正在安装依赖...")        try:            import subprocess            subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"])            print("依赖安装完成!")        except Exception as e:            print(f"安装失败: {e}")            print("请手动运行: pip install -r requirements.txt")            return False    return Truedef main():    """主函数"""    print("=" * 50)    print("GPA GUI 元素检测器")    print("=" * 50)    # 检查模型文件    if not os.path.exists("model.pt"):        print("错误: 未找到模型文件 'model.pt'")        print("请确保模型文件存在于当前目录")        return    # 安装依赖    if not install_dependencies():        return    # 启动GUI应用    print("\n启动GUI应用...")    try:        from gui_detector import main as gui_main        gui_main()    except Exception as e:        print(f"启动失败: {e}")        print("请检查错误信息并重试")if __name__ == "__main__":    main()