YOLOv5+DeepSORT实现基于检测的视频目标跟踪

2023年8月3日 407点热度 1人点赞 0条评论

提出问题

最近因为应用需求，需要实现基于视频的多目标跟踪任务。目标跟踪的发展脉络大体如下图所示，之前虽然做过类似的基于联合检测的目标跟踪，即图中的检测与特征提取联合的方式，如2020年提出的FairMOT网络等。然而，更加往前一点的基于DeepSORT的方式，即检测与特征提取分开的方式，这一类方式虽然会有累积误差的出现，但也能解决一些实际问题，因此，有必要研究一下这一类网络的操作流程。该篇博客以YOLOv5作为DeepSORT模型的Detections部分。

解决方法

YOLOv5和DeepSORT是分开训练的。因此，针对自己的数据集需要分别训练两个模型，一个基于YOLOv5（V5.0版本，其他版本需要自己调整）的检测模型，一个基于DeepSORT的跟踪模型。以下是整个项目的参考链接。

基于DarkLabel2.4自制目标检测数据集
Yolov5和DeepSORT的数据集处理和训练
xml数据集转换为跟踪数据集
 目标跟踪部分训练类别修改
这一步很重要，之前因为没将默认的类别751改为自己的类别而出错
代码参考链接
整个训练的框架参考（YOLOv5+DeepSORT）
YOLOv5 v5.0框架，这里也有其他YOLO系列版本的尝试，可以自行选择尝试。
DeepSORT算法理论研究
评价指标参考

具体实现步骤

训练YOLOv5模型
这一步参考该链接即可：https://github.com/spacewalk01/yolov5-fire-detection
训练自制数据集DeepSORT模型
这一步主要参考解决方法中的第2步，处理好数据集的格式。以下为相关代码：
（1）因为我的数据集含有中文路径，因此需要在整个xml文件中添加相关头约束，并重写相关信息。

import os

def add_xml_declaration(xml_folder):
    # 获取目录中的所有文件
    xml_files = [file for file in os.listdir(xml_folder) if file.endswith(".xml")]

    for xml_file in xml_files:
        xml_path = os.path.join(xml_folder, xml_file)

        # 读取XML文件内容
        with open(xml_path, "r", encoding="ISO-8859-1") as f:  # 使用 utf-8-sig 编码
            xml_content = f.read()

        # 检查XML文件是否已经包含XML声明
        if not xml_content.startswith('<?xml version="1.0" encoding="UTF-8"?>'):
            # 如果没有XML声明，则在内容最前面添加XML声明
            xml_content = '<?xml version="1.0" encoding="UTF-8"?>\n' + xml_content

            # 将修改后的内容写回XML文件
            with open(xml_path, "w", encoding="utf-8") as f:  # 使用 utf-8-sig 编码
                f.write(xml_content)

if __name__ == "__main__":
    xml_folder_path = "./yolo_tracking-3.0/datasets/labels"  # 替换为包含XML文件的目录路径
    add_xml_declaration(xml_folder_path)

（2）提取数据集目标部分

# -*- coding: utf-8 -*-

import cv2
import lxml.etree as ET
import numpy as np
import os
import argparse
import xml.dom.minidom

def main():
    # JPG文件的地址
    img_path = './yolo_tracking-3.0/datasets/images/'
    # XML文件的地址
    anno_path = "./yolo_tracking-3.0/datasets/labels/"
    # 存结果的文件夹

    cut_path = ".yolo_tracking-3.0/datasets/crops/"
    if not os.path.exists(cut_path):
        os.makedirs(cut_path)
    # 获取文件夹中的文件
    imagelist = os.listdir(img_path)
    #print(imagelist)
    for image in imagelist:
        image_pre, ext = os.path.splitext(image)
        #print(image_pre)
        img_file = img_path + image
        #print(img_file)
        img = cv2.imread(img_file)
        print (img.shape)
        xml_file = anno_path + image_pre + '.xml'
        #print (xml_file)
        # DOMTree = xml.dom.minidom.parse(xml_file)
        # collection = DOMTree.documentElement
        # objects = collection.getElementsByTagName("object")

        # tree = ET.parse(xml_file)
        # root = tree.getroot()
        # 创建一个自定义解析器，并指定使用支持中文编码的解析器
        parser = ET.XMLParser(encoding="utf-8")

        # 使用自定义解析器解析XML文件
        tree = ET.parse(xml_file, parser=parser)

        # 继续对解析后的XML进行操作
        root = tree.getroot()

        # if root.find('object') == None:
        #     return
        obj_i = 0
        for obj in root.iter('object'):
            obj_i += 1
            print(obj_i)
            cls = obj.find('name').text
            xmlbox = obj.find('bndbox')
            b = [int(float(xmlbox.find('xmin').text)), int(float(xmlbox.find('ymin').text)),
                 int(float(xmlbox.find('xmax').text)),
                 int(float(xmlbox.find('ymax').text))]
            img_cut = img[b[1]:b[3], b[0]:b[2], :]
            path = os.path.join(cut_path, cls)
            # 目录是否存在,不存在则创建
            mkdirlambda = lambda x: os.makedirs(x) if not os.path.exists(x) else True
            mkdirlambda(path)
            try:
                cv2.imwrite(os.path.join(cut_path, cls, '{}_{:0>2d}.jpg'.format(image_pre, obj_i)), img_cut)
            except:
                continue

            print("&&&&")

if __name__ == '__main__':
    main()

(3)对以上裁剪得到的数据集根据每个视频/种类进行分类

import os
from PIL import Image
from shutil import copyfile, copytree, rmtree, move

PATH_DATASET = "./yolo_tracking-3.0/datasets/crops/"  # 需要处理的文件夹
PATH_NEW_DATASET = "./yolo_tracking-3.0/datasets/after_crops"  # 处理后的文件夹
PATH_ALL_IMAGES = PATH_NEW_DATASET + '/all_images'
PATH_TRAIN = PATH_NEW_DATASET + '/train'
PATH_TEST = PATH_NEW_DATASET + '/test'

# 定义创建目录函数
def mymkdir(path):
    path = path.strip()  # 去除首位空格
    path = path.rstrip("\\")  # 去除尾部 \ 符号
    isExists = os.path.exists(path)  # 判断路径是否存在
    if not isExists:
        os.makedirs(path)  # 如果不存在则创建目录
        print(path + ' 创建成功')
        return True
    else:
        # 如果目录存在则不创建，并提示目录已存在
        print(path + ' 目录已存在')
        return False

class BatchRename():
    '''
    批量重命名文件夹中的图片文件
    '''

    def __init__(self):
        self.path = PATH_DATASET  # 表示需要命名处理的文件夹

    # 修改图像尺寸
    def resize(self):
        for aroot, dirs, files in os.walk(self.path):
            # aroot是self.path目录下的所有子目录（含self.path）,dir是self.path下所有的文件夹的列表.
            filelist = files  # 注意此处仅是该路径下的其中一个列表
            # print('list', list)

            # filelist = os.listdir(self.path) #获取文件路径
            total_num = len(filelist)  # 获取文件长度（个数）

            for item in filelist:
                if item.endswith('.jpg'):  # 初始的图片的格式为jpg格式的（或者源文件是png格式及其他格式，后面的转换格式就可以调整为自己需要的格式即可）
                    src = os.path.join(os.path.abspath(aroot), item)

                    # 修改图片尺寸到128宽*256高
                    im = Image.open(src)
                    out = im.resize((128, 256), Image.ANTIALIAS)  # resize image with high-quality
                    out.save(src)  # 原路径保存

    def rename(self):

        for aroot, dirs, files in os.walk(self.path):
            # aroot是self.path目录下的所有子目录（含self.path）,dir是self.path下所有的文件夹的列表.
            filelist = files  # 注意此处仅是该路径下的其中一个列表
            # print('list', list)

            # filelist = os.listdir(self.path) #获取文件路径
            total_num = len(filelist)  # 获取文件长度（个数）

            i = 1  # 表示文件的命名是从1开始的
            for item in filelist:
                if item.endswith('.jpg'):  # 初始的图片的格式为jpg格式的（或者源文件是png格式及其他格式，后面的转换格式就可以调整为自己需要的格式即可）
                    src = os.path.join(os.path.abspath(aroot), item)

                    # 根据图片名创建图片目录
                    dirname = str(item.split('_')[0])
                    # 为相同车辆创建目录
                    # new_dir = os.path.join(self.path, '..', 'bbox_all', dirname)
                    new_dir = os.path.join(PATH_ALL_IMAGES, dirname)
                    if not os.path.isdir(new_dir):
                        mymkdir(new_dir)

                    # 获得new_dir中的图片数
                    num_pic = len(os.listdir(new_dir))

                    dst = os.path.join(os.path.abspath(new_dir),
                                       dirname + 'C1T0001F' + str(num_pic + 1) + '.jpg')
                    # 处理后的格式也为jpg格式的，当然这里可以改成png格式    C1T0001F见mars.py filenames 相机ID，跟踪指数
                    # dst = os.path.join(os.path.abspath(self.path), '0000' + format(str(i), '0>3s') + '.jpg')    这种情况下的命名格式为0000000.jpg形式，可以自主定义想要的格式
                    try:
                        copyfile(src, dst)  # os.rename(src, dst)
                        print('converting %s to %s ...' % (src, dst))
                        i = i + 1
                    except:
                        continue
            print('total %d to rename & converted %d jpgs' % (total_num, i))

    def split(self):
        # ---------------------------------------
        # train_test
        images_path = PATH_ALL_IMAGES
        train_save_path = PATH_TRAIN
        test_save_path = PATH_TEST
        if not os.path.isdir(train_save_path):
            os.mkdir(train_save_path)
            os.mkdir(test_save_path)

        for _, dirs, _ in os.walk(images_path, topdown=True):
            for i, dir in enumerate(dirs):
                for root, _, files in os.walk(images_path + '/' + dir, topdown=True):
                    for j, file in enumerate(files):
                        if (j == 0):  # test dataset；每个车辆的第一幅图片
                            print("序号：%s  文件夹： %s  图片：%s 归为测试集" % (i + 1, root, file))
                            src_path = root + '/' + file
                            dst_dir = test_save_path + '/' + dir
                            if not os.path.isdir(dst_dir):
                                os.mkdir(dst_dir)
                            dst_path = dst_dir + '/' + file
                            move(src_path, dst_path)
                        else:
                            src_path = root + '/' + file
                            dst_dir = train_save_path + '/' + dir
                            if not os.path.isdir(dst_dir):
                                os.mkdir(dst_dir)
                            dst_path = dst_dir + '/' + file
                            move(src_path, dst_path)
        rmtree(PATH_ALL_IMAGES)

if __name__ == '__main__':
    demo = BatchRename()
    demo.resize()
    demo.rename()
    demo.split()

运行解决方法中的第三步整个训练的框架参考（YOLOv5+DeepSORT）中的demo.py文件，实现完整的目标跟踪。其中，根据项目需要，我们做了相应的小改动，具体如下，增加了保存跟踪结果的txt文本。当然，相应的其他代码几个部分需要做微小调动。后续将上传完整代码。

from AIDetector_pytorch import Detector
import imutils
import cv2
import os

def _xyxy_to_tlwh(bbox_xyxy):
    x1, y1, x2, y2 = bbox_xyxy

    t = x1
    l = y1
    w = int(x2-x1)
    h = int(y2-y1)
    return t, l, w, h

def main():

    name = 'demo'

    det = Detector()
    input_video = './Yolov5-Deepsort-main/yolov5-5.0/val_video/hs.mp4'
    file_path = "./Yolov5-Deepsort-main/pre_hs_yolov5s.txt"
    cap = cv2.VideoCapture(input_video)
    name_out = os.path.basename(input_video)
    fps = int(cap.get(5))
    print('fps:', fps)
    t = int(1000/fps)

    videoWriter = None

    while True:

        # try:
        flag, im = cap.read()

        if im is None:
            break

        result_ = det.feedCap(im) ##utils/BaseDetector.py

        if flag:
            pos_frame = int(cap.get(1))
            print ('当前处理帧：',pos_frame)

        content = result_['frame_cot']
        with open(file_path, 'a+') as f:
            if len(content):
                for item in content:
                    item_out = _xyxy_to_tlwh(item[:4])
                    line = '{frame},{id},{x1},{y1},{w},{h},1,-1,-1,-1\n'.format(frame=pos_frame, id=item[5], x1=item_out[0], y1=item_out[1], w=item_out[2], h=item_out[3])
                    #print (line)
                    f.write(line)

        result = result_['frame']

        result = imutils.resize(result, height=640) ###原500
        if videoWriter is None:
            fourcc = cv2.VideoWriter_fourcc(
                'm', 'p', '4', 'v')  # opencv3.0
            videoWriter = cv2.VideoWriter(
                './results/{}'.format(name_out), fourcc, fps, (result.shape[1], result.shape[0]))

        videoWriter.write(result)
        #cv2.imshow(name, result)
        #cv2.waitKey(t)

        #if cv2.getWindowProperty(name, cv2.WND_PROP_AUTOSIZE) < 1:
            # 点x退出
        #    break
        # except Exception as e:
        #     print(e)
        #     break

    cap.release()
    videoWriter.release()
    #cv2.destroyAllWindows()

if __name__ == '__main__':

    main()

Post Views: 397