基于YOLOv5s的电动车入梯识别系统

基于YOLOv5s的电动车入梯识别系统

文章目录

• • • 1. 数据准备
    • • 标注工具
    • 2. 模型训练
    • • 安装YOLOv5
      • 准备配置文件
      • 训练模型
    • 3. UI设计
    • • 安装PyQt5
      • 创建主窗口
    • 4. 数据库集成
    • • 数据库管理 (`db_utils.py`)
    • 5. 实时检测
    • 总结

1

构建一个基于YOLOv5s的电动车入梯识别系统，我们需要完成以下几个步骤：

1. 数据准备：收集和标注数据集。
2. 模型训练：使用YOLOv5进行模型训练。
3. UI设计：使用PyQt5设计用户界面。
4. 数据库集成：使用MySQL存储检测结果。
5. 实时检测：集成模型到UI中，实现实时检测。
   

1. 数据准备

首先，你需要一个包含电动车的数据集。你可以从公开数据集中获取，或者自己采集并标注数据。

标注工具

• 使用LabelImg等工具进行数据标注。

2. 模型训练

安装YOLOv5

gitclone https://github.com/ultralytics/yolov5.gitcdyolov5 pipinstall-rrequirements.txt

准备配置文件

在yolov5目录下创建data.yaml：

train:./images/trainval:./images/valtest:./images/testnc:1# number of classesnames:['electric_bicycle']

训练模型

python train.py--img640--batch16--epochs100--datadata.yaml--cfgcfg/training/yolov5s.yaml--weightsyolov5s.pt--nameyolov5s_results

3. UI设计

使用PyQt5设计用户界面。

安装PyQt5

pipinstallPyQt5

创建主窗口

importsysfromPyQt5.QtWidgetsimportQApplication,QMainWindow,QLabel,QPushButton,QVBoxLayout,QWidget,QFileDialog,QMessageBox,QTableWidget,QTableWidgetItemfromPyQt5.QtGuiimportQPixmap,QImagefromPyQt5.QtCoreimportQTimerimportcv2importnumpyasnpfromultralyticsimportYOLOimportmysql.connectorclassElectricBikeDetection(QMainWindow):def__init__(self):super().__init__()self.initUI()self.model=YOLO('best.pt')# 加载训练好的模型self.db=DatabaseManager("localhost","root","password","detection_db")definitUI(self):self.setWindowTitle("电动车入梯识别系统")self.setGeometry(100,100,800,600)self.label=QLabel(self)self.label.setText("请选择视频文件或打开摄像头")self.label.setAlignment(Qt.AlignCenter)self.button_load=QPushButton("选择文件",self)self.button_load.clicked.connect(self.load_video)self.button_start=QPushButton("开始识别",self)self.button_start.clicked.connect(self.start_detection)self.button_camera=QPushButton("打开摄像头",self)self.button_camera.clicked.connect(self.open_camera)self.button_close_camera=QPushButton("关闭摄像头",self)self.button_close_camera.clicked.connect(self.close_camera)self.table=QTableWidget(self)self.table.setColumnCount(4)self.table.setHorizontalHeaderLabels(["时间","位置","置信度","图片"])layout=QVBoxLayout()layout.addWidget(self.label)layout.addWidget(self.button_load)layout.addWidget(self.button_start)layout.addWidget(self.button_camera)layout.addWidget(self.button_close_camera)layout.addWidget(self.table)container=QWidget()container.setLayout(layout)self.setCentralWidget(container)self.cap=Noneself.timer=Nonedefload_video(self):options=QFileDialog.Options()file_name,_=QFileDialog.getOpenFileName(self,"选择视频文件","","Video Files (*.mp4 *.avi)",options=options)iffile_name:self.video_path=file_name self.cap=cv2.VideoCapture(file_name)self.timer=QTimer(self)self.timer.timeout.connect(self.update_frame)self.timer.start(30)# 每秒更新30帧defstart_detection(self):ifnothasattr(self,'video_path'):QMessageBox.warning(self,"警告","请先选择视频文件！")returnself.button_start.setEnabled(False)self.detect_electric_bikes()defdetect_electric_bikes(self):whileself.cap.isOpened():ret,frame=self.cap.read()ifnotret:breakresults=self.model(frame)annotated_frame=results[0].plot()height,width,channel=annotated_frame.shape bytes_per_line=3*width q_img=QImage(annotated_frame.data,width,height,bytes_per_line,QImage.Format_RGB888).rgbSwapped()self.label.setPixmap(QPixmap.fromImage(q_img))forresultinresults:boxes=result.boxes.xyxy.cpu().numpy()confidences=result.boxes.conf.cpu().numpy()classes=result.boxes.cls.cpu().numpy()forbox,conf,clsinzip(boxes,confidences,classes):x_min,y_min,x_max,y_max=map(int,box)confidence=conf.item()self.db.insert_detection(x_min,y_min,x_max,y_max,confidence)self.update_table()self.app.processEvents()# 更新UIdefupdate_frame(self):ret,frame=self.cap.read()ifnotret:self.timer.stop()self.cap.release()self.button_start.setEnabled(True)returnresults=self.model(frame)annotated_frame=results[0].plot()height,width,channel=annotated_frame.shape bytes_per_line=3*width q_img=QImage(annotated_frame.data,width,height,bytes_per_line,QImage.Format_RGB888).rgbSwapped()self.label.setPixmap(QPixmap.fromImage(q_img))defopen_camera(self):self.cap=cv2.VideoCapture(0)self.timer=QTimer(self)self.timer.timeout.connect(self.update_frame)self.timer.start(30)# 每秒更新30帧defclose_camera(self):ifself.capisnotNone:self.timer.stop()self.cap.release()self.cap=Nonedefupdate_table(self):detections=self.db.get_all_detections()self.table.clearContents()self.table.setRowCount(len(detections))fori,detectioninenumerate(detections):time=detection['time']position=f"({detection['x_min']},{detection['y_min']},{detection['x_max']},{detection['y_max']})"confidence=detection['confidence']image_path=detection['image_path']self.table.setItem(i,0,QTableWidgetItem(time))self.table.setItem(i,1,QTableWidgetItem(position))self.table.setItem(i,2,QTableWidgetItem(f"{confidence:.2f}"))self.table.setItem(i,3,QTableWidgetItem(image_path))if__name__=="__main__":app=QApplication(sys.argv)window=ElectricBikeDetection()window.show()sys.exit(app.exec_())

4. 数据库集成

使用MySQL存储检测结果。

数据库管理 (db_utils.py)

importmysql.connectorclassDatabaseManager:def__init__(self,host,user,password,db_name):self.conn=mysql.connector.connect(host=host,user=user,password=password,database=db_name)self.create_tables()defcreate_tables(self):cursor=self.conn.cursor()cursor.execute(''' CREATE TABLE IF NOT EXISTS detections ( id INT AUTO_INCREMENT PRIMARY KEY, time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, x_min INT, y_min INT, x_max INT, y_max INT, confidence FLOAT, image_path VARCHAR(255) ) ''')self.conn.commit()definsert_detection(self,x_min,y_min,x_max,y_max,confidence,image_path=None):cursor=self.conn.cursor()cursor.execute(''' INSERT INTO detections (x_min, y_min, x_max, y_max, confidence, image_path) VALUES (?, ?, ?, ?, ?, ?) ''',(x_min,y_min,x_max,y_max,confidence,image_path))self.conn.commit()defget_all_detections(self):cursor=self.conn.cursor()cursor.execute('SELECT * FROM detections')returncursor.fetchall()defclose(self):self.conn.close()

5. 实时检测

在上述代码中，我们实现了以下功能：

• 选择视频文件。
• 开始识别按钮触发检测过程。
• 打开和关闭摄像头。
• 实时显示检测结果。
• 将检测结果存储到MySQL数据库中。

总结

通过以上步骤，同学构建一个基于YOLOv5s的电动车入梯识别系统。这个系统包括数据准备、模型训练、UI设计、数据库集成和实时检测等功能。