当前位置：首页 > news >正文

YOLO11目标检测运行推理简约GUI界面

news 2025/9/12 6:29:44

YOLO11推理简约GUI界面

使用方法：

支持pt和onnx格式模型,并且自动检测设备，选择推理设备

选择推理图片所在的文件夹
选择推理后的结果保存地址

选择所需要的置信度阈值

点击开始推理，程序自动运行并在下方实时显示推理进度

非常方便不用每次都改代码来推理了

界面如下所示：

代码如下：

# -*- coding: utf-8 -*-
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import os, sys, threading, subprocess
from pathlib import Path
import cv2
import torch
from ultralytics import YOLOclass App(tk.Tk):def __init__(self):super().__init__()self.title("YOLOv11 批量推理 (支持 .pt / .onnx)")self.geometry("540x480")  # 增加了高度以显示设备信息self.resizable(False, False)# 设备信息显示self.device_info = self.get_device_info()tk.Label(self, text=f"检测到的设备: {self.device_info}", fg="blue").place(x=20, y=5)# ---- 权重 ----tk.Label(self, text="权重文件 (.pt / .onnx)：").place(x=20, y=40)self.ent_w = tk.Entry(self, width=45)self.ent_w.place(x=180, y=40)tk.Button(self, text="浏览", command=lambda: self.browse_file(self.ent_w, [("模型文件", "*.pt *.onnx")])).place(x=460, y=36)# ---- 图片文件夹 ----tk.Label(self, text="图片文件夹：").place(x=20, y=80)self.ent_i = tk.Entry(self, width=45)self.ent_i.place(x=180, y=80)tk.Button(self, text="浏览", command=lambda: self.browse_directory(self.ent_i)).place(x=460, y=76)# ---- 输出文件夹 ----tk.Label(self, text="结果保存到：").place(x=20, y=120)self.ent_o = tk.Entry(self, width=45)self.ent_o.place(x=180, y=120)tk.Button(self, text="浏览", command=lambda: self.browse_directory(self.ent_o)).place(x=460, y=116)# ---- 置信度 ----tk.Label(self, text="置信度阈值：").place(x=20, y=160)self.scale_conf = tk.Scale(self, from_=0.01, to=1.0, resolution=0.01,orient=tk.HORIZONTAL, length=300)self.scale_conf.set(0.35)self.scale_conf.place(x=180, y=140)# ---- 设备选择 ----tk.Label(self, text="推理设备：").place(x=20, y=200)self.device_var = tk.StringVar(value="auto")devices = self.get_available_devices()self.device_combo = ttk.Combobox(self, textvariable=self.device_var, values=devices, width=15, state="readonly")self.device_combo.place(x=180, y=200)# ---- 复选框 ----self.var_empty = tk.BooleanVar(value=True)self.var_box = tk.BooleanVar(value=True)self.var_recursive = tk.BooleanVar(value=False)tk.Checkbutton(self, text="保存无目标的图片", variable=self.var_empty).place(x=20, y=240)tk.Checkbutton(self, text="在结果图片上画框", variable=self.var_box).place(x=220, y=240)tk.Checkbutton(self, text="递归子文件夹", variable=self.var_recursive).place(x=20, y=270)# ---- 运行按钮 / 进度条 ----self.btn_run = tk.Button(self, text="开始推理", width=15, command=self.run_thread)self.btn_run.place(x=20, y=310)self.pb = ttk.Progressbar(self, length=480, mode='determinate')self.pb.place(x=20, y=350)# ---- 日志 ----self.txt = tk.Text(self, height=6, width=70, state="disabled")self.txt.place(x=20, y=380)def get_device_info(self):"""获取设备信息"""if torch.cuda.is_available():gpu_count = torch.cuda.device_count()gpu_name = torch.cuda.get_device_name(0)return f"GPU: {gpu_name} ({gpu_count}个)"else:return "CPU only"def get_available_devices(self):"""获取可用设备列表"""devices = ["auto"]if torch.cuda.is_available():for i in range(torch.cuda.device_count()):devices.append(f"cuda:{i}")devices.append("cpu")return devicesdef browse_file(self, entry, filetypes):"""浏览文件"""f = filedialog.askopenfilename(filetypes=filetypes)if f:entry.delete(0, tk.END)entry.insert(0, f)def browse_directory(self, entry):"""浏览目录"""f = filedialog.askdirectory()if f:entry.delete(0, tk.END)entry.insert(0, f)def log(self, msg):"""日志输出"""self.txt.configure(state="normal")self.txt.insert(tk.END, msg + "\n")self.txt.see(tk.END)self.txt.configure(state="disabled")self.update()# ---------- 推理 ----------def run_thread(self):"""启动推理线程"""if not self.validate():returnself.btn_run.config(state="disabled")self.pb["value"] = 0threading.Thread(target=self.infer, daemon=True).start()def validate(self):"""验证输入"""for e in (self.ent_w, self.ent_i, self.ent_o):if not e.get():messagebox.showerror("提示", "请完整填写路径！")return Falsew_path = Path(self.ent_w.get())if not w_path.exists():messagebox.showerror("错误", "权重文件不存在！")return Falseif w_path.suffix.lower() not in ['.pt', '.onnx']:messagebox.showerror("错误", "只支持 .pt 或 .onnx 格式的权重文件！")return Falsei_path = Path(self.ent_i.get())if not i_path.exists():messagebox.showerror("错误", "图片文件夹不存在！")return Falsereturn Truedef infer(self):"""执行推理"""try:# 获取设备设置device_choice = self.device_var.get()if device_choice == "auto":device = "0" if torch.cuda.is_available() else "cpu"else:device = device_choicew_path = self.ent_w.get()ext = Path(w_path).suffix.lower()self.log(f"正在加载模型，使用设备: {device}...")# 加载模型if ext == '.pt':# 对于PyTorch模型，直接在加载时指定设备[4,5](@ref)model = YOLO(w_path, device=device)elif ext == '.onnx':# ONNX模型当前版本默认使用CPU[1](@ref)model = YOLO(w_path, task='detect')device = 'cpu'self.log("注意: ONNX模型当前使用CPU推理，如需GPU加速请使用TensorRT转换")else:raise ValueError("不支持的权重格式")in_dir = Path(self.ent_i.get())out_dir = Path(self.ent_o.get())out_dir.mkdir(parents=True, exist_ok=True)# 收集图片文件pattern = '**/*' if self.var_recursive.get() else '*'img_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}imgs = [p for p in in_dir.glob(pattern)if p.suffix.lower() in img_extensions and p.is_file()]total = len(imgs)if total == 0:messagebox.showwarning("警告", "未找到任何图片！")returnself.pb["maximum"] = totalself.log(f"找到 {total} 张图片，开始推理...")# 批量推理[8](@ref)for idx, img_p in enumerate(imgs, 1):rel_path = img_p.relative_to(in_dir) if in_dir in img_p.parents else Path()save_dir = out_dir / rel_path.parentsave_dir.mkdir(parents=True, exist_ok=True)save_img = save_dir / f"{img_p.stem}_result.jpg"# 执行推理[4,8](@ref)results = model.predict(source=str(img_p),conf=self.scale_conf.get(),save=False,verbose=False,device=device  # 确保指定设备)result = results[0]if len(result.boxes) == 0 and not self.var_empty.get():continue# 处理结果图像img_out = result.plot() if self.var_box.get() else result.orig_imgcv2.imwrite(str(save_img), img_out)# 更新进度self.pb["value"] = idxself.log(f"[{idx:03d}/{total:03d}] {img_p.name}")# 完成后打开结果文件夹subprocess.Popen(f'explorer "{out_dir}"')messagebox.showinfo("完成", f"推理完成！处理了 {total} 张图片。")except Exception as e:error_msg = f"推理错误: {str(e)}"self.log(error_msg)messagebox.showerror("错误", error_msg)finally:self.btn_run.config(state="normal")self.pb["value"] = 0if __name__ == "__main__":# 在打包环境下调整路径if getattr(sys, 'frozen', False):os.chdir(os.path.dirname(sys.executable))app = App()app.mainloop()