将图像数据读写标记成txt文件

# -*- coding: UTF-8 -*- 
import os

def IsSubString(SubStrList,Str):
    flag=True
    for substr in SubStrList:
        if not(substr in Str):
            flag=False
    
    return flag

#扫面文件
def GetFileList(FindPath,FlagStr=[]):
    FileList=[]
    FileNames=os.listdir(FindPath)
    if len(FileNames)>0:
        for fn in FileNames:
            if len(FlagStr)>0:
                if IsSubString(FlagStr,fn):
                    fullfilename=os.path.join(FindPath,fn)
                    FileList.append(fullfilename)
            else:
                fullfilename=os.path.join(FindPath,fn)
                FileList.append(fullfilename)
    
    if len(FileList)>0:
        FileList.sort()
        
    return FileList



train_txt=open('train.txt','w')
imgfile=GetFileList('/home/disk/lyw/AuoPose/DATA/images/train/carbody/')#将数据集放在与.py文件相同目录下
for img in imgfile:
    str1=img+' '+'1'+'\n'        #用空格代替转义字符 \t 
    train_txt.writelines(str1)
    

imgfile=GetFileList('/home/disk/lyw/AuoPose/DATA/images/train/carhead/')
for img in imgfile:
    str2=img+' '+'0'+'\n'
    train_txt.writelines(str2)
#train_txt.close()

imgfile=GetFileList('/home/disk/lyw/AuoPose/DATA/images/train/cartrail/')
for img in imgfile:
    str3=img+' '+'2'+'\n'
    train_txt.writelines(str3)
train_txt.close()



#测试集文件列表
test_txt=open('val.txt','w')
imgfile=GetFileList('/home/disk/lyw/AuoPose/DATA/images/val/carbody/')#将数据集放在与.py文件相同目录下
for img in imgfile:
    str4=img+' '+'1'+'\n'
    test_txt.writelines(str4)
    

imgfile=GetFileList('/home/disk/lyw/AuoPose/DATA/images/val/carhead/')
for img in imgfile:
    str5=img+' '+'0'+'\n'
    test_txt.writelines(str5)
    
imgfile=GetFileList('/home/disk/lyw/AuoPose/DATA/images/val/cartrail/')
for img in imgfile:
    str6=img+' '+'2'+'\n'
    test_txt.writelines(str6)
test_txt.close()

print("ok")