import os

import torch
import xmltodict  # 用于将XML数据转换为Python字典
from PIL import Image  # Python图像处理库
from torch.utils.data import Dataset  # PyTorch数据集基类
from torchvision import transforms  # 图像预处理工具


class VOCDataset(Dataset):
    # VOC数据集初始化
    def __init__(self, image_dir, label_dir, transform, label_transform):
        """
        初始化VOC格式数据集
        Args:
            image_dir: 图像文件目录路径
            label_dir: 标签文件目录路径
            transform: 图像预处理变换
            label_transform: 标签预处理变换（代码中未实际使用）
        """
        self.image_dir = image_dir
        self.label_dir = label_dir
        self.transform = transform
        self.label_transform = label_transform
        self.imgs = os.listdir(self.image_dir)  # 获取图像目录下所有文件名
        # 定义类别列表
        self.classes_list = ["no helmet", "motor", "number", "with helmet"]

    # 返回数据集大小
    def __len__(self):
        """返回数据集中的图像数量"""
        return len(self.imgs)

    # 获取单个数据项
    def __getitem__(self, index):
        """获取索引对应的数据项"""
        # 获取图像文件名并构建完整路径
        img_name = self.imgs[index]
        img_path = os.path.join(self.image_dir, img_name)
        image = Image.open(img_path).convert('RGB')  # 打开图像并转换为RGB格式
        # 构建对应的XML标签文件路径（将图像扩展名替换为.xml）
        label_name = img_name.split('.')[0] + '.xml'
        label_path = os.path.join(self.label_dir, label_name)
        # 解析XML标签文件
        with open(label_path, "r", encoding="utf-8") as f:
            label_content = f.read()
        label_dict = xmltodict.parse(label_content)  # 将XML转换为字典
        # 提取标注信息
        objects = label_dict["annotation"]["object"]
        target = []  # 存储所有目标的标注信息
        # 遍历每个标注对象
        for obj in objects:
            obj_name = obj["name"]  # 获取对象类别名称
            obj_class_id = self.classes_list.index(obj_name)  # 将类别名称转换为索引ID
            # 提取边界框坐标（注意转换为浮点数）
            xmin = float(obj["bndbox"]["xmin"])
            ymin = float(obj["bndbox"]["ymin"])
            xmax = float(obj["bndbox"]["xmax"])
            ymax = float(obj["bndbox"]["ymax"])
            # 将当前目标的标注信息添加到列表
            target.extend([obj_class_id, xmin, ymin, xmax, ymax])
        # 将标注列表转换为张量
        target = torch.tensor(target)
        # 应用图像预处理变换
        if self.transform is not None:
            image = self.transform(image)
        return image, target  # 返回处理后的图像和标注张量


if __name__ == '__main__':
    image_dir = r'E:\deeplearning\torch\data\HelmetDataset-VOC\train\images'
    label_dir = r'E:\deeplearning\torch\data\HelmetDataset-VOC\train\labels'
    # 创建数据集实例
    train_dataset = VOCDataset(image_dir, label_dir, transforms.Compose([transforms.ToTensor()]), None)
    print(len(train_dataset))  # 打印数据集大小
    print(train_dataset[1])