使用PyTorch将文件夹下的图片分为训练集和验证集实例

yipeiwu_com5年前Python基础

PyTorch提供了ImageFolder的类来加载文件结构如下的图片数据集:

root/dog/xxx.png
root/dog/xxy.png
root/dog/xxz.png

root/cat/123.png
root/cat/nsdf3.png
root/cat/asd932_.png

使用这个类的问题在于无法将训练集(training dataset)和验证集(validation dataset)分开。我写了两个类来完成这个工作。

import os
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import ToTensor, Resize, Compose
from PIL import Image
from sklearn.model_selection import train_test_split

class ImageFolderSplitter:
  # images should be placed in folders like:
  # --root
  # ----root\dogs
  # ----root\dogs\image1.png
  # ----root\dogs\image2.png
  # ----root\cats
  # ----root\cats\image1.png
  # ----root\cats\image2.png  
  # path: the root of the image folder
  def __init__(self, path, train_size = 0.8):
    self.path = path
    self.train_size = train_size
    self.class2num = {}
    self.num2class = {}
    self.class_nums = {}
    self.data_x_path = []
    self.data_y_label = []
    self.x_train = []
    self.x_valid = []
    self.y_train = []
    self.y_valid = []
    for root, dirs, files in os.walk(path):
      if len(files) == 0 and len(dirs) > 1:
        for i, dir1 in enumerate(dirs):
          self.num2class[i] = dir1
          self.class2num[dir1] = i
      elif len(files) > 1 and len(dirs) == 0:
        category = ""
        for key in self.class2num.keys():
          if key in root:
            category = key
            break
        label = self.class2num[category]
        self.class_nums[label] = 0
        for file1 in files:
          self.data_x_path.append(os.path.join(root, file1))
          self.data_y_label.append(label)
          self.class_nums[label] += 1
      else:
        raise RuntimeError("please check the folder structure!")
    self.x_train, self.x_valid, self.y_train, self.y_valid = train_test_split(self.data_x_path, self.data_y_label, shuffle = True, train_size = self.train_size)

  def getTrainingDataset(self):
    return self.x_train, self.y_train

  def getValidationDataset(self):
    return self.x_valid, self.y_valid

class DatasetFromFilename(Dataset):
  # x: a list of image file full path
  # y: a list of image categories
  def __init__(self, x, y, transforms = None):
    super(DatasetFromFilename, self).__init__()
    self.x = x
    self.y = y
    if transforms == None:
      self.transforms = ToTensor()
    else:
      self.transforms = transforms
    
  def __len__(self):
    return len(self.x)

  def __getitem__(self, idx):
    img = Image.open(self.x[idx])
    img = img.convert("RGB")
    return self.transforms(img), torch.tensor([[self.y[idx]]])

# test code
# splitter = ImageFolderSplitter("for_test")
# transforms = Compose([Resize((51, 51)), ToTensor()])
# x_train, y_train = splitter.getTrainingDataset()
# training_dataset = DatasetFromFilename(x_train, y_train, transforms=transforms)
# training_dataloader = DataLoader(training_dataset, batch_size=2, shuffle=True)
# x_valid, y_valid = splitter.getValidationDataset()
# validation_dataset = DatasetFromFilename(x_valid, y_valid, transforms=transforms)
# validation_dataloader = DataLoader(validation_dataset, batch_size=2, shuffle=True)
# for x, y in training_dataloader:
#   print(x.shape, y.shape)

更多的代码可以在我的Github reop下找到。

相关文章

如何处理Python3.4 使用pymssql 乱码问题

在项目中发现这样一个问题:sqlserver数据库编码为gbk,使用python3.4+pymssql 查询,中文乱码,经过一番思考问题解决,下面把解决办法分享给大家: conn...

python判断windows系统是32位还是64位的方法

本文实例讲述了python判断windows系统是32位还是64位的方法。分享给大家供大家参考。具体分析如下: 通常64的windows系统program files文件夹(用来安装应用...

python 时间信息“2018-02-04 18:23:35“ 解析成字典形式的结果代码详解

python 时间信息“2018-02-04 18:23:35“ 解析成字典形式的结果代码详解

将时间信息“2018-02-04  18:23:35“ 解析成字典形式的结果 如:{‘year':2018,‘month':2,‘day':4,‘hour':18:‘minut...

关于python的bottle框架跨域请求报错问题的处理方法

关于python的bottle框架跨域请求报错问题的处理方法

在用python的bottle框架开发时,前端使用ajax跨域访问时,js代码老是进入不了success,而是进入了error,而返回的状态却是200。url直接在浏览器访问也是正常的,...

Python基础学习之函数方法实例详解

本文实例讲述了Python基础学习之函数方法。分享给大家供大家参考,具体如下: 前言 与其他编程语言一样,函数(或者方法)是组织好的,可重复使用的,用来实现单一,或相关联功能的代码段。...