在我的torchvision库里介绍的博文(https://www.cnblogs.com/yjphhw/p/9773333.html)里说了对pytorch的dataset的定义方式。
本文相当于实现一个自定义的数据集,而这正是我们在做自己工程所需要的,我们总是用自己的数据嘛。
继承 from torch.utils.data import Dataset 类
然后实现 __len__(self) ,和 __getitem__(self,idx) 两个方法。以及数据增强也可以写入,数据增强想了想还是放到了Dataset里,
习惯上可能与常用的不同,但是觉得由于每种数据都有自己的增强方法所以,增强方法可以和数据集绑定到一起的。
接上一节我们通过切割,获取了2217个图像切片。
这就是我的FarmDataset
from torch.utils.data import Dataset, DataLoader from PIL import Image,ImageEnhance from osgeo import gdal from torchvision import transforms import glob import torch as tc import numpy as np class FarmDataset(Dataset): def __init__(self,istrain=True,isaug=True): self.istrain=istrain self.trainxformat='./data/train/data1500/*.bmp' self.trainyformat='./data/train/label1500/*.bmp' self.testxformat='./data/test/*.png' self.fns=glob.glob(self.trainxformat) if istrain else glob.glob(self.testxformat) self.length=len(self.fns) self.transforms=transforms self.isaug=isaug def __len__(self): #total length is 2217 return self.length def __getitem__(self,idx): if self.istrain: imgxname=self.fns[idx] sampleimg = Image.open(imgxname) imgyname=imgxname.replace('data1500','label1500') targetimg = Image.open(imgyname).convert('L') #sampleimg.save('original.bmp') #data augmentation if self.isaug: sampleimg,targetimg=self.imgtrans(sampleimg,targetimg) #check the result of dataautmentation #sampleimg.save('sampletmp.bmp') #targetimg.save('targettmp.bmp') sampleimg=transforms.ToTensor()(sampleimg) #targetimg=transforms.ToTensor()(targetimg).squeeze(0).long() targetimg=np.array(targetimg) targetimg=tc.from_numpy(targetimg).long() #to tensor #print(sampleimg.shape,targetimg.shape) return sampleimg,targetimg else: return gdal.Open(self.fns[idx]) def imgtrans(self,x,y,outsize=1024): '''input is a PIL image image dataaugumentation return also aPIL image。 ''' #rotate should consider y degree=np.random.randint(360) x=x.rotate(degree,resample=Image.NEAREST,fillcolor=0) y=y.rotate(degree,resample=Image.NEAREST,fillcolor=0) #here should be carefull, in case of label damage #random do the input image augmentation if np.random.random()>0.5: #sharpness factor=0.5+np.random.random() enhancer=ImageEnhance.Sharpness(x) x=enhancer.enhance(factor) if np.random.random()>0.5: #color augument factor=0.5+np.random.random() enhancer=ImageEnhance.Color(x) x=enhancer.enhance(factor) if np.random.random()>0.5: #contrast augument factor=0.5+np.random.random() enhancer=ImageEnhance.Contrast(x) x=enhancer.enhance(factor) if np.random.random()>0.5: #brightness factor=0.5+np.random.random() enhancer=ImageEnhance.Brightness(x) x=enhancer.enhance(factor) #img flip transtypes=[Image.FLIP_LEFT_RIGHT,Image.FLIP_TOP_BOTTOM, Image.ROTATE_90,Image.ROTATE_180,Image.ROTATE_270] transtype=transtypes[np.random.randint(len(transtypes))] x = x.transpose(transtype) y = y.transpose(transtype) #img resize between 0.8-1.2 w,h=x.size factor=1+np.random.normal()/5 if factor>1.2: factor=1.2 if factor<0.8: factor=0.8 #print(factor,x.size) x=x.resize((int(w*factor),int(h*factor)),Image.NEAREST) y=y.resize((int(w*factor),int(h*factor)),Image.NEAREST) #random crop w,h=x.size stx=np.random.randint(w-outsize) sty=np.random.randint(h-outsize) #print((stx,sty,outsize,outsize)) x=x.crop((stx,sty,stx+outsize,sty+outsize)) #stx,sty,width,height y=y.crop((stx,sty,stx+outsize,sty+outsize)) #print(x.size,y.size) return x,y #return outsized pil image if __name__=='__main__': d=FarmDataset(istrain=True) x,y=d[2216] print(x.shape) print(y.shape)
输入的是个1500x1500的图像,输出的是增强后的1024x1024后的图像。
其实对于分割问题来看,以后这个就可以作为一个模板,修改修改就可以换到另一个数据集中。
放几张图片:
原始图像:
进行数据增强后可以得到的一系列:
经过check 发现没有的问题通过测试。
本站文章如无特殊说明,均为本站原创,如若转载,请注明出处:Pytorch dataset自定义【直播】2019 年县域农业大脑AI挑战赛—数据准备(二),Dataset定义 - Python技术站