使用PyTorch将文件夹下的图片分为训练集和验证集实例
更新时间:2020年01月08日 14:02:50 作者:xgbm_k
今天小编就为大家分享一篇使用PyTorch将文件夹下的图片分为训练集和验证集实例,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧
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下找到。
相关文章
python pyinstaller打包exe报错的解决方法
这篇文章主要给大家介绍了关于python pyinstaller打包exe报错的解决方法,文中通过示例代码介绍的非常详细,对大家的学习或者使用python具有一定的参考学习价值,需要的朋友们下面来一起学习学习吧2019-11-11python 爬取豆瓣电影短评并利用wordcloud生成词云图
这篇文章主要介绍了python 爬取豆瓣电影短评并利用wordcloud生成词云图,文章围绕主题展开详细的内容介绍,具有一定的参考价值,需要的小伙伴可以参考一下2022-06-06Python 过滤字符串的技巧,map与itertools.imap
Python中的map函数非常有用,在字符转换和字符遍历两节都出现过,现在,它又出现了,会给我们带来什么样的惊喜呢?是不是要告诉我们,map是非常棒的,以后要多找它玩呢?2008-09-09解决Keyerror ''''acc'''' KeyError: ''''val_acc''''问题
这篇文章主要介绍了解决Keyerror 'acc' KeyError: 'val_acc'问题,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧2020-06-06
最新评论