PYQT6做UI，用YOLOV8检测RTSP视频流——增加combox实现RTSP视频流的选择以及自动打开

增加combox实现RTSP视频流的选择以及自动打开

import sys
import cv2
import imutils

from PyQt6.QtCore import QThread, pyqtSignal, Qt
from PyQt6.QtGui import QImage, QPixmap
from PyQt6.QtWidgets import QApplication, QLabel, QMainWindow, QPushButton, QVBoxLayout, QWidget, QComboBox
from ultralytics import YOLO

model = YOLO("")# 你的模型地址


class VideoThread(QThread):
    model_detection_signal = pyqtSignal(QImage)

    def __init__(self, rtsp_url):
        super().__init__()
        self.rtsp_url = rtsp_url
        self.running = False

    def run(self):
        self.running = True
        cap = cv2.VideoCapture(self.rtsp_url)

        while self.running:
            # 从视频流中读取一帧
            ret, frame = cap.read()
            # 如果读取失败（例如，视频结束或连接断开），则跳出循环
            if not ret:
                break

            # 使用imutils的resize函数调整帧的大小到640像素宽
            frame = imutils.resize(frame, width=640)

            # 使用模型预测处理帧，并将结果存储在results中
            results = model.predict(source=frame)

            # 从results中获取第一个元素并绘制它
            annotated_frame = results[0].plot()
            # 获取帧的高度、宽度和通道数
            height, width, channel = annotated_frame.shape
            # 计算每行字节数，因为QImage使用每行像素数*每像素字节数的方式来表示图像数据
            bytes_per_line = 3 * width

            # 创建QImage对象，从annotated_frame中获取数据，并设置宽度、高度、每行字节数和图像格式
            # rgbSwapped()方法是因为OpenCV默认使用BGR格式，而QImage使用RGB格式
            qimage = QImage(annotated_frame.data, width, height, bytes_per_line,
                            QImage.Format.Format_RGB888).rgbSwapped()

            # 发出信号，将处理后的QImage对象发送出去，用于在GUI中显示
            self.model_detection_signal.emit(qimage)

        cap.release()

    def stop(self):
        self.running = False
        self.wait()


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()

        self.setWindowTitle("RTSP Streams with YOLOv8")
        self.setGeometry(100, 100, 800, 600)
        layout = QVBoxLayout()
        self.label = QLabel("Stream")
        self.button = QPushButton("打开")

        self.combobox = QComboBox()
        self.combobox.addItem("")#添加你的RTSP
        self.combobox.addItem("")#添加你的RTSP
        

        layout.addWidget(self.label)
        layout.addWidget(self.combobox)
        layout.addWidget(self.button)
        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

        self.thread = VideoThread(self.combobox.currentText())
        self.thread.model_detection_signal.connect(self.update_image)
        self.combobox.currentTextChanged.connect(self.on_combobox_changed)
        self.button.clicked.connect(self.start_streams)

    def on_combobox_changed(self, text):
        """
        # 通过改变combobox的内容从而管理线程
        @param text: 
        @return: 
        """
        self.thread.stop()
        self.thread = VideoThread(text)
        # 连接信号
        self.thread.model_detection_signal.connect(self.update_image)
        self.thread.start()

    def start_streams(self):
        """
        # 定义方法改变按钮的字样，同时打开/关闭线程
        @return:
        """
        if self.button.text() == "打开":
            self.button.setText("关闭")
            # 开启线程
            self.thread.start()
        else:
            self.button.setText("打开")
            # 关闭线程
            self.thread.stop()

    def update_image(self, qimage):
        """
        # 用来更新图片
        @param qimage: 获取信号的图片帧
        @return: None
        """
        # 获取标签label的尺寸
        size = self.label.size()
        # 按照标签的尺寸调整图像大小，并保持图像的宽高比
        scaled_image = qimage.scaled(size.width(), size.height(), Qt.AspectRatioMode.KeepAspectRatio)
        # 将调整大小后的QImage转换为QPixmap
        pixmap = QPixmap.fromImage(scaled_image)
        # 使用QPixmap对象设置标签label的显示内容
        self.label.setPixmap(pixmap)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())