PyTorch深度学习快速入门教程（绝对通俗易懂！）【小土堆】

P7 p6. Dataset类代码实战

对P7写入txt文档略做了些调整：

（ps：增加了检查文件夹是否存在，不存在则创建的功能，这样就省去了自己自己再去一个个创建的功能）

import os

root_dir = “your_path”

# 使用字典的“键值对”将目标图片文件夹与标签文件对应

targets = {

"ants_image": "ants_label",

"bees_image": "bees_label",

"spiders_image": "spiders_label"

}

# 创建文件夹、遍历

for target_dir, out_dir in targets.items():

# 检查标签文件夹是否存在，没有则创建

if not os.path.exists(os.path.join(root_dir, out_dir)):

os.makedirs(os.path.join(root_dir, out_dir))

img_paths = os.listdir(os.path.join(root_dir, target_dir))

label = target_dir.split('_')【0】

for img_path in img_paths:

file_name = img_path.split('.jpg')【0】

with open(os.path.join(root_dir, out_dir, f"{file_name}.txt"), 'w+') as f:

f.write(label)