深入解析:基于YOLO目标检测模型的视频推理GUI工具
功能:
在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 YOLO
class 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():
return
self.btn_run.config(state="disabled")
self.stop_flag = False
threading.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 False
return True
def 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 None
self.pb["maximum"] = total_frames
self.log(f"视频总帧数:{total_frames},开始推理...")
for idx in range(total_frames):
if self.stop_flag: # 用户可随时关闭窗口中断
break
ret, 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_img
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))
# 实时显示(支持中文路径)
# cv2.namedWindow("YOLOv11 实时推理", cv2.WINDOW_NORMAL)
cv2.imshow("YOLOv11 实时推理", out_frame)
if cv2.waitKey(1) & 0xFF == ord('q'): # 按 q 也可提前退出
self.stop_flag = True
if writer:
writer.write(out_frame)
self.pb["value"] = idx + 1
if 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()

浙公网安备 33010602011771号