[ pytorch训练 ]


1. 数据处理

from pathlib import Path
import os, csv, random

# 1. 分类
class_name_list = os.listdir('images')
class_to_num = {class_name: num for num, class_name in enumerate(class_name_list)}
# 文件夹名:排序 {'err': 0, 'suc': 1}

# 2. 合并所有数据
image_dir = []
for class_name in class_to_num:
    image_dir += Path(os.path.join("images", class_name)).glob("*.jpg")
# 合并到一起 [WindowsPath('images/err/ISIC_0211399.jpg'), ...]

# 3. 打散保存
random.shuffle(image_dir)
with open('pytorch.csv', mode='w', newline='') as f:
    write = csv.writer(f)
    for image in image_dir:
        class_name = image.parts[-2]
        label = class_to_num[class_name]
        write.writerow([image, label])