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)
可以看到下面这些都是自动生成的: