import os
import random
# 打乱数据
def shuffle_data(data_list_path):
with open(data_list_path, 'r', encoding='utf-8') as f:
lines = f.readlines()
random.shuffle(lines)
print("训练数据:%d 张" % len(lines))
with open(data_list_path, 'w', encoding='utf-8') as f:
f.writelines(lines)
# 生成数据
def run(data_dir, train_list_path, test_list_path, label_path):
f_train = open(train_list_path, 'w', encoding='utf-8')
f_test = open(test_list_path, 'w', encoding='utf-8')
f_label = open(label_path, 'w', encoding='utf-8')
label_dict = dict()
class_label = 0
class_dirs = os.listdir(data_dir)
for class_dir in class_dirs:
print(class_dir)
if class_dir not in label_dict:
if class_dir !='.DS_Store':
label_dict[class_dir] = class_label
class_sum = 0
path = data_dir + "/" + class_dir
if path !='image/.DS_Store':
img_paths = os.listdir(path)
# print(img_paths)
# if os.listdir(path):
for img_path in img_paths:
name_path = path + '/' + img_path
if class_sum % 3 == 0:
f_test.write(name_path + " %d" % class_label + "\n")
else:
f_train.write(name_path + " %d" % class_label + "\n")
class_sum += 1
class_label += 1
label_str = str(label_dict).replace("'", '"')
f_label.write(label_str)
f_label.close()
f_train.close()
f_test.close()
print('create data list done!')
# 打乱数据
shuffle_data(train_list_path)
if __name__ == '__main__':
data_dir = 'image'
train_list = 'train_list.txt'
test_list = 'query_list.txt'
label_file = 'labels.txt'
run(data_dir, train_list, test_list, label_file)
使用百度的paddleclas制作train_list,query_list,labels
最新推荐文章于 2023-10-16 23:43:56 发布