小董的BLOG
0%

制作数据集

发表于 更新于

最近看了几个算法的数据集建立方法，在此记录一下，方便以后借用

对于数量较少且每种图片数量相同的情况

若要使用该方法，图片在文件夹中需按顺序排列

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import glob
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from torchvision import transforms
from sklearn.model_selection import train_test_split


class Mydataset(Dataset):
    def __init__(self, images, labels, transform):
        self.images = images
        self.labels = labels
        self.transform = transform
        dataset = []
        for i in range(len(labels)):	#每个图片都有对应的标签元素，因此标签数量=图片数量
            temp_img = Image.open(images[i])
            temp_img = self.transform(temp_img)  # transforms.Compose对象可直接输入图像
            dataset.append((temp_img, labels[i]))#经过train_test_split的图片标签一一对应
        self.dataset = dataset

    def __getitem__(self, index):
        return self.dataset[index]

    def __len__(self):
        return len(self.labels)

#先看这个
def load_dataset(self):
    data = []
    all_imgs_path = glob.glob(r'dataset\*.png')	#查找文件夹中的所有png文件
    for ip in all_imgs_path:
        data.append(ip)		#此处是保存了每个图像路径的列表
    labels = []
    for i in range(20): #只适用于知道每个类别图片数量的情况，且数量最好相等且不多
        labels.extend([i] * 72) #图片有20种，每种有72张
      #分别生成72个0、1...20，此处i即为标签
    tr_imgs, te_imgs, tr_labs, te_labs = train_test_split(data, labels, train_size=0.9)	#因为data按照文件顺序排列，因此图片与标签可以一一对应，训练集:验证集=9:1
   
    transform = transforms.Compose([
        transforms.Resize((128, 128)),
        transforms.ToTensor()
    ])
      
    tr_set = Mydataset(tr_imgs, tr_labs, transform)
    te_set = Mydataset(te_imgs, te_labs, transform)
    tr_loader = DataLoader(tr_set, batch_size=self.args.train_batch_size, shuffle=True, num_workers=self.args.workers,
                           pin_memory=True)
    te_loader = DataLoader(te_set, batch_size=self.args.train_batch_size, shuffle=True, num_workers=self.args.workers,
                           pin_memory=True)
    return tr_loader, te_loader

对于图片已按照文件夹分好的情况

数据集划分

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
import os
from shutil import copy
import random


def mkfile(file):
    if not os.path.exists(file):
        os.makedirs(file)


# 获取 flower_photos 文件夹下除 .txt 文件以外所有文件夹名（即5种花的类名）
file_path = 'flower_data/flower_photos'
flower_class = [cla for cla in os.listdir(file_path) if ".txt" not in cla]
#os.listdir会返回文件夹下的文件夹名的列表

# 创建 训练集train 文件夹，并由5种类名在其目录下创建5个子目录
mkfile('flower_data/train')
for cla in flower_class:
    mkfile('flower_data/train/' + cla)

# 创建 验证集val 文件夹，并由5种类名在其目录下创建5个子目录
mkfile('flower_data/val')
for cla in flower_class:
    mkfile('flower_data/val/' + cla)

# 划分比例，训练集 : 验证集 = 9 : 1
split_rate = 0.1

# 遍历5种花的全部图像并按比例分成训练集和验证集
for cla in flower_class:
    cla_path = file_path + '/' + cla + '/'  # 某一类别花的子目录
    images = os.listdir(cla_path)  # iamges 列表存储了该目录下所有图像的名称
    num = len(images)
    eval_index = random.sample(images, k=int(num * split_rate))  # 从images列表中随机抽取 k 个图像名称
    for index, image in enumerate(images):  #创建验证集
        # eval_index 中保存验证集val的图像名称
        if image in eval_index:
            image_path = cla_path + image
            new_path = 'flower_data/val/' + cla
            copy(image_path, new_path)  # 将选中的图像复制到新路径

        # 其余的图像保存在训练集train中
        else:
            image_path = cla_path + image
            new_path = 'flower_data/train/' + cla
            copy(image_path, new_path)
        print("\r[{}] processing [{}/{}]".format(cla, index + 1, num), end="")  # processing bar
    print()

print("processing done!")

pytorch加载数据集

1
2
3
4
5
6
7
8
# 获取图像数据集的路径
data_root = os.path.abspath(os.path.join(os.getcwd(), "../.."))  # get data root path 返回上上层目录
image_path = data_root + "/data_set/flower_data/"  			# flower data_set path

# 导入训练集并进行预处理
train_dataset = datasets.ImageFolder(root=image_path + "/train",		
                                     transform=data_transform["train"])

ImageFolder()返回的对象是一个包含数据集所有图像及对应标签构成的二维元组容器，支持索引和迭代，可作为torch.utils.data.DataLoader的输入，其有三种方法：

  • self.classes：用一个列表保存类别名称，例如[‘dog’,’cat’]

  • self.class_to_idx：得到字典：类别 + 对应的索引值，例如{‘dog’:1,’cat’:2}

    1
    2
    3
    flower_list = train_dataset.class_to_idx
    # 将 flower_list 中的 key 和 val 调换位置
    cla_dict = dict((val, key) for key, val in flower_list.items())

  • self.imgs：保存(图像路径, 它所属于的类别index) tuple的 list

同时，imageFolder()过程只会加载index，而不会执行transform操作（transform操作是懒加载的，只有使用dataLoader的时候才会执行）

对同一数据集进行分割

1
2
3
4
5
6
7
# 从训练集的50000个样本中，取49000个作为训练集，剩余1000个作为验证集
NUM_TRAIN = 49000

loader_train = DataLoader(dataset, batch_size=64, sampler=sampler.SubsetRandomSampler(range(NUM_TRAIN)))

loader_val = DataLoader(dataset, batch_size=64, sampler=sampler.SubsetRandomSampler(range(NUM_TRAIN, 50000)))