目标跟踪数据集格式转换（按帧排序->按目标轨迹ID排序）

2023年11月4日 486点热度 3人点赞 0条评论

提出问题

因科研需求，需要对目标跟踪数据集的格式进行相关转换，方便模型的输入。这里以SAT-MTB数据集格式转为AIR-MOT数据集格式为例。即下图示意所示，将右边格式转换为左边格式。

解决方案

具体的实施方案如下面的代码步骤所示：

###第一步：改名
import os
import shutil

flder_path = r'xxx\SAT-MTB_Dataset\ship'
flder = os.listdir(flder_path)
#print (flder)
for i in flder:
    shutil.move(os.path.join(flder_path, i), 'xxx/SAT-MTB_Dataset/ship_s/' + i + '_s')

###第二步：删除原文件夹包含的det和seg任务的标签，并将“mot”改名为gt，里面的txt也改名
import os

flder_path1 = r'xxx\SAT-MTB_Dataset\ship_s'
flder1 = os.listdir(flder_path1)
print (flder1)

for j in flder1:
    shutil.rmtree(os.path.join(flder_path1, j, 'det'))
    shutil.rmtree(os.path.join(flder_path1, j, 'seg'))
    os.rename(os.path.join(flder_path1, j, 'mot'), os.path.join(flder_path1, j, 'gt'))
    os.rename(os.path.join(flder_path1, j, 'gt', j.split('_')[0]), os.path.join(flder_path1, j, 'gt', j + '.txt' ))

###第三步：修改txt格式并以AIR-MOT格式txt存储起来
import numpy as np
import os

def find_gt_folder(root_dir):
    for dirpath, dirnames, filenames in os.walk(root_dir):
        for dirname in dirnames:
            if dirname == "gt":
                # 找到以"gt"命名的文件夹
                return os.path.join(dirpath, dirname)

    # 如果没有找到，返回None
    return None

#print (find_gt_folder('G:/datasets/目标跟踪/SAT-MTB_Dataset/airplane_a/'))

fld_path = 'xxx/SAT-MTB_Dataset/ship_s/'

fld = os.listdir(fld_path)

for k in fld:
    temp = find_gt_folder(os.path.join(fld_path, k))

    data = np.loadtxt(temp + '/' + k + '.txt',dtype=np.float64, delimiter=',')

    # 根据第二列排序数据
    sorted_data = data[data[:, 1].argsort(kind='stable')]

    # 删除最后两列
    sorted_data = sorted_data[:, :-2]

    sorted_data[:, -1] = 1

    #print(sorted_data)

    # 将数据四舍五入为整数
    sorted_data = np.round(sorted_data).astype(int)

    # 将数据保存到以逗号分隔的txt文件
    np.savetxt(temp + '/' + k + 'a.txt', sorted_data, delimiter=',', fmt='%d')

###第四步：上述保存的txt同文件夹里有两个txt，需删除一个，保留一个与上上一级同名文件
import numpy as np
import os
import shutil

def find_gt_folder(root_dir):
    for dirpath, dirnames, filenames in os.walk(root_dir):
        for dirname in dirnames:
            if dirname == "gt":
                # 找到以"gt"命名的文件夹
                return os.path.join(dirpath, dirname)

    # 如果没有找到，返回None
    return None

ori_path = r'xxx\SAT-MTB_Dataset\ship_s'

ori_fld = os.listdir(ori_path)

for n in ori_fld:
    temp_ = find_gt_folder(os.path.join(ori_path, n))
    os.remove(os.path.join(temp_, n + '.txt'))
    os.rename(os.path.join(temp_, n + 'a.txt'), os.path.join(temp_, n + '.txt'))

附件
以上代码以jupyter的形式保存，可在这里下载。

Post Views: 518