目标跟踪JDE模型训练自己的数据集

前言

        之前训练CenterTrack的时候制作了一个mot格式的数据集，现在要训练JDE，下面是我的自制数据集的文件目录。基本上就是按照MOT17数据集整理的，除了视频序列的名字不一样。img目录下存放着视屏序列的图片，gt目录下存放着按照mot格式标注好的标签数据（使用Darklabel工具）。

─custom    //根目录
├─annotations
├─test
│  ├─1
│  │  ├─gt
│  │  └─img
│  ├─3
│  │  ├─gt
│  │  └─img
│  └─....
└─train
    ├─1
    │  ├─gt
    │  └─img
    ├─11
    │  ├─gt
    │  └─img
    └─...

数据集格式调整

        将数据集从mot调整为JDE需要的格式。

        按照JDE官方提供Readme文件，完成数据集调整需要完成三个部分的任务。

1. 将数据集目录结构调整为如下结构

Caltech
   |——————images
   |        └——————00001.jpg
   |        |—————— ...
   |        └——————0000N.jpg
   └——————labels_with_ids
            └——————00001.txt
            |—————— ...
            └——————0000N.txt

        按照这个结构，图片数据不需要按照视频序列的方式来保存，所有图片统一保存到了images目录下。labels_with_ids下的每个txt和图片对应，其中保存着对应图片中的目标信息（类似于yolo格式），具体格式如下。class对应就是自己的类别，本博客使用的自制数据集仅有一个类别，identity则是在图片中对应目标的id，对应到mot格式的gt文件下则是id，后面bbox数据则是根据图像大小进行归一化后的。

[class] [identity] [x_center] [y_center] [width] [height]   //JDE标签文件格式

<frame>, <id>, <bb_left>, <bb_top>, <bb_width>, <bb_height>, <conf>, <x>, <y>, <z>   //mot 标签格式

        所以首先要做的就是标签文件从下面的mot格式转为上面的JDE需要的格式。下面的代码则用来完成这个任务，并按照train_half的方式，将视频序列划分为前后两部分，分别用于训练和测试。

2.生成用于训练的xxx.train和xxx.val文件

        xxx.train和xxx.val是用来记录训练数据的路径和测试数据的路径。如下图为JDE源码中打他目录下的mot17.train的样式。

3. 编写cfg目录下的ccmcpe.json文件

        ccmcpe.json文件记录有 数据集的目录、训练用.train文件和测试用.val文件。

        train对应编写.train文件

        test对应编写.val文件

        test_emb我不知道干嘛用的/(ㄒoㄒ)/~~

{
    "root":"mydataset",
    "train":
    {
        "mot17":"./data/custom.train"
    },
    "test_emb":   
    {
    },
    "test":
    {
        "mot17":"./data/custom.val"
    }
}

 4.好了再加一步，助君一步到位。

        下面的代码可以帮你调整好数据集格式，并生成对应.train和.val文件。

import os
import random
import shutil


def convert(box, size=(1280, 800)):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[2]) / 2.0
    y = (box[1] + box[3]) / 2.0
    w = box[2] - box[0]
    h = box[3] - box[1]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h


def modify_dataset(dataset_root, train_val_root, new_dataset_root):
    # 数据集根目录
    # dataset_root = 'custom/train/'
    # train_val_root = './data'  # cfg_save_dir
    # new_dataset_root = "./mydataset"  # new_dataset_root

    # 读取训练集图片路径
    train_image_dir = os.listdir(dataset_root)
    format_img, format_gt = "img", "gt"

    train_txt = os.path.join(train_val_root, "custom.train")
    val_txt = os.path.join(train_val_root, "custom.val")
    train_txt_fp = open(train_txt, "w")
    val_txt_fp = open(val_txt, "w")

    mydataset_image_dir = os.path.join(new_dataset_root, "images")
    mydataset_labels_dir = os.path.join(new_dataset_root, "labels_with_ids")
    os.makedirs(mydataset_image_dir, exist_ok=True)
    os.makedirs(mydataset_labels_dir, exist_ok=True)

    classes = 0
    image_cnt = 0
    new_id_cnt = 0
    for one_dir in train_image_dir:
        id_dict = {}

        image_dir = os.path.join(dataset_root, one_dir, format_img)
        label_txt_path = os.path.join(dataset_root, one_dir, format_gt, "gt.txt")
        label_fp = open(label_txt_path, "r")
        label_lines = label_fp.readlines()

        image_name_list = os.listdir(image_dir)
        image_num = len(image_name_list) // 2
        train_list = image_name_list[:image_num]
        val_list = image_name_list[image_num:]

        pre_image_id = -1
        for line in label_lines:
            line_list = line.strip().split(',')
            image_id = int(line_list[0])

            if image_id != pre_image_id:
                src_image_path = os.path.join(image_dir, "{:06d}.jpg".format(image_id))

                if not os.path.exists(src_image_path):
                    continue

                image_cnt += 1
                dst_image_path = os.path.join(mydataset_image_dir, "{:06d}.jpg".format(image_cnt))
                shutil.copy(src_image_path, dst_image_path)
                if "{:06d}.jpg".format(image_id) in train_list:
                    train_txt_fp.write(dst_image_path + '\n')
                elif "{:06d}.jpg".format(image_id) in val_list:
                    val_txt_fp.write(dst_image_path + '\n')

                pre_image_id = image_id

            img_label_txt = os.path.join(mydataset_labels_dir, "{:06d}.txt".format(image_cnt))
            img_label_txt_fp = open(img_label_txt, "a")

            identity = int(line_list[1])
            if (not id_dict.get(identity)) or (id_dict.get(identity) == 0):
                id_dict[identity] = new_id_cnt
                new_id_cnt += 1
            identity = id_dict[identity]

            bbox = [int(i) for i in line_list[2:6]]
            bbox[2] = bbox[0] + bbox[2]
            bbox[3] = bbox[1] + bbox[3]

            center_x, center_y, w, h = convert(bbox)

            img_label_txt_fp.write(
                "{},{:d},{:.6f},{:.6f},{:.6f},{:.6f}\n".format(classes, identity, center_x, center_y, w, h))


if __name__ == '__main__':
    custom_mot = "./custom/train"    # 指定自己的数据集目录   如果需要可以把自己的测试集test下的视频序列也拷贝到train下进行划分
    train_val_root = "./data"        # 这个train和val文件存放的目录，默认保存在jde源代码中data目录下
    new_dataset_root = "mydataset"   # 这个目录用来保存转换后数据集

    modify_dataset(custom_mot, train_val_root, new_dataset_root)

训练

        数据集整理完便可以开始着手训练了。目前JDE的源码直接跑会报错，坑啊，参考了这篇博客做出了一些修改，才得以正常运行。主要就是两处修改。

1.在这行代码前添加一个判断

mkdir_if_missing(weights_to+"/cfg")            # 添加这行
copyfile(cfg, weights_to + '/cfg/yolo3.cfg')   # 找到这行

2.注释掉如下代码

 开始训练

python train.py --cfg cfg/yolov3_576x320.cfg --batch-size 16

测试

...未进行