人工智能学习07--pytorch02--数据集取图片+以图片名为文件名的txt文档（label）

1、取图片

from torch.utils.data import Dataset
# import cv2
from PIL import Image
import os

class MyData(Dataset):
# help() 中说，想要使用 Dataset，需要重写。__getitem__ 和 __len__

    def __init__(self,root_dir,label_dir): # 初始化函数
        # 可以把self指定的变量给后面的函数使用（成为全局变量）
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir,self.label_dir) # 把两个路径拼接在一起(\\)
        self.img_path = os.listdir(self.path) # 获得所有地址

    def __getitem__(self, idx):
        # 要获取图片，首先创建图片地址列表
        img_name = self.img_path[idx]
        img_item_path = os.path.join(self.root_dir,self.label_dir,img_name)
        img = Image.open(img_item_path)
        label = self.label_dir
        return img,label

    def __len__(self):
        return len(self.img_path)

root_dir = 'dataset/train'
ants_label_dir = 'ants_image'
bees_label_dir = 'bees_image'
ants_dataset = MyData(root_dir,ants_label_dir)
bees_dataset = MyData(root_dir,bees_label_dir)
# 当类中有__getitem__方法，实例化对象A后调用A[0]会自动调用__getitem__的

#img,label=bees_dataset[1]
#img.show()

train_dataset = ants_dataset + bees_dataset  # 把两个数据集进行拼接
#len(train_dataset)
#245
#len(ants_dataset)
#124
#len(bees_dataset)
#121
img,label = train_dataset[123] # ant
img.show()
img,label = train_dataset[124] # bee
img.show()

假设这个train文件夹里面只有ants_image和bees_image

运行结果


可以看到这段代码成功地从数据集里面取出了图片

2、加label

用图片名字命名txt文件，内容是这个图片的分类。
因为图片太多了，所以用代码自动添加
只以ant为例，bee的就是把这个改改

import os
# 把图片的label 生成以图片名为文件名的txt文档
root_dir = 'dataset/train'
target_dir = 'ants_image'
img_path = os.listdir(os.path.join(root_dir, target_dir))
label = target_dir.split('_')[0]
out_dir = 'ants_label'
for i in img_path:
    file_name = i.split('.jpg')[0]
    with open(os.path.join(root_dir, out_dir,"{}.txt".format(file_name)),'w') as f:
        f.write(label)

可以看到下面这些都是自动生成的：