基于PyQt的YOLOv5+DeepSORT可视化界面,可实现目标跟踪、模型更换、结果保存和轨迹隐藏等功能。

基于PyQt的YOLOv5+DeepSORT可视化界面,可实现目标跟踪、模型更换、结果保存和轨迹隐藏等功能。

代码示例,仅供参考,包括界面设计和功能实现。

文章目录

YOLOv5+Deepsort可视化界面

基于pyqt

预实现功能特点:

勾选跟踪目标/更换自己的模型/保存结果/隐藏轨迹

基于PyQt的YOLOv5+DeepSORT可视化界面,可实现目标跟踪、模型更换、结果保存和轨迹隐藏等功能。

代码示例,仅供参考,包括界面设计和功能实现。

如何训练呢?车辆行人检测数据集

训练YOLOv5模型以识别车辆和行人两个类别的数据集涉及几个步骤,准备数据、配置YOLOv5模型、训练模型以及使用DeepSORT进行跟踪。以下是一个详细的指南:

仅供参考吧

1. 准备数据

确保同学数据集已经按照YOLO格式准备好,每个图像文件都有一个对应的.txt标签文件,内容如下(每行代表一个对象):

复制代码
<class_index> <x_center> <y_center> <width> <height>
  • <class_index>: 类别的索引(例如,0代表行人,1代表车辆)。
  • <x_center>, <y_center>: 边界框中心点相对于图像宽度和高度的比例。
  • <width>, <height>: 边界框的宽度和高度相对于图像宽度和高度的比例。

假设你的数据结构如下:

复制代码
dataset/
├── images/
│   ├── train/
│   └── val/
└── labels/
    ├── train/
    └── val/

2. 配置YOLOv5

创建一个名为data.yaml的数据集配置文件,内容如下:

yaml 复制代码
train: ./dataset/images/train  # 训练集图片路径
val: ./dataset/images/val      # 验证集图片路径

nc: 2  # 类别数量,这里是2,因为我们有两个类别:行人和车辆
names: ['person', 'car']  # 类别名称列表

3. 开始训练

在YOLOv5目录下运行以下命令开始训练模型:

bash 复制代码
python train.py --img 640 --batch 16 --epochs 100 --data data.yaml --weights yolov5s.pt
  • --img 640: 设置输入图像的大小为640x640。
  • --batch 16: 每个批次的图像数量。
  • --epochs 100: 训练周期数。
  • --data data.yaml: 指定之前创建的数据集配置文件。
  • --weights yolov5s.pt: 从YOLOv5的小型版本(yolov5s)开始微调。

4. 使用DeepSORT进行跟踪

训练完成后,将YOLOv5与DeepSORT结合使用来进行目标跟踪。下面是一个简单的示例代码,展示了如何加载训练好的YOLOv5模型,并将其与DeepSORT集成用于视频中的目标检测和跟踪:

python 复制代码
import cv2
import torch
from deep_sort_pytorch.utils.parser import get_config
from deep_sort_pytorch.deep_sort import DeepSort
from yolov5.models.experimental import attempt_load
from yolov5.utils.general import non_max_suppression, scale_coords

# 加载YOLOv5模型
weights = 'runs/train/exp/weights/best.pt'  # 训练好的权重路径
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = attempt_load(weights, map_location=device)
model.to(device).eval()

# 加载DeepSORT
cfg = get_config()
cfg.merge_from_file("deep_sort_pytorch/configs/deep_sort.yaml")
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT, max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
                    nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
                    max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
                    use_cuda=True)

def process_video(video_path):
    cap = cv2.VideoCapture(video_path)
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(device)
        img = img.float() / 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        pred = model(img)[0]
        pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)

        for i, det in enumerate(pred):  # detections per image
            if len(det):
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()

                outputs = deepsort.update(det.cpu(), frame)

                for output in outputs:
                    bbox_2d = output[:4]
                    id = output[-1]
                    cv2.rectangle(frame, (int(bbox_2d[0]), int(bbox_2d[1])), (int(bbox_2d[2]), int(bbox_2d[3])), (0, 255, 0), 2)
                    cv2.putText(frame, f"ID: {id}", (int(bbox_2d[0]), int(bbox_2d[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

        cv2.imshow('Tracking', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    video_path = "path/to/video.mp4"
    process_video(video_path)

这个脚本会读取指定的视频文件,利用YOLOv5模型对每一帧进行物体检测,并使用DeepSORT算法对检测到的目标进行跟踪。对于每个检测结果,都会绘制边界框并在框上显示跟踪ID。

总结

训练一个能够识别车辆和行人的YOLOv5模型,并将其与DeepSORT结合起来实现多目标跟踪。

1. 安装依赖

首先,确保安装了必要的库:

bash 复制代码
pip install PyQt5 opencv-python torch torchvision deep_sort_pytorch

2. 创建主界面

使用PyQt5创建一个简单的GUI界面。

python 复制代码
import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QPushButton, QVBoxLayout, QWidget, QFileDialog, QCheckBox, QComboBox, QHBoxLayout
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import Qt, QTimer
import cv2
import numpy as np
import torch
from yolov5.models.experimental import attempt_load
from yolov5.utils.datasets import LoadImages
from yolov5.utils.general import non_max_suppression, scale_coords
from deep_sort_pytorch.utils.parser import get_config
from deep_sort_pytorch.deep_sort import DeepSort

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("YOLOv5+DeepSORT")
        self.setGeometry(100, 100, 800, 600)

        # Initialize UI components
        self.initUI()

        # Initialize YOLOv5 and DeepSORT
        self.model = None
        self.deepsort = None
        self.tracks = {}
        self.tracking_enabled = True
        self.show_trajectory = True
        self.selected_class = "car"

        # Timer for video processing
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.process_frame)
        self.cap = None

    def initUI(self):
        # Layout setup
        layout = QVBoxLayout()
        widget = QWidget()
        widget.setLayout(layout)
        self.setCentralWidget(widget)

        # Input and Output labels
        input_label = QLabel("输入")
        output_label = QLabel("输出")
        layout.addWidget(input_label)
        layout.addWidget(output_label)

        # Image labels for displaying frames
        self.input_label = QLabel(self)
        self.output_label = QLabel(self)
        layout.addWidget(self.input_label)
        layout.addWidget(self.output_label)

        # Buttons and checkboxes
        button_layout = QHBoxLayout()
        open_button = QPushButton("打开")
        open_button.clicked.connect(self.open_video)
        adapt_window_button = QPushButton("适应窗口")
        adapt_window_button.clicked.connect(self.adapt_window)
        predict_button = QPushButton("预测")
        predict_button.clicked.connect(self.start_prediction)
        stop_button = QPushButton("停止")
        stop_button.clicked.connect(self.stop_prediction)
        trajectory_checkbox = QCheckBox("绘制轨迹")
        trajectory_checkbox.setChecked(True)
        trajectory_checkbox.stateChanged.connect(self.toggle_trajectory)
        class_combobox = QComboBox()
        class_combobox.addItems(["person", "car"])
        class_combobox.currentIndexChanged.connect(self.select_class)
        person_checkbox = QCheckBox("person")
        car_checkbox = QCheckBox("car")
        person_checkbox.setChecked(False)
        car_checkbox.setChecked(True)
        person_checkbox.stateChanged.connect(lambda: self.select_class("person"))
        car_checkbox.stateChanged.connect(lambda: self.select_class("car"))

        button_layout.addWidget(open_button)
        button_layout.addWidget(adapt_window_button)
        button_layout.addWidget(predict_button)
        button_layout.addWidget(stop_button)
        button_layout.addWidget(trajectory_checkbox)
        button_layout.addWidget(QLabel("选择跟踪对象:"))
        button_layout.addWidget(person_checkbox)
        button_layout.addWidget(car_checkbox)

        layout.addLayout(button_layout)

    def open_video(self):
        options = QFileDialog.Options()
        file_name, _ = QFileDialog.getOpenFileName(self, "选择视频文件", "", "Video Files (*.mp4 *.avi);;All Files (*)", options=options)
        if file_name:
            self.cap = cv2.VideoCapture(file_name)
            self.load_model()

    def load_model(self):
        # Load YOLOv5 model
        self.model = attempt_load('yolov5s.pt', map_location='cpu')
        self.model.to('cpu').eval()

        # Load DeepSORT
        cfg = get_config()
        cfg.merge_from_file("deep_sort_pytorch/configs/deep_sort.yaml")
        self.deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT, max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
                                 nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
                                 max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
                                 use_cuda=True)

    def start_prediction(self):
        if self.cap is not None:
            self.timer.start(30)  # 30 ms delay between frames

    def stop_prediction(self):
        self.timer.stop()

    def process_frame(self):
        ret, frame = self.cap.read()
        if not ret:
            self.timer.stop()
            return

        # Process frame with YOLOv5 and DeepSORT
        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to('cpu')
        img = img.float() / 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        pred = self.model(img)[0]
        pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)

        for i, det in enumerate(pred):  # detections per image
            if len(det):
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()

                outputs = self.deepsort.update(det.cpu(), frame)

                self.tracks = {}
                for output in outputs:
                    bbox_2d = output[:4]
                    id = output[-1]
                    center = ((bbox_2d[0] + bbox_2d[2]) // 2, (bbox_2d[1] + bbox_2d[3]) // 2)

                    if id not in self.tracks:
                        self.tracks[id] = []
                    self.tracks[id].append(center)

                    if self.tracking_enabled and self.selected_class in ["person", "car"]:
                        cv2.rectangle(frame, (int(bbox_2d[0]), int(bbox_2d[1])), (int(bbox_2d[2]), int(bbox_2d[3])), (0, 255, 0), 2)
                        cv2.putText(frame, f"ID: {id}", (int(bbox_2d[0]), int(bbox_2d[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

                    if self.show_trajectory:
                        for i in range(1, len(self.tracks[id])):
                            point1 = self.tracks[id][i - 1]
                            point2 = self.tracks[id][i]
                            cv2.line(frame, point1, point2, (0, 0, 255), 2)

        # Display the processed frame
        height, width, channel = frame.shape
        bytes_per_line = 3 * width
        q_img = QImage(frame.data, width, height, bytes_per_line, QImage.Format_RGB888)
        self.output_label.setPixmap(QPixmap.fromImage(q_img))

    def toggle_trajectory(self, state):
        self.show_trajectory = state == Qt.Checked

    def select_class(self, class_name):
        self.selected_class = class_name

    def adapt_window(self):
        self.input_label.setScaledContents(True)
        self.output_label.setScaledContents(True)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
相关推荐
xiaoyaohou111 小时前
033、部署优化(四):模型编译与TVM在边缘设备上的应用
yolo
xiaoyaohou112 小时前
031、部署优化(二):TensorRT引擎构建与INT8量化推理
yolo
不熬夜的熬润之13 小时前
YOLOv5-OBB 训练避坑笔记
人工智能·yolo·计算机视觉
荪荪19 小时前
yolov8检测模型pt转rknn
人工智能·yolo·机器人·瑞芯微
时光之源19 小时前
一场关于红绿灯检测项目的辩论赛——YOLOv26
深度学习·yolo·grok·gemini·deepseek
乐迪信息21 小时前
智慧港口中AI防爆摄像机的船舶越线识别功能
大数据·人工智能·物联网·算法·目标跟踪
Uopiasd1234oo1 天前
AAttn区域注意力机制改进YOLOv26特征感知与表达能力提升
人工智能·yolo·目标跟踪
QQ676580081 天前
智慧AI甲骨文检测 目标检测图像数据集 甲骨文识别第10341期
人工智能·yolo·目标检测·目标跟踪·甲骨文检测·甲骨文识别
嵌入式吴彦祖1 天前
yolo简述和训练原理
yolo