新增类型描述文件处理脚本

dataset_process.py

@@ -16,8 +16,12 @@ def unpickle(file):
     return dict
 
 
-train_dic = './dataset/CIFAR-10/train/'
-test_dic = './dataset/CIFAR-10/test/'
+# 定义解压后batch文件夹
+file_dir = './dataset/CIFAR-10/cifar-10-batches-py'
+dataset_dir = './dataset/CIFAR-10'
+train_dic = f'{dataset_dir}/train/'
+test_dic = f'{dataset_dir}/test/'
+
 # 判断文件夹是否存在，不存在的话创建文件夹
 if not os.path.exists(train_dic):
     os.mkdir(train_dic)
@@ -63,10 +67,15 @@ def cifar10_img(file_dir):
     print('Finish transforming to image')
 
     # 处理描述文件
-    # meta_name = file_dir + '/' + 'batches.meta'
-    # meta_dict = unpickle(meta_name)
-    # label_names = [str(item) for item in meta_dict[b'label_names']]
-    # print(meta_name + ' is done')
+    meta_name = file_dir + '/' + 'batches.meta'
+    meta_dict = unpickle(meta_name)
+    label_names = [str(item) for item in meta_dict[b'label_names']]
+    print(meta_name + ' is done')
+    f = open(f'{dataset_dir}/class.txt', 'w')  # 创建类型描述文件
+    for label_name in label_names:
+        line = label_name + '\n'
+        f.write(line)
+    f.close()
 
 
 def gen_label_txt(label_txt_path, img_dir):
@@ -92,15 +101,9 @@ def write_class_list(classes, class_txt_path):
 
 
 if __name__ == '__main__':
-    # 定义解压后batch文件夹
-    file_dir = './dataset/CIFAR-10/cifar-10-batches-py'
-    dataset_dir = './dataset/CIFAR-10'
-    train_txt = dataset_dir + '/train.txt'
-    test_txt = dataset_dir + '/test.txt'
-
     # 处理解压后文件
     cifar10_img(file_dir)
 
     # 生成标签文件
-    gen_label_txt(train_txt, train_dic)
-    gen_label_txt(test_txt, test_dic)
+    gen_label_txt(dataset_dir + '/train.txt', train_dic)
+    gen_label_txt(dataset_dir + '/test.txt', test_dic)