基于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_())
相关推荐
探物 AI1 天前
【感知·yolov11】零基础友好:大白话拆解 YOLOv11,像素变检测框底层逻辑一遍过
yolo
前网易架构师-高司机2 天前
带标注的升降杆识别数据集,识别率93.7%,3496张图,支持yolo,coco json,voc xml,文末有模型训练代码
yolo·数据集·高速·升降杆·抬杆··
广州灵眸科技有限公司2 天前
瑞芯微(EASY EAI)RV1126B yolov11-track多目标跟踪部署教程
linux·开发语言·网络·人工智能·yolo·机器学习·目标跟踪
深度学习lover2 天前
<数据集>yolo食物分类检测<目标检测>
人工智能·深度学习·yolo·目标检测·计算机视觉·食物分类识别
sheyuDemo2 天前
关于小土堆目标检测YOLOv5的一些报错
人工智能·深度学习·yolo·目标检测
深度学习lover2 天前
<数据集>yolo 桃子识别<目标检测>
人工智能·python·yolo·目标检测·计算机视觉·桃子识别
乐迪信息2 天前
乐迪信息:智慧港口安全,船舶AI视觉检测实现主动预警
大数据·人工智能·安全·计算机视觉·目标跟踪·视觉检测
白羊by2 天前
YOLOv8 官方损失函数详解(按任务分类)
人工智能·深度学习·算法·yolo·分类
赵药师2 天前
Cityscape数据集转YOLO
人工智能·深度学习·yolo
我材不敲代码3 天前
YOLOv4 核心内容详解
yolo