深度学习制作自己的数据集—为数据集打上标签保存为txt文件，并进行划分和加载数据集

目录

0 前言

1 为图片数据集打上标签并保存为txt文件

2 将txt文件中的图片标签数据集随机划分为训练集和测试集

3 加载txt文件中的图片标签数据集

0 前言

目前是被封控的第四天了，只能呆在宿舍不能出去，记得上次这样子还是一年前大四快毕业那时候了……

这几天在宿舍没有什么事干，实验也暂时做不了了，将部分数据处理完后，就把之前的这个内容做一下笔记吧，这也不是什么新的知识了，简单记录一下，方便以后可以查看。

1 为图片数据集打上标签并保存为txt文件

由于这里我做的是用深度学习回归预测，所以我的标签保存在（.csv）文件中，这时候需要将图片和标签一一对应起来，并且要分好文件夹，下面是我分好的文件夹（images保存的是图片，label.csv保存的是对应的标签，这里可以根据个人的数据集更改文件名称）：

下面是为图片数据集打上标签并保存为txt文件的代码（文件的路劲需要根据自己文件所在位置进行更改）：

import osimport numpy as npimport pandas as pdlabel = 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 = 0listText = 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,则继续遍历循环，否则退出循环continuename = 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文件中写入nametarget = ''i += 1listText.write(name)#在创建的txt文件中写入namelistText.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 osimport random# 划分比例，训练集 : 验证集 = 8 : 2split_rate = 0.2class 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_numdef 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 = 1file = 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 += 1except 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_pathif not os.path.exists(file_folder):# 如果临时目录不存在则创建os.makedirs(file_folder)part_file_name = file_folder + "/" + str(part_name) + "_list_1.txt"return part_file_namedef 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 osimport numpy as npimport torchfrom torchvision import transformsfrom PIL import Imagefrom torch.utils.data import DataLoaderfrom 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 = imgsself.label = labelsself.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按比例随机切分数据

pytorch加载自己的图片数据集的两种方法__-周-_的博客-CSDN博客_pytorch读取图片数据集