Dataset类代码实战

from torch.utils.data import Dataset
from PIL import Image
import os

class MyData(Dataset):
    def __init__(self, root_dir, label_dir):
        """
        初始化数据集
        
        参数:
            root_dir (str): 数据集的根目录。
            label_dir (str): 该标签的图像子目录。
        """
        self.root_dir = root_dir
        self.label_dir = label_dir
        # 构建标签目录的路径
        self.path = os.path.join(self.root_dir, self.label_dir)
        # 列出标签目录中的所有图像文件
        self.img_path = os.listdir(self.path)

    def __getitem__(self, idx):
        """
        获取图像及其标签
        
        参数:
            idx (int): 要获取的项目索引。
        
        返回:
            tuple: (图像, 标签)
        """
        # 获取图像文件名
        img_name = self.img_path[idx]
        # 构建完整的图像文件路径
        img_item_path = os.path.join(self.root_dir, self.label_dir, img_name)
        # 打开图像文件
        img = Image.open(img_item_path)
        # 图像的标签
        label = self.label_dir
        return img, label

    def __len__(self):
        """
        返回数据集中图像的数量
        
        返回:
            int: 图像数量
        """
        return len(self.img_path)


# 定义数据集路径
root_dir = "dataset/train"
ants_label_dir = "ants"
bees_label_dir = "bees"

# 创建蚂蚁和蜜蜂的数据集实例
ants_dataset = MyData(root_dir, ants_label_dir)
bees_dataset = MyData(root_dir, bees_label_dir)

# 合并数据集（注意：PyTorch 的 Dataset 类不支持直接用 + 连接）
train_dataset = ants_dataset + bees_dataset

代码标记：

import os

root_dir = 'dataset/train'  # 数据集的根目录
target_dir = 'ants_image'   # 包含图像的子目录

# 获取目标目录中的所有图像文件名
img_path = os.listdir(os.path.join(root_dir, target_dir))  

# 从目标目录名中提取标签（例如，从 'ants_image' 提取 'ants'）
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)  # 将标签写入文件

实战编写代码注意点

1.__int__与__init__的区别，别被Python编译器自动补充搞错了

2.绝对路径写时要加\\不能是\