分类任务数据集制作

众所周知，在深度学习图像分类、目标检测等任务的训练过程中需要大量的训练图片；为了方便快速批量的读取数据，往往都会先把图片数据集转化成特殊的文件格式，常见的如：.mat, .npy, .pkl, .csv 等。本文主要介绍如何由现有图片生成.pkl文件。

图片的整理

首先，将图片按类别整理成如下目录格式： 
- 第一层目录为/jpg 
- 第二层目录如图 
 
显然每个文件夹下放的是各类的图片。

转换程序

# img2pkl.py
from __future__ import division, print_function, absolute_import
import os
import sys
# import tarfile
import numpy as np
import pickle
import random
from PIL import Image

# 主要函数接口
def load_data(dirname="17flowers", pklpath='17flowers/jpg/',
                resize_pics=(224, 224), shuffle=True, one_hot=False):

    dataset_file = os.path.join(dirname, 'dataset.pkl')
    X, Y = build_image_dataset_from_dir(directory=pklpath,
                                        dataset_file=dataset_file,
                                        resize=resize_pics,
                                        filetypes=['.jpg', '.jpeg'],
                                        convert_gray=False,
                                        shuffle_data=shuffle,
                                        categorical_Y=one_hot)

    return X, Y

####
def build_image_dataset_from_dir(directory,
                                 dataset_file="dataset.pkl",
                                 resize=None, convert_gray=None,
                                 filetypes=None, shuffle_data=False,
                                 categorical_Y=False):
    try:
        X, Y = pickle.load(open(dataset_file, 'rb'))
    except Exception:
        X, Y = image_dirs_to_samples(directory, resize, convert_gray, filetypes)
        if categorical_Y:
            Y = to_categorical(Y, np.max(Y) + 1) # First class is '0'
        if shuffle_data:
            X, Y = shuffle(X, Y)
        pickle.dump((X, Y), open(dataset_file, 'wb'))
    return X, Y

########
def image_dirs_to_samples(directory, resize=None, convert_gray=None,
                          filetypes=None):
    print("Starting to parse images...")
    if filetypes:
        if filetypes not in [list, tuple]: filetypes = list(filetypes)
    samples, targets = directory_to_samples(directory, flags=filetypes)
    for i, s in enumerate(samples):
        samples[i] = load_image(s)
        if resize:
            samples[i] = resize_image(samples[i], resize[0], resize[1])
        if convert_gray:
            samples[i] = convert_color(samples[i], 'L')
        samples[i] = pil_to_nparray(samples[i])
        samples[i] /= 255.
    print("Parsing Done!")
    return samples, targets

# 类别one-hot编码
def to_categorical(y, nb_classes):
    """ to_categorical.

    Convert class vector (integers from 0 to nb_classes)
    to binary class matrix, for use with categorical_crossentropy.

    Arguments:
        y: `array`. Class vector to convert.
        nb_classes: `int`. Total number of classes.

    """
    y = np.asarray(y, dtype='int32')
    if not nb_classes:
        nb_classes = np.max(y)+1
    Y = np.zeros((len(y), nb_classes))
    for i in range(len(y)):
        Y[i, y[i]] = 1.
    return Y

# 将数据打乱
def shuffle(*arrs):
    """ shuffle.

    Shuffle given arrays at unison, along first axis.

    Arguments:
        *arrs: Each array to shuffle at unison as a parameter.

    Returns:
        Tuple of shuffled arrays.

    """
    arrs = list(arrs)
    for i, arr in enumerate(arrs):
        assert len(arrs[0]) == len(arrs[i])
        arrs[i] = np.array(arr)
    p = np.random.permutation(len(arrs[0]))
    return tuple(arr[p] for arr in arrs)

# 遍历各文件夹和文件夹中的图片
def directory_to_samples(directory, flags=None):
    """ Read a directory, and list all subdirectories files as class sample """
    samples = []
    targets = []
    label = 0
    classes = sorted(os.walk(directory).next()[1])
    for c in classes:
        c_dir = os.path.join(directory, c)
        for sample in os.walk(c_dir).next()[2]:
            if not flags or any(flag in sample for flag in flags):
                    samples.append(os.path.join(c_dir, sample))
                    targets.append(label)
        label += 1
    return samples, targets

# 使用PIL库读取图片
def load_image(in_image):
    """ Load an image, returns PIL.Image. """
    img = Image.open(in_image)
    return img

# 图片尺度变换函数
def resize_image(in_image, new_width, new_height, out_image=None,
                 resize_mode=Image.ANTIALIAS):
    """ Resize an image.

    Arguments:
        in_image: `PIL.Image`. The image to resize.
        new_width: `int`. The image new width.
        new_height: `int`. The image new height.
        out_image: `str`. If specified, save the image to the given path.
        resize_mode: `PIL.Image.mode`. The resizing mode.

    Returns:
        `PIL.Image`. The resize image.

    """
    img = in_image.resize((new_width, new_height), resize_mode)
    if out_image:
        img.save(out_image)
    return img

# 
def convert_color(in_image, mode):
    """ Convert image color with provided `mode`. """
    return in_image.convert(mode)

# 加载图片并转换为numpy多维数组
def pil_to_nparray(pil_image):
    """ Convert a PIL.Image to numpy array. """
    pil_image.load()
    return np.asarray(pil_image, dtype="float32")

if __name__ == "__main__":
    load_data(dirname="dataSet", resize_pics=(224, 224), True, True)
 
 
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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
 
 
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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

使用

如下说明了如何使用这个数据转换模块，如果是初次调用将会先生成.pkl文件。

import image2pkl
X, Y = image2pkl.load_data(one_hot=True, resize_pics=(224,224))
# 得到的X是训练样本，Y是对应的标签；用法和MNIST手写体识别的数据集一样
 
 
1
2
3
 
 

  
  
 
 
 
 
1
2
3

也可以直接使用生成的.pkl文件，如下：

'''using .pkl file directly'''
import cPickle as pickle

file = open("dataset.pkl")
data, labels = pickle.load(file)
print data.shape
print labels.shape
# alpha控制训练和测试数据的比例
alpha = 0.3
trainIndex = data.shape[0] - int(data.shape[0] * alpha)
testIndex = int(data.shape[0] * alpha)
trainData = data[:trainIndex]
testData  = data[-testIndex:]
trainLabels = labels[:trainIndex]
testLabels = labels[-testIndex:]