基于 YOLOv5 和 PyQt5 的煤矿运输带安全检测平台开发

一、引言

在煤矿生产过程中，传输带作为重要的物料运输设备，其运行状态直接关系到生产效率和安全。然而，传输带在运行过程中常常会出现跑偏、打滑、火灾等安全事故，这些事故不仅会导致设备损坏和生产中断，还可能引发人员伤亡和重大经济损失。因此，实时监测传输带的运行状态并及时预警，对于保障煤矿生产安全具有重要意义。目标检测是计算机视觉领域的重要任务之一，广泛应用于安防监控、自动驾驶、智能机器人等多个领域。YOLOv5 是一种高效、轻量级的目标检测算法，而 PyQt5 是一个用于开发跨平台桌面应用程序的 Python 库。本文将介绍如何结合 YOLOv5 和 PyQt5 开发一个简单的目标检测平台，实现图片、视频和摄像头的目标检测功能，并将检测结果保存到数据库中，为煤矿传输带的安全监测提供技术支持。

二、开发环境准备及模型训练

（一）开发环境

1. Python 环境： Python 3.6 或更高版本。

2. PyQt5：安装 PyQt5：

   pip install PyQt5

3. YOLOv5：按照 YOLOv5 官方仓库 的说明进行安装和配置。

4. PyTorch：根据系统环境选择合适的 PyTorch 版本进行安装。

5. OpenCV：用于图像和视频处理，通过以下命令安装

   pip install opencv-python

6. 其他依赖库：安装代码中用到的其他依赖库，如 numpy、qdarkstyle 等。

7. 数据库：安装 MySQL 数据库，并创建一个名为 demo 的数据库，用于存储检测结果。

（二）模型训练 

1.数据准备

• 收集包含传输带故障（如跑偏、打滑、火灾等）的图像和视频数据。

• 对数据进行标注，使用工具 LabelImg 标注目标的位置和类别，类别主要分为断带，停带，着火，烟雾四类。

• 将数据按照分为训练集和验证集和测试集，比例为8：1：1。

2. 模型训练

• 在any.yaml文件里替换数据集路径和类别信息。

• 更改yolov5s.yaml里的训练数为4。

• 可在train.py文件下找到 以下代码，设置迭代次数，这里迭代80次 

•  parser.add_argument("--epochs", type=int, default=80 , help="total training epochs")

• 运行train.py文件，模型权重文件将保存在 runs/train/weights 目录下。

三、代码实现

1. 主界面设计

主界面使用 PyQt5 设计，包含图片检测、视频检测、摄像头检测和打开输出文件夹四个按钮，以及一个用于显示检测结果的区域。

class Ui_MainWindow(QtWidgets.QMainWindow):
    def __init__(self, parent=None):
        super(Ui_MainWindow, self).__init__(parent)
        self.timer_video = QtCore.QTimer()
        self.setupUi(self)
        self.init_logo()
        self.init_slots()
        self.cap = cv2.VideoCapture()
        self.out = None
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.half = self.device.type != 'cpu'  # half precision only supported on CUDA
        # 其他初始化代码...

2. 图片检测功能

点击“图片检测”按钮，通过文件对话框选择图片，调用 YOLOv5 模型进行检测，并将结果显示在主界面的显示区域。

def button_image_open(self):
    print('打开图片')
    name_list = []

    img_name, _ = QtWidgets.QFileDialog.getOpenFileName(
        self, "打开图片", "", "*.jpg;;*.png;;All Files(*)")
    if not img_name:
        return
    img = cv2.imread(img_name)
    print(img_name)
    showimg = img
    with torch.no_grad():
        img = letterbox(img, new_shape=self.imgsz)[0]
        # Convert
        # BGR to RGB, to 3x416x416
        img = img[:, :, ::-1].transpose(2, 0, 1)
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(self.device)
        img = img.half() if self.half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        # Inference
        pred = self.model(img)[0]
        # Apply NMS
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)
        # Process detections
        for i, det in enumerate(pred):
            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_boxes(
                    img.shape[2:], det[:, :4], showimg.shape).round()

                for *xyxy, conf, cls in reversed(det):
                    label = '%s %.2f' % (self.names[int(cls)], conf)
                    # print(label.split()[0])  # 打印各目标名称
                    name_list.append(self.names[int(cls)])
                    plot_one_box(xyxy, showimg, label=label,
                                 color=self.colors[int(cls)], line_thickness=2)

    cv2.imwrite('result/prediction.jpg', showimg)
    self.result = cv2.cvtColor(showimg, cv2.COLOR_BGR2BGRA)
    self.result = cv2.resize(self.result, (640, 480), interpolation=cv2.INTER_AREA)
    self.QtImg = QtGui.QImage(self.result.data, self.result.shape[1], self.result.shape[0],
                              QtGui.QImage.Format_RGB32)
    self.label.setPixmap(QtGui.QPixmap.fromImage(self.QtImg))

3. 视频检测功能

点击“视频检测”按钮，通过文件对话框选择视频文件，调用 YOLOv5 模型进行实时检测。

def button_video_open(self):
    video_name, _ = QtWidgets.QFileDialog.getOpenFileName(
        self, "打开视频", "", "*.mp4;;*.avi;;All Files(*)")

    if not video_name:
        return

    flag = self.cap.open(video_name)
    if flag == False:
        QtWidgets.QMessageBox.warning(
            self, u"Warning", u"打开视频失败", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok)
    else:
        self.out = cv2.VideoWriter('result/vedio_prediction.avi', cv2.VideoWriter_fourcc(
            *'MJPG'), 20, (int(self.cap.get(3)), int(self.cap.get(4))))
        self.timer_video.start(30)
        self.pushButton_video.setDisabled(True)
        self.pushButton_img.setDisabled(True)
        self.pushButton_camera.setDisabled(True)

4. 摄像头检测功能

点击“摄像头检测”按钮，打开摄像头进行实时检测。

def button_camera_open(self):
    if not self.timer_video.isActive():
        # 默认使用第一个本地camera
        flag = self.cap.open(0)
        if flag == False:
            QtWidgets.QMessageBox.warning(
                self, u"Warning", u"打开摄像头失败", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok)
        else:
            self.out = cv2.VideoWriter('result/camera_prediction.avi', cv2.VideoWriter_fourcc(
                *'MJPG'), 20, (int(self.cap.get(3)), int(self.cap.get(4))))
            self.timer_video.start(30)
            self.pushButton_video.setDisabled(True)
            self.pushButton_img.setDisabled(True)
            self.pushButton_camera.setText(u"关闭摄像头")
    else:
        self.timer_video.stop()
        self.cap.release()
        self.out.release()
        self.label.clear()
        self.init_logo()
        self.pushButton_video.setDisabled(False)
        self.pushButton_img.setDisabled(False)
        self.pushButton_camera.setText(u"摄像头检测")

5. 数据库保存功能

将检测结果保存到 MySQL 数据库中。

def save_detection_to_db(self, detection_info):
    # detection_info 
    # detection_info = {
    #     'class_name': class_name,
    #     'confidence': confidence,
    #     'xyxy': xyxy
    # }
    x1, y1, x2, y2 = [int(item) for item in detection_info['xyxy']]
    SQLHelper.fetch_one(
        "INSERT INTO detections (class_name, confidence, x1, y1, x2, y2, detection_time) VALUES (%s, %s, %s, %s, %s, %s, %s)",
        (detection_info['class_name'], detection_info['confidence'], x1, y1, x2, y2, detection_info['time'])
    )

6. 数据库配置和操作类

class Config(object):
    PYMYSQL_POOL = PooledDB(
        creator=pymysql,  # 使用链接数据库的模块
        maxconnections=6,  # 连接池允许的最大连接数，0和None表示不限制连接数
        mincached=2,  
        maxcached=5,  
        maxshared=3,
        blocking=True,  
        maxusage=None,  # 一个链接最多被重复使用的次数，None表示无限制
        setsession=[],  # 开始会话前执行的命令列表。
        ping=0,
        host='localhost',
        port=3306,
        user='root',
        password='060811',  # 你自己的密码
        database='demo',  # 提前建好库
        charset='utf8'
    )


class SQLHelper(object):
    @staticmethod
    def open(cursor):
        POOL = Config.PYMYSQL_POOL
        conn = POOL.connection()
        cursor = conn.cursor(cursor=cursor)
        return conn, cursor

    @staticmethod
    def close(conn, cursor):
        conn.commit()
        cursor.close()
        conn.close()

    @classmethod
    def fetch_one(cls, sql, args, cursor=pymysql.cursors.DictCursor):
        conn, cursor = cls.open(cursor)
        cursor.execute(sql, args)
        obj = cursor.fetchone()
        cls.close(conn, cursor)
        return obj

四、总结

在本项目中，为了提升煤矿传输带的安全监测能力，我们引入了基于 YOLOv5 的目标检测技术。通过数据准备、模型训练与优化，成功开发出高效检测传输带故障（跑偏、打滑、火灾等）的系统。该系统集成到 PyQt5 平台，实现了对图片、视频或摄像头输入的实时检测功能，并将检测结果实时显示在界面上，为操作人员提供了直观的监测信息，极大地提高了传输带故障检测的实时性和准确性，为煤矿生产安全提供了有力的技术支持。