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}...")# 關鍵修改：初始化時不傳遞device參數[1,3,4](@ref)model = YOLO(w_path)# 對于PyTorch模型，使用.to()方法遷移到指定設備[6,7](@ref)if ext == '.pt' and device != 'cpu':model.to(device)self.log(f"PyTorch模型已遷移到設備: {device}")elif ext == '.onnx':self.log("注意: ONNX模型當前使用CPU推理，如需GPU加速請使用TensorRT轉換")device = 'cpu'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} 張圖片，開始推理...")# 批量推理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"# 執行推理，在predict方法中指定設備[1,3,4](@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()