Python 生成训练集的数据格式

简介

一般我们在用深度学习进行训练时，首先就是把数据处理成我们想要的格式，一些分类网络的图片训练格式一般都是train.txt 里面保存这图片的路径，以及对应的label。本篇博客就来介绍一下如何用python 来生成相应的txt文件。

1，准备数据

首先我们要有一个数据集，这里以market1501数据集为例，该数据集是行人重识别（reid）的主流数据集，该数据集的格式如下所示：

前面四位数都是表示同一类，后面的其实无所谓，每个字段都是用“_"来分开，所以我们要把它生成这样的格式：
/path/ro/your/market1501/bounding_box_train/0002_c1s1_000451_03.jpg 0
前面是图片的路径，后面是该图对应的标签。

2 python 代码

import os
import sys
train_path = "/path/to/your/market1501/bounding_box_train"
# 返回每个图片的路径
def file_name(file_dir):
    L=[]
    for root, dirs, files in os.walk(file_dir):
        for file in files:
            if os.path.splitext(file)[1] == '.jpg' or os.path.splitext(file)[1] == '.JPG' or os.path.splitext(file)[1] == '.JPEG' or os.path.splitext(file)[1] == '.jpeg'or os.path.splitext(file)[1] == '.png':
                fpath=os.path.join(root, file)
                L.append(fpath)
    return L
label_dic={}

def main():
    #打开一个txt文件，用来保存生成的图片路径和相应的label
    list_file_train = open("./train_reid.txt", 'w')
    #label必须是从0开始，不然会报错！
    labelidx=0
    files=file_name(train_path)
    allrecoder=[]
    for file in files:
        alls=file.split('/')
        #获取图片的名字
        name=alls[-1]
        #用”_“作为分隔符，得到前面的0002
        name=name.split('_')
        labn=name[0]
        #放到字典里作为key
        ret=label_dic.get(labn)
        label=None
        #生成相应的标签
        if(ret==None):
            label_dic[labn]=labelidx
            label=labelidx
            labelidx = labelidx+1
        else:
            #find label index at ret
            label=ret
        #把图片的路径以及标签存放在一个列表中
        recode=file+" "+str(label)
        allrecoder.append(recode)
    flen=len(allrecoder)
    indices = np.arange(flen)
    #随机重新进行排列
    np.random.shuffle(indices)
    for idx in range(flen):
        flen=flen-1
        print(flen)
       #把字符串写到list_file_train打开的文件中
        list_file_train.write(allrecoder[indices[idx]]+"\n")
    list_file_train.close()
 
    print("finished!",label_dic)
main()

上述代码便是把数据集生成训练格式的方法。