增强这款YOLO人物检测工具的功能,增加更多实用按钮,并获得详细的步骤说明和完整可运行代码。我将为你扩展保存检测结果、清空界面、切换YOLO模型、调整检测置信度这4个核心功能,同时优化界面布局和用户体验。
一、功能扩展详细说明
新增功能列表
- 清空当前界面(重置原始图片和检测结果,方便重新操作)
- 切换YOLO模型(支持yolov5s/yolov8s轻量模型,兼顾速度和精度)
- 调整检测置信度阈值(支持0.3-0.9区间调整,过滤低置信度误检)
实现步骤拆解
步骤1:扩展类属性,新增必要的全局变量
需要添加置信度阈值、当前模型名称等属性,用于保存用户的操作配置,避免功能之间相互干扰。
步骤2:优化界面布局,新增功能按钮
在原有按钮框架基础上,新增4个功能按钮,采用网格布局合理排列,保证界面整洁美观。
步骤3:实现新增功能的核心方法
分别编写保存图片、清空界面、切换模型、调整置信度的对应方法,处理边界情况(如无检测结果时禁止保存)。
步骤4:兼容原有功能,优化细节体验
修改原有检测方法,适配新的置信度阈值和切换后的模型;优化提示信息,让用户操作更清晰。
步骤5:完善辅助方法,保证功能稳定性
确保图片格式转换、尺寸调整等辅助方法兼容新增功能,避免出现报错。
二、完整增强版代码
import tkinter as tkfrom tkinter import filedialog, messagebox, Scale, Labelimport cv2from PIL import Image, ImageTkfrom ultralytics import YOLOclassYOLOHumanDetectionApp:def__init__(self, root): self.root = root self.root.title("YOLO 高精度人物检测工具(增强版)") self.root.geometry("1200x700") # 放大窗口,适配更多功能# 原始属性保留 self.original_image = None self.detected_image = None# 新增扩展属性 self.confidence_threshold = tk.DoubleVar(value=0.5) # 置信度阈值,默认0.5 self.current_model_name = tk.StringVar(value="yolov5s") # 当前使用模型 self.model = None# 延迟加载模型,支持切换 self._load_yolo_model() # 初始化加载默认模型# 搭建增强版界面 self._create_widgets()def_load_yolo_model(self):"""加载指定名称的YOLO模型(支持yolov5s/yolov8s)"""try: model_name = self.current_model_name.get() self.model = YOLO(f"{model_name}.pt") messagebox.showinfo("模型加载成功", f"已成功加载 {model_name} 模型!")except Exception as e: messagebox.showerror("模型加载失败", f"加载 {self.current_model_name.get()} 模型出错:\n{str(e)}") self.model = Nonedef_create_widgets(self):"""搭建增强版界面,新增更多功能按钮"""# 按钮总框架 btn_main_frame = tk.Frame(self.root) btn_main_frame.pack(pady=20)# 第一排按钮:原有功能 + 新增核心功能 btn_frame1 = tk.Frame(btn_main_frame) btn_frame1.pack(pady=10) self.select_btn = tk.Button( btn_frame1, text="选择图片", command=self.select_image, width=18, height=2, font=("Arial", 10) ) self.select_btn.grid(row=0, column=0, padx=8) self.detect_btn = tk.Button( btn_frame1, text="执行人物检测", command=self.detect_humans, width=18, height=2, font=("Arial", 10) ) self.detect_btn.grid(row=0, column=1, padx=8) self.save_btn = tk.Button( btn_frame1, text="保存检测结果", command=self.save_detected_image, width=18, height=2, font=("Arial", 10) ) self.save_btn.grid(row=0, column=2, padx=8) self.clear_btn = tk.Button( btn_frame1, text="清空当前界面", command=self.clear_interface, width=18, height=2, font=("Arial", 10) ) self.clear_btn.grid(row=0, column=3, padx=8)# 第二排按钮:模型切换 + 置信度调整 btn_frame2 = tk.Frame(btn_main_frame) btn_frame2.pack(pady=10)# 模型切换按钮 self.switch_model_btn = tk.Button( btn_frame2, text="切换YOLO模型(yolov5s/yolov8s)", command=self.switch_yolo_model, width=30, height=2, font=("Arial", 10) ) self.switch_model_btn.grid(row=0, column=0, padx=8)# 置信度阈值调整 self.conf_label = Label(btn_frame2, text="检测置信度阈值:", font=("Arial", 10)) self.conf_label.grid(row=0, column=1, padx=5) self.conf_scale = Scale( btn_frame2, from_=0.3, to=0.9, resolution=0.05, variable=self.confidence_threshold, orient=tk.HORIZONTAL, length=200, font=("Arial", 9) ) self.conf_scale.grid(row=0, column=2, padx=5) self.conf_value_label = Label( btn_frame2, textvariable=self.confidence_threshold, font=("Arial", 10), width=8 ) self.conf_value_label.grid(row=0, column=3, padx=5)# 图片显示框架(保持原有布局,略微放大显示尺寸) img_frame = tk.Frame(self.root) img_frame.pack(pady=10) self.original_label = tk.Label(img_frame, text="原始图片", font=("Arial", 12)) self.original_label.grid(row=0, column=0, padx=30) self.original_img_label = tk.Label(img_frame) self.original_img_label.grid(row=1, column=0, padx=30) self.detected_label = tk.Label(img_frame, text="检测结果", font=("Arial", 12)) self.detected_label.grid(row=0, column=1, padx=30) self.detected_img_label = tk.Label(img_frame) self.detected_img_label.grid(row=1, column=1, padx=30)defselect_image(self):"""选择并显示原始图片(保留原有功能,无修改)""" file_path = filedialog.askopenfilename( title="选择图片文件", filetypes=(("图片文件", "*.jpg *.jpeg *.png *.bmp"), ("所有文件", "*.*")) )ifnot file_path:return self.original_image = cv2.imread(file_path) rgb_image = cv2.cvtColor(self.original_image, cv2.COLOR_BGR2RGB) display_image = self._resize_image(rgb_image, width=450) # 略微放大显示尺寸 pil_image = Image.fromarray(display_image) tk_image = ImageTk.PhotoImage(image=pil_image) self.original_img_label.tk_image = tk_image self.original_img_label.config(image=tk_image)defdetect_humans(self):"""YOLO人物检测(增强:适配置信度阈值和切换后的模型)"""if self.original_image isNone: messagebox.showwarning("警告", "请先选择一张图片!")returnif self.model isNone: messagebox.showerror("错误", "YOLO模型未加载成功,无法执行检测!")return# 增强:传入置信度阈值,仅保留person类别(0对应person) conf = self.confidence_threshold.get()try: results = self.model( self.original_image, classes=[0], conf=conf # 新增:使用调整后的置信度阈值 )except Exception as e: messagebox.showerror("检测失败", f"执行人物检测时出错:\n{str(e)}")return# 绘制检测结果 self.detected_image = results[0].plot() # YOLO自动绘制标注框和类别信息# 格式转换与显示(略微放大显示尺寸) rgb_detected = cv2.cvtColor(self.detected_image, cv2.COLOR_BGR2RGB) display_detected = self._resize_image(rgb_detected, width=450) pil_detected = Image.fromarray(display_detected) tk_detected = ImageTk.PhotoImage(image=pil_detected) self.detected_img_label.tk_detected = tk_detected self.detected_img_label.config(image=tk_detected)# 统计检测结果(兼容增强后的检测参数) human_count = len(results[0].boxes[results[0].boxes.cls == 0]) messagebox.showinfo("检测完成", f"共检测到 {human_count} 个人物!\n使用模型:{self.current_model_name.get()}\n置信度阈值:{conf:.2f}" )defsave_detected_image(self):"""新增:保存检测结果图片到本地"""if self.detected_image isNone: messagebox.showwarning("警告", "暂无检测结果可供保存,请先执行人物检测!")return# 弹出保存文件对话框 save_path = filedialog.asksaveasfilename( title="保存检测结果图片", defaultextension=".jpg", filetypes=(("JPG图片", "*.jpg"), ("PNG图片", "*.png"), ("所有文件", "*.*")) )ifnot save_path:return# 保存图片(使用cv2.imwrite,保持BGR格式正确)try: cv2.imwrite(save_path, self.detected_image) messagebox.showinfo("保存成功", f"检测结果已成功保存到:\n{save_path}")except Exception as e: messagebox.showerror("保存失败", f"保存图片时出错:\n{str(e)}")defclear_interface(self):"""新增:清空当前界面的图片和检测结果"""# 重置类属性 self.original_image = None self.detected_image = None# 清空图片显示标签 self.original_img_label.config(image="") self.detected_img_label.config(image="")# 清空PhotoImage引用(避免内存泄漏)if hasattr(self.original_img_label, "tk_image"):del self.original_img_label.tk_imageif hasattr(self.detected_img_label, "tk_detected"):del self.detected_img_label.tk_detected messagebox.showinfo("清空完成", "已成功清空当前界面,可重新选择图片进行检测!")defswitch_yolo_model(self):"""新增:切换YOLO模型(yolov5s <-> yolov8s)"""# 切换模型名称 current_model = self.current_model_name.get() new_model = "yolov8s"if current_model == "yolov5s"else"yolov5s" self.current_model_name.set(new_model)# 重新加载模型 self._load_yolo_model()def_resize_image(self, image, width=None, height=None):"""按比例调整图片大小(保留原有功能,无修改)""" h, w = image.shape[:2]if width isNoneand height isNone:return imageif width isNone: ratio = height / h new_w = int(w * ratio) new_h = heightelse: ratio = width / w new_w = width new_h = int(h * ratio)return cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)if __name__ == "__main__": root = tk.Tk() app = YOLOHumanDetectionApp(root) root.mainloop()
三、使用说明与注意事项
1. 环境依赖(需安装完整包)
pip install tkinter opencv-python pillow ultralytics torch
注意:tkinter在Python自带环境中通常已包含,若缺失可根据系统单独安装(如Ubuntu:sudo apt-get install python3-tk)。
2. 操作流程
- 运行代码后,自动加载默认模型
yolov5s.pt(首次运行会自动下载预训练模型,需联网) - 点击「选择图片」,导入本地jpg/png/bmp格式图片
- (可选)调整置信度阈值(0.3-0.9),数值越高检测越严格,误检越少但可能漏检
- (可选)点击「切换YOLO模型」,在
yolov5s和yolov8s之间切换(yolov8s精度更高,首次切换需下载) - 点击「执行人物检测」,查看检测结果并获取人物统计信息
- (可选)点击「保存检测结果」,将标注后的图片导出到本地
- 如需重新检测,点击「清空当前界面」,重复步骤2-6
3. 核心优化点
- 界面采用分排布局,按钮和控件排列整齐,视觉体验更好
- 检测结果提示包含模型名称和置信度阈值,方便用户调试
- 保留原有核心功能,向下兼容,无需改变用户的基础操作习惯
四、扩展功能拓展建议
若需进一步增强功能,可参考以下方向:
- 导出检测结果为Excel表格(包含人物坐标、置信度等信息)
运行上述完整代码,即可获得一款功能更完善、操作更便捷的YOLO人物检测工具,满足更多实际使用场景的需求。