pytorch自己创建数据集

一、数据准备

假设图像在".\datasets\DogCat\train"中，文件名称是‘dog’和‘cat’，数据集是二分类，现在需要创建文件名和标签对应关系。

fold=r'.\datasets\DogCat\train'
wfold=r'.\datasets\DogCat'
flist=os.listdir(fold) #文件名称列表

#文件描述
with open(os.path.join(wfold,'train.txt'),'w') as fw:
	for f in flist:
		if 'dog' in f:
			string=f+' 0\n' #标签和文件名用空格隔开
		else:
			string=f+' 1\n'
		fw.write(string)

二、继承Dataset类

from torch.utils import data
class MyDogCat(data.Dataset):
	#初始化方法
	def __init__(self,root,datatxt):
		super(MyDogCat,self).__init__()
		with open(datatxt) as fr: #打开描述文件
			lines=fr.readlines()
		self.names=[] #名称
		self.labels=[] #标签
		self.root=root #图像文件夹目录
		for line in lines:
			l=line.rstrip('\n')
			fname,lab=l.split(' ') #以空格分隔
			self.names.append(fname)
			self.labels.append(lab)
		
	#根据索引取出图像和标签
	def __getitem__(self,index):
		fname=self.names[index]
		lab=self.labels[index]
		img=Image.open(os.path.join(self.root,fname)).convert('RGB')
		return img,lab
	
	#定义数据集规模
	def __len__(self):
		return len(self.names)
		

资料来源：《深度卷积神经网络原理与实践》