从xml数据集到FairMOT数据集转换

最新推荐文章于 2022-10-20 13:29:23 发布

zsffuture

最新推荐文章于 2022-10-20 13:29:23 发布

阅读量1.6k

点赞数 8

分类专栏：深度学习文章标签： FairMOT数据制作

本文链接：https://blog.csdn.net/weixin_42398658/article/details/110873985

版权

深度学习专栏收录该内容

61 篇文章 583 订阅

订阅专栏

这篇博客介绍了如何将UA-DETRAC数据集转换为FairMOT数据集的格式，用于多目标跟踪训练。首先获取每个视频的宽高信息，然后解析XML标签文件，按照FairMOT的要求整理成txt格式，包括目标类别、ID、归一化的中心坐标和宽高。最后，将转换后的标签文件保存到指定目录。

摘要由CSDN通过智能技术生成

这是我写的数据集，大家需要修改路径即可，当然希望大家可以看一下转换思路，以后就会转换啦，授之以鱼不如授之以渔

"""
把UA-DETRAC数据集制作FairMOT数据集的格式，以此进行训练
FairMOT数据集的格式： class、id、x_center/img_width、y_center/img_height、w/img_width、h/img_height 
    class               ：目标类别
    id                  ：目标id
    x_center/img_width  ：归一化中心列坐标
    y_center/img_height ：归一化中心行坐标
    w/img_width         ：归一化宽
    h/img_height        ：归一化高
"""
import xml.etree.ElementTree as et
import os.path as osp
import os
import shutil
import numpy as np
import cv2 as cv
import xml.etree.ElementTree as et

def mkdirs(d):
    # if not osp.exists(d):
    if not osp.isdir(d):
        os.makedirs(d)

# 获取每个视频的宽高信息
def get_img_WH(imgpath):
    dirname = os.listdir(imgpath)
    print('---dirname: ', dirname)
    # 遍历图片名称，分别获取每个文件夹图片的宽高
    dicWH = {}
    for name in dirname:
        path = os.path.join(imgpath,name)
        print("---imgpath ：", path)
        img = cv.imread(path+'/img00001.jpg')
        h, w, c = img.shape
        dicWH[name] = [w,h]
    return dicWH

# 获取每个视频的标签xml信息,返回是二维列表[[帧号，id号，左上角顶点列坐标，左上角顶点行坐标，宽，高]]
def get_xml_data(xmlpath,xmlname):
    tree = et.parse(os.path.join(xmlpath,xmlname+'.xml'))
    root = tree.getroot()
    datalist = []
    nodes = root.findall("frame")
    for node in nodes:
        # print("node.tag-->{0}, node.attrib-->{1}".format(node.tag, node.attrib))
        sub_node = node.findall("target_list")
        for node1 in sub_node:
            # print("node1.tag-->{0}, node1.attrib-->{1}".format(node1.tag, node1.attrib))
            three_node = node1.findall("target")
            for node2 in three_node:
                # print("node2.tag-->{0}, node2.attrib-->{1}".format(node2.tag, node2.attrib))
                four_node1 = node2.findall("box")
                for node3 in four_node1:
                    # print("node3.tag-->{0}, node3.attrib-->{1}".format(node3.tag, node3.attrib))
                    datalist.append([int(node.attrib['num']), int(node2.attrib['id']), float(node3.attrib['left']), float(node3.attrib['top']), float(node3.attrib['width']), float(node3.attrib['height'])])
    return datalist                



data_root = r"/home/ubuntu/zsfprj/MOTdata/DETRAC/"
seq_root = data_root + 'images/train'
label_root = data_root + 'labels_with_ids/train'
# 原始标签文件目录
label_ori = "/home/ubuntu/zsfprj/MOTdata/DETRAC-Train-Annotations-XML/"
namelist = os.listdir(label_ori)
# 获取图片的宽高信息
dicwh = {}
dicwh = get_img_WH(seq_root)

if not os.path.isdir(label_root):
    mkdirs(label_root)
else:  # 如果之前已经生成过: 递归删除目录和文件, 重新生成目录
    shutil.rmtree(label_root)
    os.makedirs(label_root)


print("Dir %s made" % label_root)

seqs = os.listdir(seq_root)

print(seqs)

tid_curr = 0
tid_last = -1
total_track_id_num = 0  # 计算数据集中的所有跟踪目标的数目
for seq in seqs:  # 每段视频都对应一个gt.txt
    print("Process %s, \n" % seq, end='')
    # seq_root = data_root + 'MOT/images/train'
    # label_root = data_root + 'MOT/labels_with_ids/train'
    # seq_info_path = osp.join(seq_root, seq, 'seqinfo.ini')  # 提取每个数据的info信息 /media/ckq/data/kitti/MOT/images/train
    # print(seq_info_path)
    # with open(seq_info_path) as seq_info_h:  # 读取 *.ini 文件
    #     seq_info = seq_info_h.read()
    #     seq_width = int(seq_info[seq_info.find('imWidth=') + 8:seq_info.find('\nimHeight')])  # 视频的宽
    #     seq_height = int(seq_info[seq_info.find('imHeight=') + 9:seq_info.find('\nimExt')])  # 视频的高
        # print('seq_width:',seq_width)
        # print('seq_height:', seq_height)
    # 获取每个文件夹的图片的宽高信息
    seq_width = dicwh[seq][0]
    seq_height = dicwh[seq][1]
    #gt_txt = osp.join(seq_root, seq, 'gt', 'gt.txt')  # 读取GT文件
    # print(gt_txt)  #打印路径
    # gt = np.loadtxt(gt_txt, dtype=np.str, delimiter=',')  # 加载成np格式
    # 获取原始标签的信息
    # gt = np.loadtxt(gt_txt, dtype=np.float64, delimiter=',')  # 加载成np格式
    gt = np.array(get_xml_data(label_ori, seq))
    print(gt)  # 打印文本内容
    print('gt.T\n')
    print(gt.T)  # 也是打印文本内容
    temp = gt.T[:2, :]
    print("temp = \n")
    print(temp)
    idx = np.lexsort(temp)  # 优先按照track id排序(对视频帧进行排序, 而后对轨迹ID进行排序),这里是numpy的二维排序的接口
    print(idx)
    gt = gt[idx, :]

    tr_ids = set(gt[:, 1])
    print("%d track ids in seq %s" % (len(tr_ids), seq))
    total_track_id_num += len(tr_ids)  # track id统计数量如何正确计算？

    # seq_label_root = osp.join(label_root, seq, 'img1')
    seq_label_root = osp.join(label_root, seq)
    mkdirs(seq_label_root)

    # 读取GT数据的每一行(一行即一条数据)
    # for fid, tid, x, y, w, h, mark, cls, vis_ratio in gt:
    for fid, tid, bbox_left, bbox_top, bbox_width, bbox_height in gt:
        # height, width, length , location_x,location_y,location_z , rotation_y in gt:
        # frame_id, track_id, top, left, width, height, mark, class, visibility ratio
        # if cls != 3:  # 我们需要Car的标注数据
        # if type != 3:  # 我们需要Car的标注数据
        #     continue

        # if mark == 0:  # mark为0时忽略(不在当前帧的考虑范围)
        #     continue

        # if vis_ratio <= 0.2:
        #     continue

        fid = int(fid)
        tid = int(tid)

        # 判断是否是同一个track, 记录上一个track和当前track
        if not tid == tid_last:  # not 的优先级比 == 高
            tid_curr += 1
            tid_last = tid
        # 由于kitti标签与训练标签参数有点不同 需要自己计算 x y w h
        w = bbox_width
        h = bbox_height
        x = int(bbox_left + 0.5)
        y = int(bbox_top + 0.5)

        # bbox中心点坐标
        x += w / 2
        y += h / 2

        # 网label中写入track id, bbox中心点坐标和宽高(归一化到0~1)
        # 第一列的0是默认只对一种类别进行多目标检测跟踪(0是类别)
        label_str = '0 {:d} {:.6f} {:.6f} {:.6f} {:.6f}\n'.format(
            tid_curr,
            x / seq_width,  # center_x
            y / seq_height,  # center_y
            w / seq_width,  # bbox_w
            h / seq_height)  # bbox_h
        # print(label_str.strip())

        label_f_path = osp.join(seq_label_root, 'img{:05d}.txt'.format(fid))
        with open(label_f_path, 'a') as f:  # 以追加的方式添加每一帧的label
            f.write(label_str)

print("Total %d track ids in this dataset" % total_track_id_num)
print('Done')