在这里插入代码片
```# -*- coding: utf-8 -*-
import os
import subprocess
from datetime import datetime
import yaml
# Form implementation generated from reading ui file 'single-train.ui'
#
# Created by: PyQt5 UI code generator 5.15.9
#
# WARNING: Any manual changes made to this file will be lost when pyuic5 is
# run again. Do not edit this file unless you know what you are doing.
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import QProcess
from base_utils.rename_class import LabelRenamer
from base_utils.split_dataset import DataSplitter
from base_utils.xmlTotxt import VOCtoYOLOConverter
class Ui_MainWindow(object):
def setupUi(self, MainWindow):
MainWindow.setObjectName("MainWindow")
MainWindow.resize(1032, 648)
MainWindow.setStyleSheet("background-image: url(bg.jpg);")
self.centralwidget = QtWidgets.QWidget(MainWindow)
self.centralwidget.setObjectName("centralwidget")
self.label = QtWidgets.QLabel(self.centralwidget)
self.label.setGeometry(QtCore.QRect(230, 0, 121, 41))
font = QtGui.QFont()
font.setFamily("微软雅黑")
font.setPointSize(22)
self.label.setFont(font)
self.label.setObjectName("label")
self.pushButton = QtWidgets.QPushButton(self.centralwidget)
self.pushButton.setGeometry(QtCore.QRect(24, 92, 81, 31))
font = QtGui.QFont()
font.setFamily("Arial")
font.setPointSize(12)
self.pushButton.setFont(font)
self.pushButton.setObjectName("pushButton")
self.pushButton.clicked.connect(self.train)
self.pushButton_2 = QtWidgets.QPushButton(self.centralwidget)
self.pushButton_2.setGeometry(QtCore.QRect(24, 202, 81, 31))
font = QtGui.QFont()
font.setFamily("Arial")
font.setPointSize(12)
self.pushButton_2.setFont(font)
self.pushButton_2.setObjectName("pushButton_2")
self.pushButton_2.clicked.connect(self.test)
self.label_3 = QtWidgets.QLabel(self.centralwidget)
self.label_3.setGeometry(QtCore.QRect(130, 180, 71, 31))
font = QtGui.QFont()
font.setFamily("微软雅黑")
font.setPointSize(12)
self.label_3.setFont(font)
self.label_3.setObjectName("label_3")
self.model_path = QtWidgets.QTextEdit(self.centralwidget)
self.model_path.setGeometry(QtCore.QRect(220, 180, 311, 31))
font = QtGui.QFont()
font.setFamily("Arial")
font.setPointSize(11)
self.model_path.setFont(font)
self.model_path.setObjectName("model_path")
self.label_4 = QtWidgets.QLabel(self.centralwidget)
self.label_4.setGeometry(QtCore.QRect(130, 230, 71, 31))
font = QtGui.QFont()
font.setFamily("微软雅黑")
font.setPointSize(12)
self.label_4.setFont(font)
self.label_4.setObjectName("label_4")
self.test_image = QtWidgets.QTextEdit(self.centralwidget)
self.test_image.setGeometry(QtCore.QRect(220, 230, 311, 31))
font = QtGui.QFont()
font.setFamily("Arial")
font.setPointSize(11)
self.test_image.setFont(font)
self.test_image.setObjectName("test_image")
self.label_5 = QtWidgets.QLabel(self.centralwidget)
self.label_5.setGeometry(QtCore.QRect(20, 320, 81, 31))
font = QtGui.QFont()
font.setFamily("微软雅黑")
font.setPointSize(12)
self.label_5.setFont(font)
self.label_5.setObjectName("label_5")
self.log_info = QtWidgets.QTextBrowser(self.centralwidget)
self.log_info.setGeometry(QtCore.QRect(10, 360, 1011, 241))
font = QtGui.QFont()
font.setFamily("Arial")
font.setPointSize(11)
self.log_info.setFont(font)
self.log_info.setObjectName("log_info")
self.label_6 = QtWidgets.QLabel(self.centralwidget)
self.label_6.setGeometry(QtCore.QRect(690, 0, 121, 41))
font = QtGui.QFont()
font.setFamily("微软雅黑")
font.setPointSize(22)
self.label_6.setFont(font)
self.label_6.setObjectName("label_6")
self.groupBox = QtWidgets.QGroupBox(self.centralwidget)
self.groupBox.setGeometry(QtCore.QRect(10, 50, 541, 271))
self.groupBox.setTitle("")
self.groupBox.setObjectName("groupBox")
self.label_7 = QtWidgets.QLabel(self.groupBox)
self.label_7.setGeometry(QtCore.QRect(120, 70, 71, 31))
font = QtGui.QFont()
font.setFamily("微软雅黑")
font.setPointSize(12)
self.label_7.setFont(font)
self.label_7.setObjectName("label_7")
self.input_root = QtWidgets.QTextEdit(self.groupBox)
self.input_root.setGeometry(QtCore.QRect(210, 70, 311, 31))
self.input_root.setObjectName("input_root")
self.camer_num = QtWidgets.QTextEdit(self.groupBox)
self.camer_num.setGeometry(QtCore.QRect(210, 20, 311, 31))
font = QtGui.QFont()
font.setFamily("Arial")
font.setPointSize(11)
self.camer_num.setFont(font)
self.camer_num.setObjectName("camer_num")
self.label_2 = QtWidgets.QLabel(self.groupBox)
self.label_2.setGeometry(QtCore.QRect(120, 20, 71, 31))
font = QtGui.QFont()
font.setFamily("微软雅黑")
font.setPointSize(12)
self.label_2.setFont(font)
self.label_2.setObjectName("label_2")
self.pushButton_3 = QtWidgets.QPushButton(self.groupBox)
self.pushButton_3.setGeometry(QtCore.QRect(10, 230, 81, 31))
font = QtGui.QFont()
font.setFamily("Arial")
font.setPointSize(12)
self.pushButton_3.setFont(font)
self.pushButton_3.setObjectName("pushButton_3")
self.pushButton_3.clicked.connect(self.pt2onnx)
self.label_8 = QtWidgets.QLabel(self.groupBox)
self.label_8.setGeometry(QtCore.QRect(120, 230, 71, 31))
font = QtGui.QFont()
font.setFamily("微软雅黑")
font.setPointSize(12)
self.label_8.setFont(font)
self.label_8.setObjectName("label_8")
self.any_model_path = QtWidgets.QTextEdit(self.groupBox)
self.any_model_path.setGeometry(QtCore.QRect(210, 230, 311, 31))
font = QtGui.QFont()
font.setFamily("Arial")
font.setPointSize(11)
self.camer_num.setFont(font)
self.any_model_path.setObjectName("any_model_path")
self.test_result = QtWidgets.QGroupBox(self.centralwidget)
self.test_result.setGeometry(QtCore.QRect(560, 50, 461, 271))
self.test_result.setTitle("")
self.test_result.setObjectName("test_result")
self.groupBox.raise_()
self.label.raise_()
self.pushButton.raise_()
self.pushButton_2.raise_()
self.label_3.raise_()
self.model_path.raise_()
self.label_4.raise_()
self.test_image.raise_()
self.label_5.raise_()
self.log_info.raise_()
self.label_6.raise_()
self.test_result.raise_()
MainWindow.setCentralWidget(self.centralwidget)
self.menubar = QtWidgets.QMenuBar(MainWindow)
self.menubar.setGeometry(QtCore.QRect(0, 0, 1032, 23))
self.menubar.setObjectName("menubar")
MainWindow.setMenuBar(self.menubar)
self.statusbar = QtWidgets.QStatusBar(MainWindow)
self.statusbar.setObjectName("statusbar")
MainWindow.setStatusBar(self.statusbar)
self.retranslateUi(MainWindow)
QtCore.QMetaObject.connectSlotsByName(MainWindow)
self.process = QProcess()
self.process.setProcessChannelMode(QProcess.MergedChannels)
def retranslateUi(self, MainWindow):
_translate = QtCore.QCoreApplication.translate
MainWindow.setWindowTitle(_translate("MainWindow", "YOLOV5"))
self.label.setText(_translate("MainWindow", "<html><head/><body><p>目标检测</p></body></html>"))
self.pushButton.setText(_translate("MainWindow", "训练"))
self.pushButton_2.setText(_translate("MainWindow", "测试"))
self.label_3.setText(_translate("MainWindow", "模型路径:"))
self.label_4.setText(_translate("MainWindow", "图像路径:"))
self.label_5.setText(_translate("MainWindow", "状态"))
self.label_6.setText(_translate("MainWindow", "测试结果"))
self.label_7.setText(_translate("MainWindow", "输入路径:"))
self.label_2.setText(_translate("MainWindow", "相机编号"))
self.pushButton_3.setText(_translate("MainWindow", "格式转换"))
self.label_8.setText(_translate("MainWindow", "模型路径:"))
def preprocess(self):
rootdir = self.input_root.toPlainText()
img_path = rootdir + r'\images'
lab_path = rootdir + r'\labels'
print('xml转txt')
converter = VOCtoYOLOConverter(lab_path)
converter.convert_all()
print('检查类别')
renamer = LabelRenamer(lab_path, 0)
renamer.rename_class()
print("划分数据")
splitter = DataSplitter(img_path, lab_path, rootdir)
splitter.split_data(train_rate=0.85, val_rate=0.15, test_rate=0, need_unlabels=True)
print('处理完成')
def train(self):
rootdir = self.input_root.toPlainText()
img_path = rootdir + r'\images'
lab_path = rootdir + r'\labels'
print('xml转txt')
converter = VOCtoYOLOConverter(lab_path)
converter.convert_all()
print('检查类别')
renamer = LabelRenamer(lab_path, 0)
renamer.rename_class()
print("划分数据")
splitter = DataSplitter(img_path, lab_path, rootdir)
splitter.split_data(train_rate=0.85, val_rate=0.15, test_rate=0, need_unlabels=True)
print('处理完成')
# 训练状态与日志
# self.process.readyReadStandardOutput.connect(self.process)
self.process.readyReadStandardOutput.connect(self.StartTrain)
self.process.readyReadStandardOutput.connect(self.Loging)
# 获取当前脚本所在的目录
script_directory = os.path.dirname(os.path.abspath(__file__))
# 创建子目录的完整路径
target_directory = os.path.join(script_directory, 'data')
# 创建目录
if not os.path.exists(target_directory):
os.makedirs(target_directory)
# 创建文件的绝对路径
camera_num = self.camer_num.toPlainText()
yaml_file_path = os.path.join(target_directory, f"base{camera_num}.yaml")
# YAML内容
now_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
rootdir = self.input_root.toPlainText()
yaml_data = {
'path': rootdir,
'train': 'train/images',
'val': 'valid/images',
'nc': 1,
'names': ['defect'],
'time': now_time
}
# 写入YAML文件
with open(yaml_file_path, 'w') as file:
yaml.dump(yaml_data, file, default_flow_style=False)
print(f"{yaml_file_path} 创建完成")
train_path = 'runs/train'
self.process.start('python train.py'
f" --data {yaml_file_path}"
f" --name {camera_num}"
f" --project {train_path}")
self.process.finished.connect(self.FinishTrain)
def test(self):
test_path = self.test_image.toPlainText()
model_path = self.model_path.toPlainText()
camera_num = self.camer_num.toPlainText()
detect_path = 'runs/detect'
# 测试状态与日志
self.process.readyReadStandardOutput.connect(self.StartTest)
self.process.readyReadStandardOutput.connect(self.Loging)
self.process.start(' python detect.py'
f" --weights {model_path}"
f" --source {test_path}"
f" --name base{camera_num}"
f" --project {detect_path}")
self.process.finished.connect(self.FinishTest)
def pt2onnx(self):
model_path = self.any_model_path.toPlainText()
self.process.readyReadStandardOutput.connect(self.StartConvert)
self.process.readyReadStandardOutput.connect(self.Loging)
self.process.start('python export.py'
f" --weights {model_path}")
self.process.finished.connect(self.FinishConvert)
def Loging(self):
data = self.process.readAllStandardOutput().data().decode('utf-8')
self.log_info.append(data)
def StartTrain(self):
self.label_5.setText('训练开始')
def FinishTrain(self):
self.label_5.setText('训练完成')
def StartTest(self):
self.label_5.setText('测试开始')
def FinishTest(self):
self.label_5.setText('测试完成')
def StartConvert(self):
self.label_5.setText('转换开始')
def FinishConvert(self):
self.label_5.setText('转换完成')
if __name__ == "__main__":
import sys
app = QtWidgets.QApplication(sys.argv)
MainWindow = QtWidgets.QMainWindow()
ui = Ui_MainWindow()
ui.setupUi(MainWindow)
MainWindow.show()
sys.exit(app.exec_())
yolov5训练测试转格式的pyqt界面
于 2023-12-21 17:10:38 首次发布