机器学习复习（9）——自定义dataset

目录

第一种dataset(文件夹名即为标签)

用于将格式（1）转换为格式（2）

第二种dataset(标签在labels文件夹下的对应的txt文件里面)

---

第一种dataset(文件夹名即为标签)

数据组织格式（1）

--data

----train

------class1(文件夹名字即为标签)

--------image1.jpg

------class2

dataset

from torch.utils.data import Dataset
from  PIL import Image
class Mydata(Dataset):def __init__(self,root_dir,label_dir):self.root_dir = root_dirself.label_dir = label_dirself.path = os.path.join(self.root_dir,self.label_dir)self.img_path= os.listdir(self.path)def __getitem__(self, idx):img_name = self.img_path[idx]img_item_path = os.path.join(self.path,img_name)img = Image.open(img_item_path)label =self.label_dirreturn img,labeldef __len__(self):return len(self.img_path)root_dir="../PATH/TO/train"
class1_label_dir="class1"
class2_label_dir="class2"class1_data =Mydata(root_dir,class1_label_dir)
class2_data =Mydata(root_dir,class2_label_dir)    train_dataset= class1_data+class2_data

用于将格式（1）转换为格式（2）

数据集格式转换

import os
root_dir = 'root_path'
target_dir = 'target_image'
img_path = os.listdir(os.path.join(root_dir, target_dir))
label = target_dir.split('_')[0]
out_dir = 'ants_label'
for i in img_path:file_name = i.split('.jpg')[0]with open(os.path.join(root_dir, out_dir,"{}.txt".format(file_name)),'w') as f:f.write(label)

第二种dataset(标签在labels文件夹下的对应的txt文件里面)

数据组织格式（2）

--data

----train

------images

--------01.jpg

------labels

--------01.txt        (txt里面的内容是label内容：目标检测，分类等）

# 导入PyTorch的数据集工具和其他必要的库
from torch.utils.data import Dataset
import os
from PIL import Image# 自定义的数据集类，继承自torch.utils.data.Dataset
class Mydate(Dataset):def __init__(self, dir_root, dir_image, dir_label):# 初始化函数，设置数据集的根目录、图像目录和标签目录self.root = dir_root  # 数据集的根目录self.image_dir = dir_image  # 存放图像的子目录self.image_path = os.path.join(self.root, self.image_dir)  # 图像的完整路径self.label_dir = dir_label  # 存放标签的子目录self.label_path = os.path.join(self.root, self.label_dir)  # 标签的完整路径# 获取图像文件列表和标签文件列表self.image_list = os.listdir(os.path.join(dir_root, dir_image))  # 根据图像目录列出所有图像文件self.label_list = os.listdir(os.path.join(dir_root, dir_label))  # 根据标签目录列出所有标签文件def __getitem__(self, idx):# 通过索引获取数据集中的单个样本，包括图像和标签image_name = self.image_list[idx]  # 根据索引获取图像文件名image_path = os.path.join(self.image_path, image_name)  # 构造图像的完整路径label_name = self.label_list[idx]  # 根据索引获取标签文件名label_path = os.path.join(self.label_path, label_name)  # 构造标签的完整路径img = Image.open(image_path).convert('RGB')  # 打开图像文件并转换为RGB格式# 读取标签文件with open(label_path, 'r') as f:label = f.read().strip()  # 读取标签内容，并去除可能的空白字符return img, labeldef __len__(self):# 返回数据集中样本的总数return len(self.image_list)  # 由于图像列表的长度代表了数据集大小，直接返回其长度##############测试代码################
if __name__ == "__main__":# 指定数据集的根目录、图像目录和标签目录dir_root = "root\\path\\train"dir_image = "images"dir_label = "labels"test_data = Mydate(dir_root, dir_image, dir_label)  # 创建数据集实例idx = 0  # 指定要获取的样本索引img, label = test_data[idx]  # 获取指定索引的样本print(label)  # 打印样本的标签