代码实现效果:
代码:
import os
root_dir = r"C:\Users\ninan\Pytorch_study\torch_study\dataset\train"
# 这个地方如果需要修改成bees_iamge的话,把这个注释一下,用下面一行代码就行
target_dir = "ants_image"
# target_dir = "bees_image"
img_path = os.listdir(os.path.join(root_dir,target_dir))
#这里是把target_dir分割,得到ants 或者 bees 这样的数据集名称,以便于后期修改数据集名称
label = target_dir.split("_")[0]
#根据前面输入的target_dir得到对应的标签
out_dir = f"{label}_label"
#保存文件
for i in img_path:
file_name = i.split(".jpg")[0]
with open(os.path.join(root_dir,out_dir,f"{file_name}.txt"),"w") as fp:
fp.write(label)