pytorch蜜蜂蚂蚁数据集处理python代码

代码实现效果：

代码：

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)