Pytorch 实现数据集自定义读取
以读取VOC2012语义分割数据集为例，具体见代码注释：
VocDataset.py
fromPILimportImage
importtorch
importtorch.utils.dataasdata
importnumpyasnp
importos
importtorchvision
importtorchvision.transformsastransforms
importtime
#VOC数据集分类对应颜色标签
VOC_COLORMAP=[[0,0,0],[128,0,0],[0,128,0],[128,128,0],
[0,0,128],[128,0,128],[0,128,128],[128,128,128],
[64,0,0],[192,0,0],[64,128,0],[192,128,0],
[64,0,128],[192,0,128],[64,128,128],[192,128,128],
[0,64,0],[128,64,0],[0,192,0],[128,192,0],
[0,64,128]]
#颜色标签空间转到序号标签空间，就他妈这里浪费巨量的时间,这里还他妈的有问题
defvoc_label_indices(colormap,colormap2label):
"""AssignlabelindicesforPascalVOC2012Dataset."""
idx=((colormap[:,:,2]*256+colormap[:,:,1])*256+colormap[:,:,0])
#out=np.empty(idx.shape,dtype=np.int64)
out=colormap2label[idx]
out=out.astype(np.int64)#数据类型转换
end=time.time()
returnout
classMyDataset(data.Dataset):#创建自定义的数据读取类
def__init__(self,root,is_train,crop_size=(320,480)):
self.rgb_mean=(0.485,0.456,0.406)
self.rgb_std=(0.229,0.224,0.225)
self.root=root
self.crop_size=crop_size
images=[]#创建空列表存文件名称
txt_fname='%s/ImageSets/Segmentation/%s'%(root,'train.txt'ifis_trainelse'val.txt')
withopen(txt_fname,'r')asf:
self.images=f.read().split()
#数据名称整理
self.files=[]
fornameinself.images:
img_file=os.path.join(self.root,"JPEGImages/%s.jpg"%name)
label_file=os.path.join(self.root,"SegmentationClass/%s.png"%name)
self.files.append({
"img":img_file,
"label":label_file,
"name":name
})
self.colormap2label=np.zeros(256**3)
#整个循环的意思就是将颜色标签映射为单通道的数组索引
fori,cminenumerate(VOC_COLORMAP):
self.colormap2label[(cm[2]*256+cm[1])*256+cm[0]]=i
#按照索引读取每个元素的具体内容
def__getitem__(self,index):
datafiles=self.files[index]
name=datafiles["name"]
image=Image.open(datafiles["img"])
label=Image.open(datafiles["label"]).convert('RGB')#打开的是PNG格式的图片要转到rgb的格式下，不然结果会比较要命
#以图像中心为中心截取固定大小图像，小于固定大小的图像则自动填0
imgCenterCrop=transforms.Compose([
transforms.CenterCrop(self.crop_size),
transforms.ToTensor(),
transforms.Normalize(self.rgb_mean,self.rgb_std),#图像数据正则化
])
labelCenterCrop=transforms.CenterCrop(self.crop_size)
cropImage=imgCenterCrop(image)
croplabel=labelCenterCrop(label)
croplabel=torch.from_numpy(np.array(croplabel)).long()#把标签数据类型转为torch
#将颜色标签图转为序号标签图
mylabel=voc_label_indices(croplabel,self.colormap2label)
returncropImage,mylabel
#返回图像数据长度
def__len__(self):
returnlen(self.files)
Train.py
importmatplotlib.pyplotasplt
importtorch.utils.dataasdata
importtorchvision.transformsastransforms
importnumpyasnp
fromPILimportImage
fromVocDatasetimportMyDataset
#VOC数据集分类对应颜色标签
VOC_COLORMAP=[[0,0,0],[128,0,0],[0,128,0],[128,128,0],
[0,0,128],[128,0,128],[0,128,128],[128,128,128],
[64,0,0],[192,0,0],[64,128,0],[192,128,0],
[64,0,128],[192,0,128],[64,128,128],[192,128,128],
[0,64,0],[128,64,0],[0,192,0],[128,192,0],
[0,64,128]]
root='../data/VOCdevkit/VOC2012'
train_data=MyDataset(root,True)
trainloader=data.DataLoader(train_data,4)
#从数据集中拿出一个批次的数据
fori,datainenumerate(trainloader):
getimgs,labels=data
img=transforms.ToPILImage()(getimgs[0])
labels=labels.numpy()#tensor转numpy
labels=labels[0]#获得批次标签集中的一张标签图像
labels=labels.transpose((1,0))#数组维度切换，将第1维换到第0维，第0维换到第1维
##将单通道索引标签图片映射回颜色标签图片
newIm=Image.new('RGB',(480,320))#创建一张与标签大小相同的图片，用以显示标签所对应的颜色
foriinrange(0,480):
forjinrange(0,320):
sele=labels[i][j]#取得坐标点对应像素的值
newIm.putpixel((i,j),(int(VOC_COLORMAP[sele][0]),int(VOC_COLORMAP[sele][1]),int(VOC_COLORMAP[sele][2])))
#显示图像和标签
plt.figure("image")
ax1=plt.subplot(1,2,1)
ax2=plt.subplot(1,2,2)
plt.sca(ax1)
plt.imshow(img)
plt.sca(ax2)
plt.imshow(newIm)
plt.show()
以上这篇Pytorch实现数据集自定义读取就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持毛票票。
声明：本文内容来源于网络，版权归原作者所有，内容由互联网用户自发贡献自行上传，本网站不拥有所有权，未作人工编辑处理，也不承担相关法律责任。如果您发现有涉嫌版权的内容，欢迎发送邮件至：czq8825#qq.com（发邮件时，请将#更换为@）进行举报，并提供相关证据，一经查实，本站将立刻删除涉嫌侵权内容。