YOLOv10结合StrongSORT+OSNet实现目标跟踪于重识别【附代码】

已于 2024-06-30 00:24:56 修改
阅读量2.6k 收藏 11
点赞数 15
分类专栏： YOLOv10基础 文章标签： YOLO 目标跟踪 目标检测 深度学习
于 2024-06-29 23:53:14 首次发布
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/Dora_blank/article/details/140072160
版权

文章目录

前言

本文将介绍如何使用YOLOv10和StrongSORT+OSNet实现目标跟踪+重识别+轨迹绘制，并详细讲解代码中的每个参数和函数的作用。

功能概述

1. 选择需要跟踪的视频路径
2. 加载 YOLOv10 模型并进行目标检测
3. 使用 StrongSORT+OSNet实现目标跟踪于重识别
4. 根据跟踪内容可视化绘制跟踪轨迹

必要环境

  1. 配置yolov10环境 可参考往期博客
    地址：https://blog.csdn.net/Dora_blank/article/details/139302363?spm=1001.2014.3001.5502

一、代码结构

1. 参数定义

parser = argparse.ArgumentParser()
# 检测参数1`
parser.add_argument('--weights', default=r"yolov10n.pt", type=str, help='weights path')
parser.add_argument('--source', default=r"video_1.mp4", type=str, help='video(.mp4)path')
parser.add_argument('--save', default=r"./save", type=str, help='save img or video path')
parser.add_argument('--vis', default=True, action='store_true', help='visualize image')
parser.add_argument('--conf_thre', type=float, default=0.5, help='conf_thre')
parser.add_argument('--iou_thre', type=float, default=0.5, help='iou_thre')

# 跟踪参数
parser.add_argument('--track_model', default=r"./track_models/osnet_x0_25_msmt17.pt", type=str, help='track model')
parser.add_argument('--max_dist', type=float, default=0.2, help='max dist')
parser.add_argument('--max_iou_distance', type=float, default=0.7, help='max_iou_distance')
parser.add_argument('--max_age', type=int, default=30, help='max_age')
parser.add_argument('--n_init', type=int, default=3, help='n_init')
parser.add_argument('--nn_budget', type=int, default=100, help='nn_budget')
# 解析参数
opt = parser.parse_args()

参数作用如下：
目标检测：
–weights：YOLOv10权重路径
–source：输入视频路径
–save：结果保存路径
–vis：可视化跟踪结果
–conf_thre：置信度阈值
–iou_thre：IoU阈值

目标跟踪：
–track_model：选择osnet重识别权重路径
–max_dist: 关联检测框和跟踪框的最大距离阈值，当距离小于这个阈值时,检测框和跟踪框会被认为是匹配的
–max_iou_distance: 关联的最大IOU距离阈值，当IOU小于这个阈值时,检测框和跟踪器会被认为是匹配的
–max_age: 跟踪框最大保留的帧数 超过时则会删除
–n_init: 定义了跟踪器需要被匹配到的最小连续帧数,只有连续被匹配到这么多帧的跟踪器才会被认为是一个有效的跟踪目标
–nn_budget: 每个跟踪器保留的外观特征向量的最大数量,用于计算外观相似性

2. 设备选择

根据是否有GPU可用，选择使用CUDA或CPU

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

3. 定义检测器类

定义一个Detector类来封装检测和跟踪逻辑并初始化模型参数

class Detector(object):
    def __init__(self, weight_path, conf_threshold=0.2, iou_threshold=0.5):
        self.device = device
        self.model = YOLOv10(weight_path).to(self.device)
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold

        self.tracker = StrongSORT(
            opt.track_model,
            self.device,
            max_dist=opt.max_dist,
            max_iou_distance=opt.max_iou_distance,
            max_age=opt.max_age,
            n_init=opt.n_init,
            nn_budget=opt.nn_budget,
        )

        self.trajectories = {}
        self.max_trajectory_length = 5

4. 获取目标框颜色

在绘制轨迹和边框时，我们需要为每个目标分配不同的颜色

@staticmethod
def get_color(idx):
    idx = idx * 3
    color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
    return color

5. 检测与跟踪

用于对图像中检测到的目标进行跟踪，并绘制跟踪框和跟踪轨迹

def detect_image(self, img_bgr):
    image = img_bgr.copy()
    results = self.model(img_bgr, conf=self.conf_threshold, iou=self.iou_threshold)
    boxes = results[0].boxes.xyxy.cpu().numpy()  # xyxy format
    confidences = results[0].boxes.conf
    class_preds = results[0].boxes.cls

    confidences_expanded = confidences.unsqueeze(1)
    class_preds_expanded = class_preds.unsqueeze(1)
    boxes_tensor = torch.from_numpy(boxes).to(class_preds_expanded.device)
    xywhs = xyxy2xywh(boxes_tensor)
    online_targets = self.tracker.update(xywhs.cpu(), confidences_expanded.cpu(), class_preds_expanded.cpu(), image)

    for t in online_targets:
        tlwh = [t[0], t[1], t[2] - t[0], t[3] - t[1]]
        tid = t[4]
        cls = t[5]
        xmin, ymin, xmax, ymax = tlwh[0], tlwh[1], tlwh[0] + tlwh[2], tlwh[1] + tlwh[3]
        class_pred = int(cls)
        color = self.get_color(class_pred + 4)
        center = (int(xmin+xmax)//2, int(ymin + ymax)//2)
        bbox_label = results[0].names[class_pred]
        cv2.rectangle(img_bgr, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color, 2)

        if tid not in self.trajectories:
            self.trajectories[tid] = deque(maxlen=self.max_trajectory_length)
        self.trajectories[tid].appendleft(center)

        # 截断轨迹长度
        if len(self.trajectories[tid]) > self.max_trajectory_length:
            self.trajectories[tid] = self.trajectories[tid][:self.max_trajectory_length]

        for i in range(1, len(self.trajectories[tid])):

            if self.trajectories[tid][i - 1] is None or self.trajectories[tid][i] is None:
                continue

            thickness = int(np.sqrt(64 / float(i + 1)))
            cv2.line(img_bgr, self.trajectories[tid][i-1],
                     self.trajectories[tid][i], color,
                     thickness)

        # 显示类名和跟踪ID
        cv2.putText(img_bgr, f"{bbox_label} {int(tid)}",
                    (int(xmin), int(ymin-5)), cv2.FONT_HERSHEY_COMPLEX, 0.6, color, 2)

    return img_bgr

6. 绘制轨迹线代码详解

这段代码会遍历每个跟踪框的中心点，使用OpenCV的cv2.line函数绘制轨迹线，线的长度以及厚度会随距离增加而减小，轨迹长度大于self.max_trajectory_length会被截断，以便更清晰地显示在图像中

for j in range(1, len(self.trajectories[tid])):
    if self.trajectories[tid][j - 1] is None or self.trajectories[tid][j] is None:
        continue

    thickness = int(np.sqrt(64 / float(j + 1)))
    cv2.line(img_bgr, self.trajectories[tid][j - 1],
             self.trajectories[tid][j], color,
             thickness)

7. 目标框格式转换详解

这段代码块目的是将检测结果转化为可传入StrongSORT中的格式

1. 扩展置信度和类别预测的维度

confidences_expanded = confidences.unsqueeze(1)
class_preds_expanded = class_preds.unsqueeze(1)

2. 将检测框转换为Tensor并移动到与预测相同的设备

boxes_tensor = torch.from_numpy(boxes).to(class_preds_expanded.device)

3. 将坐标格式转为中心点、宽度、高度

xywhs = xyxy2xywh(boxes_tensor)

4. 使用 StrongSORT的更新函数 得到跟踪框信息

online_targets = self.tracker.update(xywhs.cpu(), confidences_expanded.cpu(), class_preds_expanded.cpu(), image)

二、完整代码

完整代码如下：

import cv2
import torch
from ultralytics import YOLOv10
import os
import argparse
from strong_sort.strong_sort import StrongSORT
from collections import deque
import numpy as np


parser = argparse.ArgumentParser()
# 检测参数1`
parser.add_argument('--weights', default=r"yolov10n.pt", type=str, help='weights path')
parser.add_argument('--source', default=r"video_1.mp4", type=str, help='video(.mp4)path')
parser.add_argument('--save', default=r"./save", type=str, help='save img or video path')
parser.add_argument('--vis', default=True, action='store_true', help='visualize image')
parser.add_argument('--conf_thre', type=float, default=0.5, help='conf_thre')
parser.add_argument('--iou_thre', type=float, default=0.5, help='iou_thre')

# 跟踪参数
parser.add_argument('--track_model', default=r"./track_models/osnet_x0_25_msmt17.pt", type=str, help='track model')
parser.add_argument('--max_dist', type=float, default=0.2, help='max dist')
parser.add_argument('--max_iou_distance', type=float, default=0.7, help='max_iou_distance')
parser.add_argument('--max_age', type=int, default=30, help='max_age')
parser.add_argument('--n_init', type=int, default=3, help='n_init')
parser.add_argument('--nn_budget', type=int, default=100, help='nn_budget')
# 解析参数
opt = parser.parse_args()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

def xyxy2xywh(x):
    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
    y[:, 2] = x[:, 2] - x[:, 0]  # width
    y[:, 3] = x[:, 3] - x[:, 1]  # height
    return y

class Detector(object):
    def __init__(self, weight_path, conf_threshold=0.2, iou_threshold=0.5):
        self.device = device
        self.model = YOLOv10(weight_path).to(self.device)
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold

        self.tracker = StrongSORT(
            opt.track_model,
            self.device,
            max_dist=opt.max_dist,
            max_iou_distance=opt.max_iou_distance,
            max_age=opt.max_age,
            n_init=opt.n_init,
            nn_budget=opt.nn_budget,
        )

        self.trajectories = {}
        self.max_trajectory_length = 5

    @staticmethod
    def get_color(idx):
        idx = idx * 3
        color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
        return color

    def detect_image(self, img_bgr):
        image = img_bgr.copy()
        results = self.model(img_bgr, conf=self.conf_threshold, iou=self.iou_threshold)
        boxes = results[0].boxes.xyxy.cpu().numpy()  # xyxy format
        confidences = results[0].boxes.conf
        class_preds = results[0].boxes.cls

        confidences_expanded = confidences.unsqueeze(1)
        class_preds_expanded = class_preds.unsqueeze(1)
        boxes_tensor = torch.from_numpy(boxes).to(class_preds_expanded.device)
        xywhs = xyxy2xywh(boxes_tensor)
        online_targets = self.tracker.update(xywhs.cpu(), confidences_expanded.cpu(), class_preds_expanded.cpu(), image)

        for t in online_targets:
            tlwh = [t[0], t[1], t[2] - t[0], t[3] - t[1]]
            tid = t[4]
            cls = t[5]
            xmin, ymin, xmax, ymax = tlwh[0], tlwh[1], tlwh[0] + tlwh[2], tlwh[1] + tlwh[3]
            class_pred = int(cls)
            color = self.get_color(class_pred + 4)
            center = (int(xmin+xmax)//2, int(ymin + ymax)//2)
            bbox_label = results[0].names[class_pred]
            cv2.rectangle(img_bgr, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color, 2)

            if tid not in self.trajectories:
                self.trajectories[tid] = deque(maxlen=self.max_trajectory_length)
            self.trajectories[tid].appendleft(center)

            # 截断轨迹长度
            if len(self.trajectories[tid]) > self.max_trajectory_length:
                self.trajectories[tid] = self.trajectories[tid][:self.max_trajectory_length]

            for i in range(1, len(self.trajectories[tid])):

                if self.trajectories[tid][i - 1] is None or self.trajectories[tid][i] is None:
                    continue

                thickness = int(np.sqrt(64 / float(i + 1)))
                cv2.line(img_bgr, self.trajectories[tid][i-1],
                         self.trajectories[tid][i], color,
                         thickness)

            # 显示类名和跟踪ID
            cv2.putText(img_bgr, f"{bbox_label} {int(tid)}",
                        (int(xmin), int(ymin-5)), cv2.FONT_HERSHEY_COMPLEX, 0.6, color, 2)

        return img_bgr


# Example usage
if __name__ == '__main__':
    model = Detector(weight_path=opt.weights, conf_threshold=opt.conf_thre, iou_threshold=opt.iou_thre)

    capture = cv2.VideoCapture(opt.source)
    fps = capture.get(cv2.CAP_PROP_FPS)
    size = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    outVideo = cv2.VideoWriter(os.path.join(opt.save, os.path.basename(opt.source).split('.')[-2] + "_out.mp4"), fourcc,
                               fps, size)

    while True:
        ret, frame = capture.read()
        if not ret:
            break
        img_vis = model.detect_image(frame)
        outVideo.write(img_vis)

        img_vis = cv2.resize(img_vis, None, fx=1, fy=1, interpolation=cv2.INTER_NEAREST)
        cv2.imshow('track', img_vis)
        cv2.waitKey(30)

    capture.release()
    outVideo.release()

三、效果展示

StrongSORT跟踪效果

在这里插入图片描述

总结

本期博客就到这里啦，喜欢的小伙伴们可以点点关注，感谢！

最近经常在b站上更新一些有关目标检测的视频，大家感兴趣可以来看看 https://b23.tv/1upjbcG

学习交流群：995760755

[空--白]
关注
  • 15
    点赞
  • 踩
  • 11
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 0
    评论
专栏目录
YOLOv10结合DeepSORT实现目标检测与轨迹跟踪
Dora_blank的博客
05-30 2108
本文将介绍如何结合YOLOv10和DeepSORT来实现目标检测与轨迹跟踪，并详细讲解代码中的每个参数和函数的作用。
推荐项目：Yolov7_StrongSORT_OSNet - 实时目标检测与跟踪的高效解决方案
gitblog_00062的博客
04-26 429
推荐项目：Yolov7_StrongSORT_OSNet - 实时目标检测与跟踪的高效解决方案 项目地址:https://gitcode.com/mikel-brostrom/Yolov7_StrongSORT_OSNet 该项目是一个基于先进的计算机视觉模型和算法的开源实现，结合了YOLOv7的高精度目标检测、StrongSORT的强大追踪能力和OSNet的优秀特征学习。通过这个项目，开发者和研...
YOLOv5 + StrongSORT with OSNet
sc1434404661的博客
09-21 6343
YOLOv5 + StrongSORT with OSNet：YOLOv5检测器 + StrongSORT跟踪算法 + OSNet行人重识别模型
手把手教你搭建YOLOv5 + StrongSORT with OSNet环境
AI浩
11-04 2832
YOLOv5 + StrongSORT with OSNet模型由YOLOv5检测器 、StrongSORT跟踪算法、OSNet行人重识别模型组成项目地址：由于现在的代码中没有yolov5和reid代码，所以需要下载源码加入到YOLOv5 + StrongSORT with OSNet项目中。yolov5下载地址：下载代码后，放到yolov5的文件夹中，如下图：”
机器视觉教学：YOLOv5 + StrongSORT with OSNet
noopier的博客
08-16 435
YOLOv5 + StrongSORT with OSNet：YOLOv5检测器 + StrongSORT跟踪算法 + OSNet行人重识别模型项目地址：https://github.com/mikel-brostrom/Yolov5_StrongSORT_OSNet。
毕设教学：YOLOv5 + StrongSORT with OSNet
m0_12378的博客
07-31 159
YOLOv5 + StrongSORT with OSNet：YOLOv5检测器 + StrongSORT跟踪算法 + OSNet行人重识别模型项目地址：https://github.com/mikel-brostrom/Yolov5_StrongSORT_OSNet。
YOLOV7 + StrongSORT 实现目标检测与跟踪
08-16 8988
环境windows 10 64bitpython 3.8pytorch1.7.1 + cu101简介前面，我们介绍过 基于YOLOv5和DeepSort的目标跟踪，最近大神又在刚出的 YOLOv7 的基础上加上了目标跟踪，跟踪部分使用的是基于 OSNet 的 StrongSORT，项目地址: https://github.com/mikel-brostrom/Yolov7_StrongSORT_O...
yolov8 strongSORT多目标跟踪工具箱BOXMOT
白起
10-12 5324
由mikel brostrom提供。在以前的版本中，有yolov5+deepsort（版本v3-v5）， yolov8+strongsort（版本v6-v9），直至演变到v10，名称BOXMOT。BOXMOT提供三种对象检测器：yolov8， yolo_nas， yolox；支持多个跟踪器：BoTSORT, DeepOCSORT, OCSORT, Hybridsort, ByteTrack, StrongSORT。以前常见的DeepSort在此由增强型StrongSORT替代。
YOLOv10结合StrongSORT+OSNet实现目标跟踪于重识别附代码
06-29
YOLOv10结合StrongSORT+OSNet实现目标跟踪于重识别【附代码】 实现如下功能： 参数解析与设置：通过argparse模块解析命令行参数，包括YOLOv10模型权重路径、视频源路径、输出保存路径、置信度阈值、IOU阈值以及...
移动跟踪-基于YOLOv10+DeepSort实现视频中移动目标跟踪算法-附项目源码+流程教程-优质项目实战.zip
06-07
本项目以YOLOv10和DeepSORT两种技术为核心，实现了对视频中移动目标的有效跟踪，提供了详细的源码和教程，是学习和实践这一技术的良好资源。 首先，YOLO（You Only Look Once）是一种实时的目标检测系统，其设计...
基于YBTrack YOLOv5 + BYTE实现的无人机目标跟踪系统源码+文档说明
06-20
基于YBTrack YOLOv5 + BYTE实现的无人机目标跟踪系统源码+文档说明，适用于希望学习不同技术领域的小白或进阶学习者。可作为毕设项目、课程设计、大作业、工程实训或初期项目立项。 【项目介绍】：YBTrack是一种...
使用YOLOv9+DeepSort实现的目标跟踪算法python源码.zip
05-15
使用YOLOv9+DeepSort实现的目标跟踪算法python源码.zip使用YOLOv9+DeepSort实现的目标跟踪算法python源码.zip使用YOLOv9+DeepSort实现的目标跟踪算法python源码.zip使用YOLOv9+DeepSort实现的目标跟踪算法python源码...
基于YOLOV5-7.0+DeepSort的目标追踪算法
10-30
DeepSORT是一种强大的多目标跟踪算法，结合YOLOv5 7.0版本的目标检测能力，可以实现高效准确的实时目标跟踪。 基于 YOLOV5 和 DeepSort 的目标追踪算法是一种结合了目标检测和运动预测的方法，用于在视频中实现多...
YOLO训练触发早停
最新发布
weixin_52236586的博客
07-25 306
日志：or use。
YOLO训练过程日志，数值都小于1原因？
weixin_52236586的博客
07-24 374
在YOLO等目标检测模型中，总损失是多个损失项的组合。通常，总损失会进行归一化以平衡各个损失项的影响。这种组合的归一化处理通常会使总损失值较小，但它提供了一个有效的衡量标准来评估模型的整体性能。总之，损失值小于1是正常的，归一化处理是为了确保训练过程中的稳定性和有效性。关注损失值的变化趋势比关注绝对值更为重要，这样可以更好地了解模型的训练情况。如果你的YOLOv5训练日志显示损失值如。
yolov3 ImageNet 预训练权重下载
qq_46396470的博客
07-22 171
记得一定要翻到页面最底下的Assets那里才能看到这些权重。
基于yolov8的实例分割全过程详细教学
qq_36738025的博客
07-25 180
（3）修改数据集的配置文件ultralytics-main\ultralytics\cfg\datasets\coco128-seg.yaml ，换上自己数据集的路径以及标签名称，和目标检测一样；修改model的配置文件ultralytics-main\ultralytics\cfg\models\v8\yolov8-seg.yaml 下的类别，将80类换成自己的类别个数。这里由于保存的是json所以要更新为txt格式，记得修改路径。在当前文件夹下运行，记得更换一下下面的路径。
yolov7+strongsort结合起来的python代码
06-10
以下是使用YOLOv7和Strong Sort结合起来实现目标检测和多目标跟踪的Python代码示例： ```python import cv2 from yolov7.detect import detect_image from strong_sort.sort import Sort # 初始化Strong Sort跟踪器 tracker = Sort() # 加载YOLOv7模型 model = 'yolov7/yolov7.weights' config = 'yolov7/yolov7.cfg' classes = 'yolov7/coco.names' # 定义阈值和NMS参数 conf_threshold = 0.5 nms_threshold = 0.4 # 加载模型和类别信息 net = cv2.dnn.readNetFromDarknet(config, model) classes = open(classes).read().strip().split('\n') # 打开摄像头或视频文件 cap = cv2.VideoCapture(0) while True: # 读取视频帧 ret, frame = cap.read() if not ret: break # 调用YOLOv7进行目标检测 detections = detect_image(frame, net, classes, conf_threshold, nms_threshold) # 预测每个物体的位置和类别 boxes = [] scores = [] for detection in detections: x, y, w, h, score, label = detection boxes.append([x, y, x+w, y+h]) scores.append(score) # 调用Strong Sort进行多目标跟踪 boxes = tracker.update(boxes, scores) # 在图像中绘制跟踪结果 for box in boxes: x1, y1, x2, y2, id = box cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) cv2.putText(frame, str(id), (int(x1), int(y1)), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 255, 0), 2) # 显示图像 cv2.imshow('frame', frame) # 按下q键退出程序 if cv2.waitKey(1) == ord('q'): break # 释放摄像头或视频文件 cap.release() # 关闭所有窗口 cv2.destroyAllWindows() ``` 这段代码中，我们首先初始化了Strong Sort跟踪器，并加载了YOLOv7模型和类别信息。然后，我们使用OpenCV读取摄像头或视频文件中的每一帧，在每一帧中调用YOLOv7进行目标检测，得到每个物体的位置和类别信息。接着，我们将检测结果传递给Strong Sort跟踪器进行多目标跟踪，得到每个物体的ID和最新的位置信息。最后，我们在图像中绘制跟踪结果，并显示在屏幕上。
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

[空--白]

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付：¥1
获取中
扫码支付

您的余额不足，请更换扫码支付或充值

打赏作者

实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值