代码收藏家 深度学习制作自己的数据集—为数据集打上标签保存为txt文件，并进行划分和加载数据集
目录
0 前言
目前是被封控的第四天了，只能呆在宿舍不能出去，记得上次这样子还是一年前大四快毕业那时候了……
这几天在宿舍没有什么事干，实验也暂时做不了了，将部分数据处理完后，就把之前的这个内容做一下笔记吧，这也不是什么新的知识了，简单记录一下，方便以后可以查看。
1 为图片数据集打上标签并保存为txt文件
由于这里我做的是用深度学习回归预测，所以我的标签保存在（.csv）文件中，这时候需要将图片和标签一一对应起来，并且要分好文件夹，下面是我分好的文件夹（images保存的是图片，label.csv保存的是对应的标签，这里可以根据个人的数据集更改文件名称）：
下面是为图片数据集打上标签并保存为txt文件的代码（文件的路劲需要根据自己文件所在位置进行更改）：
import os
import numpy as np
import pandas as pd
label = pd.read_csv('../dataset_1/label_1.csv')
label = np.array(label)
label = label.tolist()
target = ''
# for i in range(len(label)):
# for j in range(len(label[i])):
# target += str(label[i][j]) + ' '
# print(target)
# target = ''
def generate(dir):
files = os.listdir(dir) #os.listdir() 方法用于返回指定的文件夹包含的文件或文件夹的名字的列表。
# files.sort() #对文件或文件夹进行排序
files.sort(key=lambda x: int(x.replace("frame", "").split('.')[0]))
print('****************')
print('input :', dir)
print('start...')
target = ''
i = 0
listText = open('H:/代码练习/Deeplearning/data_txt_path/all_data_list_1.txt', 'a+') #创建并打开一个txt文件，a+表示打开一个文件并追加内容
listText.truncate(0)#清空txt文件里的内容
for file in files: #遍历文件夹中的文件
fileType = os.path.split(file) #os.path.split（）返回文件的路径和文件名，【0】为路径，【1】为文件名
if fileType[1] == '.txt': #若文件名的后缀为txt,则继续遍历循环，否则退出循环
continue
name = outer_path + folder + '/' +file #name 为文件路径和文件名+空格+label+换行
for j in range(len(label[i])):
target += str(label[i][j]) + ' '
name = name + ' ' + target + '\n'
# print(name)
# listText.write(name) # 在创建的txt文件中写入name
target = ''
i += 1
listText.write(name) #在创建的txt文件中写入name
listText.close() #关闭txt文件
print('down!')
print('****************')
outer_path = 'H:/代码练习/Deeplearning/dataset_1/' # 这里是你的图片路径
if __name__ == '__main__': #主函数
folderlist = os.listdir(outer_path)# 列举文件夹
for folder in folderlist: #遍历文件夹中的文件夹(若engagement文件夹中存在txt或py文件，则后面会报错）
generate(os.path.join(outer_path, folder))#调用generate函数，函数中的参数为：（图片路径+文件夹名，标签号）
代码运行后结果如下图：序号1为图片的路径，序号2为对应的标签，因为我一张图片对应3个标签，所以有后面3个值。
2 将txt文件中的图片标签数据集随机划分为训练集和测试集
完成第一个步骤后，需要将txt文件中的图片标签数据集随机划分为训练集和测试集，划分后生成训练集和测试集两个txt文件，可以根据自己的需求，更改训练集和测试集的比例。下面为划分数据集的代码（文件的路劲需要根据自己文件所在位置进行更改）：
import os
import random
# 划分比例，训练集 : 验证集 = 8 : 2
split_rate = 0.2
class SplitFiles():
"""按行分割文件"""
def __init__(self, file_name):
"""初始化要分割的源文件名和分割后的文件行数"""
self.file_name = file_name
# def get_random(self):
# """生成随机数组，随机划分 （0，190001）txt标签行数， 7600测试集标签行数"""
# random_num = random.sample(range(0, 19001), 108)
#
# return random_num
def split_file(self):
if self.file_name and os.path.exists(self.file_name):
try:
with open(self.file_name) as f: # 使用with读文件
# temp_count = 1
file = f.readlines()
count = len(file)
eval_index = random.sample(file, k=int(count * split_rate)) # 从images列表中随机抽取 k 个图像名称
for index,image_path in enumerate(file):
if image_path in eval_index:
self.write_file('test', image_path)
else:
self.write_file('train', image_path)
# temp_count += 1
except IOError as err:
print(err)
else:
print("%s is not a validate file" % self.file_name)
def get_part_file_name(self, part_name):
""""获取分割后的文件名称：在源文件相同目录下建立临时文件夹temp_part_file，然后将分割后的文件放到该路径下"""
temp_path = os.path.dirname(self.file_name) # 获取文件的路径（不含文件名）
file_folder = temp_path
if not os.path.exists(file_folder): # 如果临时目录不存在则创建
os.makedirs(file_folder)
part_file_name = file_folder + "/" + str(part_name) + "_list_1.txt"
return part_file_name
def write_file(self, part_num, line):
"""将按行分割后的内容写入相应的分割文件中"""
part_file_name = self.get_part_file_name(part_num)
try:
with open(part_file_name, "a") as part_file:
part_file.writelines(line)
except IOError as err:
print(err)
if __name__ == "__main__":
file = SplitFiles(r'H:/代码练习/Deeplearning/data_txt_path/all_data_list_1.txt')
file.split_file()
我这里将总的数据文件和划分好的数据集存在一个文件夹里，方便后面管理（本来是只有3个txt文件的，我弄了两个数据集，所以就有了6个文件）。
3 加载txt文件中的图片标签数据集
在完成步骤1和2后，最后是对数据进行加载，下面为加载数据的代码，后面读取数据调用这个类函数就可以：
import os
import numpy as np
import torch
from torchvision import transforms
from PIL import Image
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
# 我们读取图片的根目录， 在根目录下有所有图片的txt文件， 拿到txt文件后， 先读取txt文件， 之后遍历txt文件中的每一行， 首先去除掉尾部的换行符， 在以空格切分，前半部分是图片名称， 后半部分是图片标签， 当图片名称和根目录结合，就得到了我们的图片路径
class MyDataset(Dataset):
def __init__(self, img_path, transform=None):
super(MyDataset, self).__init__()
self.root = img_path
# self.txt_root = self.root + 'all_list.txt'
f = open(self.root, 'r')
data = f.readlines()
imgs = []
labels = []
# label_1,label_2,label_3 = [],[],[]
for line in data:
line = line.rstrip()
word = line.split()
imgs.append(os.path.join(self.root, word[1],word[2],word[3],word[0]))
# labels.append([float(word[1]),float(word[2]),float(word[3])])
labels.append([word[1],word[2],word[3]])
# label_1,label_2,label_3 = word[1],word[2],word[3]
# labels.append([[label_1],[label_2],[label_3]])
self.img = imgs
self.label = labels
self.transform = transform
# print(self.img)
# print(self.label)
def __len__(self):
return len(self.label)
return len(self.img)
def __getitem__(self, item):
img = self.img[item]
label = self.label[item]
# print(img)
img = Image.open(img).convert('RGB')
# 此时img是PIL.Image类型 label是str类型
if transforms is not None:
img = self.transform(img)
# print(img.max())
label = np.array(label).astype(np.float32)
label = torch.from_numpy(label)
return img, label
注意：上面/下图代码是我同时加载了三个标签，因为一张图片我是同时对应了三个标签，假如是一个图片对应一个标签，可在以下图片的函数中进行更改：
因为我的标签是浮点数，所以我在这里将其变为浮点数类型，假如是整形，可以在上面代码下图位置更改。
在执行完步骤1的代码文件后，将图片数据集打上标签并保存为txt文件；在执行步骤2的代码文件将txt文件中的图片标签数据集随机划分为训练集和测试集；最后编写步骤3加载txt文件中的图片标签数据集代码，就可加载自己的数据集 。下面是深度学习训练时，调用上面加载数据的类实现对数据的加载，也可根据自己的代码进行编写，可以参考一下下面的例子：
root_train = r'H:/代码练习/Deeplearning/data_txt_path/train_list_1.txt'
root_test = r'H:/代码练习/Deeplearning/data_txt_path/test_list_1.txt'
#将图像的像素值归一化到[-1,1]之间
normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
train_transform = transforms.Compose([
transforms.Resize((224,224)),
# transforms.RandomVerticalFlip(),
transforms.ToTensor(),
normalize])
val_transform = transforms.Compose([
transforms.Resize((224,224)),
# transforms.RandomVerticalFlip(),
transforms.ToTensor(),
normalize])
train_dataset = MyDataset(root_train,transform=train_transform)
val_dataset = MyDataset(root_test,transform=val_transform)
train_dataloader = DataLoader(dataset=train_dataset,batch_size=16,shuffle=True)
val_dataloader = DataLoader(dataset=val_dataset,batch_size=16,shuffle=True)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
for batch, (x, y) in enumerate(data_loader):
image, y= x.to(device), y.to(device)
参考来源：制作数据集（二）–为图片数据集打上标签并保存为txt文件_困坤的小菜鼠的博客-CSDN博客
python 划分数据集文件(txt标签文件按比例随机切分)_努力学习DePeng的博客-CSDN博客_python按比例随机切分数据
