额,这里我们在网上找了10类花朵的数据,将数据进行分类,放在各个文件夹,文件名是花朵的标签,然后对图片大小统一为256*256。
将数据集分成训练集(train)、验证集(validation)、测试集(test)
分别为训练集800张,验证集100张,测试集100张,训练集和验证集的需要进行灰度处理,测试集不需要。
1.准备数据集好后,将文件路径和标签保存在txt文件中
from torchvision import transforms
import os
path = './validation/'
a = []
data_tf = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])])
#图片名和label信息的文本
def photo_text(path, name):
table = os.listdir(path)
for i in table:
if i =='train.txt' or i == 'test.txt' or i == 'validation.txt':
continue
else:
a.append(path+i +' '+i.split('_')[0])
try:
with open(path + name,'x', encoding='utf-8') as f:
for j in a:
f.write(j+'\n')
except:
os.remove(path+name)
with open(path + name,'x', encoding='utf-8') as f:
for j in a:
f.write(j+'\n')
if __name__ =='__main__':
photo_text(path,'validation.txt')
2.创建数据集,可以得到
from PIL import Image
from torchvision import transforms
import torch
# 创建数据集
class MyDataset(torch.utils.data.Dataset): # 创建自己的类:MyDataset,这个类是继承的torch.utils.data.Dataset
def __init__(self, root, datatxt, transform=None, target_transform=None): # 初始化一些需要传入的参数
super(MyDataset, self).__init__()
fh = open(root + datatxt, 'r', encoding='utf-8') # 按照传入的路径和txt文本参数,打开这个文本,并读取内容
imgs = [] # 创建一个名为img的空列表,一会儿用来装东西
for line in fh: # 按行循环txt文本中的内容
line = line.rstrip() # 删除 本行string 字符串末尾的指定字符,这个方法的详细介绍自己查询python
words = line.split() # 通过指定分隔符对字符串进行切片,默认为所有的空字符,包括空格、换行、制表符等
imgs.append((words[0], int(words[1]))) # 把txt里的内容读入imgs列表保存,具体是words几要看txt内容而定
# 很显然,刚才截图所示txt的内容,words[0]是图片信息,words[1]是lable
self.imgs = imgs
self.transform = transform
self.target_transform = target_transform
def __getitem__(self, index):
# 这个方法是必须要有的,用于按照索引读取每个元素的具体内容
fn, label = self.imgs[index] # fn是图片path #fn和label分别获得imgs[index]也即是刚才每行中word[0]和word[1]的信息
img = Image.open(fn).convert('RGB') # 按照path读入图片from PIL import Image # 按照路径读取图片
if self.transform is not None:
img = self.transform(img) # 是否进行transform
return img, label # return很关键,return回哪些内容,那么我们在训练时循环读取每个batch时,就能获得哪些内容
def __len__(self): # 这个函数也必须要写,它返回的是数据集的长度,也就是多少张图片,要和loader的长度作区分
return len(self.imgs)
path_train = './train/'
path_validation = './validation/'
# 数据预处理。transforms.ToTensor()将图片转换成PyTorch中处理的对象Tensor,并且进行标准化(数据在0~1之间)
# transforms.Normalize()做归一化。它进行了减均值,再除以标准差。两个参数分别是均值和标准差
# transforms.Compose()函数则是将各种预处理的操作组合到了一起
data_tf = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])])
#个人数据装载
train_data = MyDataset(path_train ,'train.txt', transform=data_tf)
test_data = MyDataset(path_validation ,'validation.txt', transform=data_tf)
#数据装载
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, pin_memory = True)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=True)