看了一下网上的划分图像数据集大多数都是存在一个文件里,按比例把这些图片分开。但是遇到多分类问题怎么办?划分后能否直接与pytorch的ImageFolder对接?这里我给出来我自己划分分类任务图像数据集的code,通过我的代码划分完数据集后,就可以直接与ImageFolder来对接训练了。
话不多说直接上代码:
def partition_dataset(label_csv_path, root_path=None, save_path=None, shuffle=False, proportion=0.20):
"""
根据标签将图片分类,保存到不同文件夹下,文件夹名称为标签名
:param label_csv_path: 标签存放路径
:param root_path: 指定标签中图片的根目录
:param save_path: 保存图片的路径
:param proportion: 测试集所占比例
:param shuffle: 是否shuffle
:return: None
"""
import random
label_csv_list = [] # 将标签文件转化为嵌套列表 [['1', '100.PNG'], ... ,['5', '1000.PNG']]
file_name_list = []
label_file_dict = {} # 嵌套列表的字典,key:标签 value:标签对应的所有样本
label_count_dict = {} # 每个标签对应的样本数
label_csv_file = open(label_csv_path)
csv_reader_line = csv.reader(label_csv_file)
for line in csv_reader_line:
label_csv_list.append(line[::-1])
file_name_list.append(line[::-1][0])
file_name_list = set(file_name_list) # 去重,得到标签,可以作为文件名和字典的key值 {1,2,3,4,5,6}
for item in file_name_list:
label_file_dict.setdefault(item, [])
label_count_dict.setdefault(item, 1)
for item in label_csv_list:
label_file_dict[item[0]].append(item[1])
for item in file_name_list:
label_count_dict[item] = len(label_file_dict[item])
print(label_count_dict)
# 创建数据集的目录结构
train_path = os.path.join(save_path, "train")
val_path = os.path.join(save_path, "val")
if not os.path.exists(train_path): # 如果路径不存在
os.makedirs(train_path)
if not os.path.exists(val_path):
os.makedirs(val_path)
# 在train和val目录下分别创建与标签对应的文件夹
for item in file_name_list:
if not os.path.exists(os.path.join(train_path, item)):
os.makedirs(os.path.join(train_path, item))
if not os.path.exists(os.path.join(val_path, item)):
os.makedirs(os.path.join(val_path, item))
# 是否shuffle
if shuffle is True:
for k in label_file_dict:
random.shuffle(label_file_dict[k])
# 开始划分
for k in label_file_dict:
for i in range(0, int((len(label_file_dict[k])+1)*(1.0-proportion))):
shutil.copy(os.path.join(root_path, label_file_dict[k][i]),
os.path.join(train_path + "\\" + k + "\\", label_file_dict[k][i]))
print("\033[31m copy: %s to %s \033[0m" % (
label_file_dict[k][i], os.path.join(train_path + "\\" + k + "\\", label_file_dict[k][i])))
for i in range(int((len(label_file_dict[k])+1)*(1.0-proportion)), len(label_file_dict[k])):
shutil.copy(os.path.join(root_path, label_file_dict[k][i]),
os.path.join(val_path + "\\" + k + "\\", label_file_dict[k][i]))
print("\033[33m copy: %s to %s \033[0m" % (
label_file_dict[k][i], os.path.join(val_path + "\\" + k + "\\", label_file_dict[k][i])))
该方法的各种参数以及具体实现过程已经注释的很清楚了,这里我再把我的标签放上来,如果你的标签和我的格式一样,那么应该只需要修改少量的代码就可以成功运行了:
标签就长这个样子,是一个csv,没有head,左侧是图片文件名,右侧是对应的标签。
现在调用 partition_dataset 方法,进行数据集划分:
partition_dataset(r"E:\complabel.csv",
r"E:\dataset\image"+"\\",
r"E:\comp_dataset",
shuffle=True,
proportion=0.3)
划分完成后,我们看一下划分的结果:
comp_dataset文件夹中的内容:
train文件夹中的内容:
val文件夹中的内容:
可以看到数据集已经按标签划分为训练集和验证集了。
那我们验证一下第一类的图像是否全部被划分,并且是按照我们给定的百分比30%划分的:
首先打开train/1 文件夹,该文件夹下存放了所有标签为1的训练集图片:
trian/1 文件夹下一共74个sample
val/1文件夹下一共31个sample:
我的标签中的第一类一共有105个sample
74+31=105 74/105≈0.704,说明划分没有问题,那么我们可以直接使用pytorch的imageFolder读取数据集进行训练了:
train_dataset = datasets.ImageFolder(root=r"E:\dataset\comp_dataset" + "\\train",
transform=data_transform["train"])
train_num = len(train_dataset)
validate_dataset = datasets.ImageFolder(root=r"E:\dataset\comp_dataset" + "\\val",
transform=data_transform["val"])
使用ImageFolde的好处是可以直观的看到自己数据集样本分布的情况,尤其是对一个分类任务而言,样本的均衡性是至关重要的,当然另个一好处就是调用很方便,不需要去继承Dataset类,重写读数据方法了。
最后说明一下,每个人的标签或者数据集图片格式不一样,需要自己对代码进行微调。