Pytorch 实现数据集自定义读取


Posted in Python onJanuary 18, 2020

以读取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()

以上这篇Pytorch 实现数据集自定义读取就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持三水点靠木。

Python 相关文章推荐
Python计算程序运行时间的方法
Dec 13 Python
Python3处理文件中每个词的方法
May 22 Python
Python编程之黑板上排列组合,你舍得解开吗
Oct 30 Python
Python中常用信号signal类型实例
Jan 25 Python
Python实现EXCEL表格的排序功能示例
Jun 25 Python
Python Django中间件,中间件函数,全局异常处理操作示例
Nov 08 Python
Python栈的实现方法示例【列表、单链表】
Feb 22 Python
python 插入日期数据到Oracle实例
Mar 02 Python
如何配置关联Python 解释器 Anaconda的教程(图解)
Apr 30 Python
Python多分支if语句的使用
Sep 03 Python
python 实现学生信息管理系统的示例
Nov 28 Python
宝塔更新Python及Flask项目的部署
Apr 11 Python
使用pytorch搭建AlexNet操作(微调预训练模型及手动搭建)
Jan 18 #Python
selenium 多窗口切换的实现(windows)
Jan 18 #Python
pytorch实现建立自己的数据集(以mnist为例)
Jan 18 #Python
使用PyTorch实现MNIST手写体识别代码
Jan 18 #Python
Pytorch之finetune使用详解
Jan 18 #Python
pytorch 修改预训练model实例
Jan 18 #Python
Pytorch自己加载单通道图片用作数据集训练的实例
Jan 18 #Python
You might like
在PHP中执行系统外部命令
2006/10/09 PHP
第四章 php数学运算
2011/12/30 PHP
sql注入与转义的php函数代码
2013/06/17 PHP
解析htaccess伪静态的规则
2013/06/18 PHP
thinkphp实现面包屑导航(当前位置)例子分享
2014/05/10 PHP
php创建类并调用的实例方法
2019/09/25 PHP
Javascript里使用Dom操作Xml
2007/01/22 Javascript
DWZ table的原生分页浅谈
2013/03/01 Javascript
简单选项卡 js和jquery制作方法分享
2014/02/26 Javascript
JavaScript极简入门教程(三):数组
2014/10/25 Javascript
JavaScript实现自动弹出窗口并自动关闭窗口的方法
2015/08/06 Javascript
zTree树形插件异步加载方法详解
2017/06/14 Javascript
vue移动端裁剪图片结合插件Cropper的使用实例代码
2017/07/10 Javascript
JS自定义函数实现时间戳转换成date的方法示例
2017/08/27 Javascript
浅谈在Vue-cli里基于axios封装复用请求
2017/11/06 Javascript
IE浏览器下JS脚本提交表单后,不能自动提示问题解决方法
2019/06/04 Javascript
优雅的使用javascript递归画一棵结构树示例代码
2019/09/22 Javascript
python通过pil为png图片填充上背景颜色的方法
2015/03/17 Python
使用python在本地电脑上快速处理数据
2017/06/22 Python
利用selenium 3.7和python3添加cookie模拟登陆的实现
2017/11/20 Python
pycharm执行python时,填写参数的方法
2018/10/29 Python
Django 解决开发自定义抛出异常的问题
2020/05/21 Python
python中pow函数用法及功能说明
2020/12/04 Python
使用CSS3制作饼状旋转载入效果的实例
2015/06/23 HTML / CSS
详解HTML5中的picture元素响应式处理图片
2018/01/03 HTML / CSS
微软香港官网及网上商店:Microsoft HK
2016/09/01 全球购物
HomeAway的巴西品牌:Alugue Temporada
2018/04/10 全球购物
寒假思想汇报
2014/01/10 职场文书
《唯一的听众》教学反思
2014/02/20 职场文书
应届毕业生自荐信例文
2014/02/26 职场文书
大专生自荐书范文
2014/06/22 职场文书
工作失误检讨书范文
2015/01/26 职场文书
酒店开业主持词
2015/07/02 职场文书
《小乌鸦爱妈妈》教学反思
2016/02/19 职场文书
六年级作文之关于梦
2019/10/22 职场文书
vue的项目如何打包上线
2022/04/13 Vue.js