Pytorch实现数据集自定义读取-创新互联
以读取VOC2012语义分割数据集为例,具体见代码注释:
创新互联公司是一家专业提供辉南企业网站建设,专注与成都网站设计、网站制作、H5响应式网站、小程序制作等业务。10年已为辉南众多企业、政府机构等服务。创新互联专业网站设计公司优惠进行中。VocDataset.py
from PIL import Image import torch import torch.utils.data as data import numpy as np import os import torchvision import torchvision.transforms as transforms import time #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]] #颜色标签空间转到序号标签空间,就他妈这里浪费巨量的时间,这里还他妈的有问题 def voc_label_indices(colormap, colormap2label): """Assign label indices for Pascal VOC2012 Dataset.""" 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() return out class MyDataset(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' if is_train else 'val.txt') with open(txt_fname, 'r') as f: self.images = f.read().split() #数据名称整理 self.files = [] for name in self.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) #整个循环的意思就是将颜色标签映射为单通道的数组索引 for i, cm in enumerate(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) return cropImage,mylabel #返回图像数据长度 def __len__(self): return len(self.files)
另外有需要云服务器可以了解下创新互联scvps.cn,海内外云服务器15元起步,三天无理由+7*72小时售后在线,公司持有idc许可证,提供“云服务器、裸金属服务器、高防服务器、香港服务器、美国服务器、虚拟主机、免备案服务器”等云主机租用服务以及企业上云的综合解决方案,具有“安全稳定、简单易用、服务可用性高、性价比高”等特点与优势,专为企业上云打造定制,能够满足用户丰富、多元化的应用场景需求。
网站题目:Pytorch实现数据集自定义读取-创新互联
当前地址:http://pwwzsj.com/article/dogspo.html