当前位置 主页 > 网站技术 > 代码类 >

    Pytorch 实现数据集自定义读取

    栏目:代码类 时间:2020-01-18 18:07

    以读取VOC2012语义分割数据集为例,具体见代码注释:

    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)
    

    Train.py

    import matplotlib.pyplot as plt
    import torch.utils.data as data
    import torchvision.transforms as transforms
    import numpy as np
    
    from PIL import Image
    from VocDataset import MyDataset
    
    #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)
    
    #从数据集中拿出一个批次的数据
    for i, data in enumerate(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))#创建一张与标签大小相同的图片,用以显示标签所对应的颜色
      for i in range(0, 480):
        for j in range(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()