Tensorflow2.0之自定义数据集
Brief
在学习Deep Learning的过程中我们难免会因为需求而要使用自定义的DS，本文就简要的介绍如何自定义DS
PoKemon Dataset
这次我们使用的是PoKemon Dataset ，在自定义数据集之后会训练一下参数效果
数据集使用的是龙龙老师课上的数据集下面是百度网盘的链接：
▪
链接 : https://pan.baidu.com/s/1V_ZJ7ufjUUFZwD2NHSNMFw
▪
提取码： dsxl
Splitting
steps
Load data in csv
首先让我们来处理图片，下面是我们的数据集文件夹，每一个文件夹中对应了相应的PoKemo
现在我们想要把他们的数据集写在csv文件中，效果大约如下图：
先让我们捋一下思路，首先我们应该读取有pokemon图片的文件夹，然后获取标签，然后写入csv文件
我们用字典来保存我们的图片地址和标签
def load_pokemon(root, model='train'):
"""
load pokemon dataset info
:param root: the root path of dataset
:param model: train val or test
:return: image,labels,name2able
"""
name2label = {}
for name in sorted(os.listdir(os.path.join(root))):
if not os.path.exists(os.path.join(root, name)):
continue
# code each category (use the length of the key)
name2label[name] = len(name2label.keys())
...
ok 现在我们获得了有宝可梦和标签的字典，下面将他们写入到csv文件中
def load_csv(root, filename, name2label):
# this will create a csv file when I first run it.
if not os.path.exists(os.path.join(root, filename)):
images = []
for name in name2label.keys():
# 'pokemon\\name\\00001.png
images += glob.glob(os.path.join(root, name, '*.png'))
images += glob.glob(os.path.join(root, name, '*.jpg'))
images += glob.glob(os.path.join(root, name, '*.jpeg'))
print(len(images), images)
random.shuffle(images)
with open(os.path.join(root, filename), mode='w', newline='') as f:
writer = csv.writer(f)
for img in images:
name = img.split(os.sep)[-2]
lable = name2label[name]
writer.writerow([img, lable])
print('written into csv file', filename)
这样我们就得到了有宝可梦图片地址和标签的地址
现在我们在load_csv 中读取csv文件
images, labels = [], []
# read from csv file
with open(os.path.join(root, filename), 'r') as f:
reader = csv.reader(f)
for row in reader:
img, label = row
label = int(label)
images.append(img)
labels.append(label)
assert len(images) == len(labels)
return images, labels
之后我们在load_pokemon中调用load_csv函数
images, labels = load_csv(root, 'image.csv', name2label)
if model == 'train': # 60%
images = images[:int(0.6 * len(images))]
labels = labels[:int(0.6 * len(labels))]
elif model == 'val': # 20%
images = images[:int(0.2 * len(images))]
labels = labels[:int(0.2 * len(labels))]
else: # 20%
images = images[:int(0.2 * len(images))]
labels = labels[:int(0.2 * len(labels))]
return images, labels
Preprocess the data
现在我们得到的只是图片的地址，并不能直接用来训练，所以我们需要进行预处理
Read and Resize
首先我们先读取图片信息并进行裁剪
定义一个preprocess函数
def preprocess(x, y):
"""
preprocess the data
:param x: the path of the images
:param y: labels
"""
# data augmentation, 0~255
x = tf.io.read_file(x)
x = tf.image.decode_jpeg(x, channels=3)
# resize the image,you can change the value in the another net
x = tf.image.resize(x, [224, 224])
# turn around images
x = tf.image.random_crop(x, [224, 224, 3])
Data Augmentation
# # x: [0,255]=> 0~1
x = tf.cast(x, dtype=tf.float32) / 255
Normalize
这里的img_mean和img_std 都是imgNet 中几百万张数据集中得到的值，我们直接拿过来用就好
img_mean = tf.constant([0.485, 0.456, 0.406])
img_std = tf.constant([0.229, 0.224, 0.225])
def normalize(x, mean=img_mean, std=img_std):
# x: [224, 224, 3]
# mean: [224, 224, 3], std: [3]
x = (x - mean) / std
return x
好的现在我们可以在preprocess后面加上：
# 0~1 => D(0,1)
x = normalize(x)
y = tf.convert_to_tensor(y)
return x,y
现在我们就可以加载了
if __name__ == '__main__':
image_train, lab_train = load_pokemon('pokemon', model='train')
image_val, lab_val = load_pokemon('pokemon', model='val')
image_test, lab_test = load_pokemon('pokemon', model='test')
train_db = tf.data.Dataset.from_tensor_slices((image_train, lab_train))
train_db = train_db.shuffle(1000).map(preprocess).batch(32)
val_db = tf.data.Dataset.from_tensor_slices((image_val, lab_val))
val_db = val_db.map(preprocess).batch(32)
test_db = tf.data.Dataset.from_tensor_slices((image_test, lab_test))
test_db = test_db.map(preprocess).batch(32)
print(train_db)
print(val_db)
print(test_db)
完整的代码可以去我的github上下载：
https://github.com/Jupiter-king/tf2.0-Custom-data-set
转载：https://blog.csdn.net/python_LC_nohtyp/article/details/104685251
查看评论