在深度学习应用中,YOLO 系列模型因其快速与高效而被广泛应用于目标检测任务。然而,实际场景往往需要在不同权重或模型间来回切换,例如在轻量化模型与高精度模型之间平衡性能,或针对不同任务加载不同的检测权重。如果每次切换都要手动修改代码,不仅效率低下,还可能导致显存占用冲突或用户体验不佳。
为了提升使用体验,我们在 PyQt 图形界面中集成了“模型热切换”功能。用户只需通过下拉框选择不同模型,系统即可自动完成权重加载、显存释放,并在必要时对已有图片重新推理,实现真正的一键切换。本篇文章将详细介绍实现思路与关键代码,包括统一加载函数、信号与槽的连接、模型重载的安全处理,以及可选的“自定义权重选择”功能,帮助你快速打造一个支持多模型切换的智能桌面应用。
一、新增:统一的加载函数
把它放到 MainWindow 类里(建议放在 initMain 后面、其它成员方法之前):
self.current_model_path = model_path
二、修改:initMain 里初始化/填充模型下拉框
在 initMain 末尾、或合适位置,加入这段代码(如果你已有类似变量名,可酌情替换):
说明:如果你打算在 Config.py 里显式维护清单,可加上:
三、修改:原来初始化模型的地方改用 load_model
在 initMain 中原来这两行:
替换为:
四、新增:响应模型切换的方法 + 信号连接
1.在 signalconnect 里加上对 comboBox_2 的连接
2.在 MainWindow 里新增回调方法
def on_model_combo_changed(self, text: str):
“””
当用户在 comboBox_2 切换模型时触发
1) 安全处理:如果相机/视频在运行,先停止
2) 重载模型并预热
3) 如当前界面有图片,自动用新模型重新跑一次,保持体验连续
“””
try:
if text not in self.model_choices:
QMessageBox.warning(self, “提示”, f”未找到模型项:{text}”)
return
new_model_path = self.model_choices[text]
# 如果正在读视频或摄像头,先停掉
if self.cap and self.cap.isOpened():
self.video_stop()
self.ui.VideoBtn.setText(‘打开视频’)
if self.is_camera_open:
self.is_camera_open = False
self.ui.CapBtn.setText(‘打开摄像头’)
if self.cap:
self.cap.release()
cv2.destroyAllWindows()
# 切换并预热
self.load_model(new_model_path)
# 如果当前界面已经有图片,自动用新模型再跑一遍
if getattr(self, “org_path”, None):
if os.path.isfile(self.org_path):
# 单张图
self.ui.comboBox.setDisabled(False)
self.org_img = tools.img_cvread(self.org_path)
t1 = time.time()
self.results = self.model(self.org_path, conf=self.conf, iou=self.iou)[0]
t2 = time.time()
self.ui.time_lb.setText(‘{:.3f} s’.format(t2 – t1))
location_list = self.results.boxes.xyxy.tolist()
self.location_list = [list(map(int, e)) for e in location_list]
cls_list = self.results.boxes.cls.tolist()
self.cls_list = [int(i) for i in cls_list]
self.conf_list = self.results.boxes.conf.tolist()
self.conf_list = [‘%.2f %%’ % (each * 100) for each in self.conf_list]
self.id_list = [i for i in range(len(self.location_list))]
now_img = self.results.plot()
self.draw_img = now_img
self.img_width, self.img_height = self.get_resize_size(now_img)
resize_cvimg = cv2.resize(now_img,(self.img_width, self.img_height))
pix_img = tools.cvimg_to_qpiximg(resize_cvimg)
self.ui.label_show.setPixmap(pix_img)
self.ui.label_show.setAlignment(Qt.AlignCenter)
# 目标数目 & 下拉框
target_nums = len(self.cls_list)
self.ui.label_nums.setText(str(target_nums))
choose_list = [‘全部’] + [Config.CH_names[id] + ‘_’ + str(index) for index, id in enumerate(self.cls_list)]
self.ui.comboBox.clear()
self.ui.comboBox.addItems(choose_list)
if target_nums >= 1:
self.ui.type_lb.setText(Config.CH_names[self.cls_list[0]])
self.ui.label_conf.setText(str(self.conf_list[0]))
self.ui.label_xmin.setText(str(self.location_list[0][0]))
self.ui.label_ymin.setText(str(self.location_list[0][1]))
self.ui.label_xmax.setText(str(self.location_list[0][2]))
self.ui.label_ymax.setText(str(self.location_list[0][3]))
else:
self.ui.type_lb.setText(”)
self.ui.label_conf.setText(”)
self.ui.label_xmin.setText(”)
self.ui.label_ymin.setText(”)
self.ui.label_xmax.setText(”)
self.ui.label_ymax.setText(”)
# 刷新表格 & 告警
self.ui.tableWidget.setRowCount(0)
self.ui.tableWidget.clearContents()
self.tabel_info_show(self.location_list, self.cls_list, self.conf_list, self.id_list, path=self.org_path)
self.is_alarm_on = False
self.set_alarm_action()
else:
# 如果 org_path 记录的是文件夹,就不做自动批处理了,避免卡顿
pass
QMessageBox.information(self, “提示”, f”模型已切换为:{text}\n路径:{new_model_path}”)
except Exception as e:
QMessageBox.critical(self, “错误”, f”切换模型失败:\n{e}”)
五、可选:允许用户“浏览本地”选择自定义权重
如果你想给 comboBox_2 增加一个“自定义…”项(用户点了就弹文件选择框),在步骤 2 的 model_choices 里加上:
self.model_choices[“选择本地权重…”] = “__PICK_LOCAL_FILE__”
然后在 on_model_combo_changed 的开头插入处理:
if self.model_choices[text] == “__PICK_LOCAL_FILE__”:
picked, _ = QFileDialog.getOpenFileName(self, “选择 YOLO 权重”, “./”, “Weights (*.pt *.onnx)”)
if not picked:
return
# 记入字典并替换下拉显示为文件名
base = os.path.basename(picked)
self.model_choices[base] = picked
idx = self.ui.comboBox_2.currentIndex()
self.ui.comboBox_2.setItemText(idx, base)
text = base # 继续走正常流程
六、小贴士
切模型前把视频/摄像头停掉,能避免显存占用导致的加载失败。
torch.cuda.empty_cache() 不是必须,但对频繁切换很有帮助。
如果你以后支持分割/姿态等任务,可在 load_model 里根据权重名或配置选择 task=’segment’/’pose’ 等。
如果 Ui_MainWindow 里还没有 comboBox_2,请在 Qt Designer 里放一个 QComboBox,objectName 设为 comboBox_2 即可,无需额外改 UI 代码。
评论(0)