学习Pytorch+Python之自定义数据集(创建CSV文件)

1、CSV文件
逗号分隔值（Comma-Separated Values，CSV，有时也称为字符分隔值，因为分隔字符也可以不是逗号），其文件以纯文本形式存储表格数据（数字和文本）。文件格式如下图，也可用Excel打开。

2、简单分类数据集
本文数据集是一个蚂蚁(ants)和蜜蜂(bees)的简单的二分类数据，分别存放于dataset文件夹下的ants_image和bees_image文件夹下，用文件夹名字作为标签。
3、创建CSV文件代码

# 导入相关库
import csv
import random
import os
import glob
# 数据地址，注意是分类文件夹的上层文件夹
# dataset文件夹中的两个子文件夹分别对应两个类名
root = 'dataset'
# 对数据地址中的文件夹进行遍历，将类名存放于列表names中
names = os.listdir(root)
# 创建名为images的空列表用于存放图像地址
images = []
# 创建名称、标签字典，用于存放二分类的类名和标签
names_labels = {}

# 遍历类名
for name in names:
	# 遍历第一个类名的时候，这个时候仅有一个键，为ants_image
	# 遍历第二个类名的时候，这个时候有两个键，分别为ants_image和bees_image
	# 那么字典names_label的键的长度在第一次遍历的时候为0，第二次遍历的时候为1
    names_labels[name] = len(names_labels.keys())
    # glob.glob返回所有匹配的文件路径列表。
	# 相对于本文而言，返回的是路径dataset/ants_image/xx...xxx.jpg的文件
	# 用*号匹配复杂没有规律的所有格式为jpg的图片名，且*表示匹配0个或多个字符
    images += glob.glob(os.path.join(root, name, '*.jpg'))

print(names_labels)
print(images)

# 对csv文件进行写操作，如果没有csv文件会自动创建
# (1)图中，逗号前为文件名，逗号后为标签
# newline=''的作用是防止每一行数据后面都自动增加了一个空行
with open(os.path.join(root, 'img.csv'), mode='w', newline='') as f:
    writer = csv.writer(f)
    for img in images:
    	# 用os.sep切割具有通用性，自动识别路径分隔符windows和linus
        name = img.split(os.sep)[-2]
        label = names_labels[name]
        writer.writerow([name, label])

4、读取CSV文件

# 创建images和labels空列表用于存放数据和标签
images, labels = [], []
with open(os.path.join(root, filename)) as f:
    reader = csv.reader(f)
    for row in reader:
        img, label = row
        label = int(label)
        images.append(img)
        labels.append(label)