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

基于YOLO目标检测模型的视频推理GUI工具

news 2025/9/6 6:51:36

功能：

在YOLO目标检测推理测试过程中，频繁修改视频路径和模型路径会带来不便。以下方法可以简化这一流程，避免重复操作。

通过GUI页面动态设置路径，减少硬编码带来的麻烦。这种方法提高了代码的灵活性和可维护性，这样可以避免每次测试时手动修改源代码。

使用方法：

将下方代码复制后直接执行，程序会弹出交互界面。在界面中选择权重文件、视频文件及保存路径（可选设置置信度阈值），即可开始推理。

功能说明

权重文件：加载训练好的模型权重
视频文件：指定待处理的输入视频
保存路径：自定义结果输出位置
置信度：可调节检测结果的置信度阈值（默认0.5）

import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import os, sys, threading, subprocess, cv2
from pathlib import Path
from ultralytics import YOLOclass App(tk.Tk):def __init__(self):super().__init__()self.title("YOLOv11 视频推理（实时显示）")self.geometry("540x360")self.resizable(False, False)# ---------- 权重 ----------tk.Label(self, text="权重文件 (.pt / .onnx)：").place(x=20, y=20)self.ent_w = tk.Entry(self, width=45)self.ent_w.place(x=180, y=20)tk.Button(self, text="浏览", command=lambda: self.browse(self.ent_w, "pt_onnx")).place(x=460, y=16)# ---------- 视频文件 ----------tk.Label(self, text="视频文件：").place(x=20, y=60)self.ent_i = tk.Entry(self, width=45)self.ent_i.place(x=180, y=60)tk.Button(self, text="浏览", command=lambda: self.browse(self.ent_i, "video")).place(x=460, y=56)# ---------- 输出文件夹 ----------tk.Label(self, text="结果保存到：").place(x=20, y=100)self.ent_o = tk.Entry(self, width=45)self.ent_o.place(x=180, y=100)tk.Button(self, text="浏览", command=lambda: self.browse(self.ent_o, "dir")).place(x=460, y=96)# ---------- 置信度 ----------tk.Label(self, text="置信度阈值：").place(x=20, y=140)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=120)# ---------- 复选框 ----------self.var_box = tk.BooleanVar(value=True)tk.Checkbutton(self, text="在结果上画框", variable=self.var_box).place(x=20, y=180)self.var_save = tk.BooleanVar(value=False)tk.Checkbutton(self, text="保存结果视频", variable=self.var_save).place(x=220, y=180)# ---------- 运行 / 进度 ----------self.btn_run = tk.Button(self, text="开始推理", width=15, command=self.run_thread)self.btn_run.place(x=20, y=220)self.pb = ttk.Progressbar(self, length=480, mode='determinate')self.pb.place(x=20, y=260)# ---------- 日志 ----------self.txt = tk.Text(self, height=4, width=70, state="disabled")self.txt.place(x=20, y=300)# 线程退出标志self.stop_flag = False# ---------------- 工具 ----------------def browse(self, entry, kind):if kind == "pt_onnx":f = filedialog.askopenfilename(filetypes=[("权重文件", "*.pt *.onnx")])elif kind == "video":f = filedialog.askopenfilename(filetypes=[("视频文件", "*.mp4 *.avi *.mov *.mkv")])else: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")# ---------------- 推理 ----------------def run_thread(self):if not self.validate():returnself.btn_run.config(state="disabled")self.stop_flag = Falsethreading.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 Falsereturn Truedef infer(self):try:w_path = self.ent_w.get()ext = Path(w_path).suffix.lower()if ext == '.pt':model = YOLO(w_path)elif ext == '.onnx':model = YOLO(w_path, task='detect')else:raise ValueError("权重必须是 .pt 或 .onnx")video_path = Path(self.ent_i.get())out_dir = Path(self.ent_o.get())out_dir.mkdir(parents=True, exist_ok=True)cap = cv2.VideoCapture(str(video_path))if not cap.isOpened():raise ValueError("无法打开视频文件")fps = int(cap.get(cv2.CAP_PROP_FPS))width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))fourcc = cv2.VideoWriter_fourcc(*'mp4v')save_path = out_dir / f"{video_path.stem}_result.mp4"writer = cv2.VideoWriter(str(save_path), fourcc, fps, (width, height)) if self.var_save.get() else Noneself.pb["maximum"] = total_framesself.log(f"视频总帧数：{total_frames}，开始推理...")for idx in range(total_frames):if self.stop_flag:          # 用户可随时关闭窗口中断breakret, frame = cap.read()if not ret:break# 推理results = model.predict(frame, conf=self.scale_conf.get(), save=False, verbose=False)[0]if self.var_box.get():out_frame = results.plot()h, w = out_frame.shape[:2]max_h = 720  # 想要的最大高度if h > max_h:new_w = int(w * max_h / h)out_frame = cv2.resize(out_frame, (new_w, max_h))else:out_frame = results.orig_imgh, w = out_frame.shape[:2]max_h = 720  # 想要的最大高度if h > max_h:new_w = int(w * max_h / h)out_frame = cv2.resize(out_frame, (new_w, max_h))# 实时显示（支持中文路径）# cv2.namedWindow("YOLOv11 实时推理", cv2.WINDOW_NORMAL)cv2.imshow("YOLOv11 实时推理", out_frame)if cv2.waitKey(1) & 0xFF == ord('q'):   # 按 q 也可提前退出self.stop_flag = Trueif writer:writer.write(out_frame)self.pb["value"] = idx + 1if idx % fps == 0:          # 每 1 秒写一次日志，避免刷屏self.log(f"已推理 {idx + 1}/{total_frames} 帧")cap.release()if writer:writer.release()cv2.destroyAllWindows()if self.stop_flag:self.log("用户中断推理")else:self.log("推理完成！")if self.var_save.get():subprocess.Popen(f'explorer /select,"{save_path}"')messagebox.showinfo("完成", "推理结束！")except Exception as e:messagebox.showerror("错误", str(e))finally:self.btn_run.config(state="normal")if __name__ == "__main__":if getattr(sys, 'frozen', False):os.chdir(sys._MEIPASS)App().mainloop()