Files
ieemoo-ai-zhanting/trainVal.py
2022-04-08 16:44:51 +08:00

25 lines
768 B
Python
Executable File

import random
import shutil
import os
def trainVal(path):
Totaldata = []
for dir in os.listdir(path):
if dir == 'images':
for fileName in os.listdir(os.sep.join([path, dir])):
Totaldata.append(fileName)
imagePath = path+'/'+'images'
random.shuffle(Totaldata)
val = Totaldata[:int(len(Totaldata)/10)]
train = Totaldata[int(len(Totaldata)/10):]
with open(path+'/'+'train.txt','w') as train_:
for data in train:
train_.write(imagePath+'/'+data)
with open(path+'/'+'val.txt','w') as val_:
for data in val:
val_.write(imagePath+'/'+data)
# return train, val
if __name__ == '__main__':
path = '/home/data/New_project/totaldata/shops_185'
trainVal(path)