将代码脚本与train、val、test文件夹放在同一目录下,在该目录下打开终端运行脚本即可
import os
import sys
import random
dir_path = os.getcwd()
print(dir_path)
lst = ['test', 'train', 'val']
for i in lst :
f_path = os.path.join(dir_path ,i)
with open(f_path + '.txt', 'w') as ft:
print(f_path + '.txt')
for path, dir_lst, file_lst in os.walk(f_path):
for file in file_lst:
if os.path.splitext(file)[1] in ['.jpg', '.jpeg', '.png', '.JPG']:
root = os.path.join(path, file+'\n')
print(root)
ft.writelines(root)
# 如果没有对应的txt文件
txt_name = os.path.splitext(file)[0] + ".txt"
txt = os.path.join(path, txt_name)
if not os.path.exists(txt):
with open(txt, 'w') as t:
t.close()
# 打乱
os.system("shuf train.txt -o train.txt")
print('shuf train.txt -o train.txt')
os.system("shuf test.txt -o test.txt")
print('shuf test.txt -o test.txt')
os.system("shuf val.txt -o val.txt")
print('shuf val.txt -o val.txt')