### 自定义Dataset类加载图像和边界框数据并结合DataLoader实现批量加载与预处理 在PyTorch中，`torch.utils.data.Dataset` 是一个抽象类，用户可以通过继承该类来创建自定义的数据集。为了加载图像和边界框数据，需要重写 `__init__`、`__len__` 和 `__getitem__` 方法。此外，通过使用 `torch.utils.data.DataLoader`，可以实现数据的批量加载、打乱顺序以及多进程处理。 以下是一个完整的实现示例： #### 1. 自定义Dataset类 ```python import torch from torch.utils.data import Dataset from PIL import Image import os import json class CustomDataset(Dataset): def __init__(self, root_dir, transform=None, target_transform=None): """ 初始化自定义数据集。 :param root_dir: 数据集根目录，包含图像和边界框信息。 :param transform: 对图像进行的变换操作。 :param target_transform: 对边界框进行的变换操作。 """ self.root_dir = root_dir self.transform = transform self.target_transform = target_transform self.image_dir = os.path.join(root_dir, "images") self.annotation_dir = os.path.join(root_dir, "annotations") # 获取所有图像文件名 self.image_files = [f for f in os.listdir(self.image_dir) if f.endswith(".jpg")] def __len__(self): """ 返回数据集中的样本数量。 """ return len(self.image_files) def __getitem__(self, idx): """ 根据索引获取单个样本。 :param idx: 样本索引。 :return: 图像和边界框数据。 """ image_name = self.image_files[idx] image_path = os.path.join(self.image_dir, image_name) annotation_path = os.path.join(self.annotation_dir, image_name.replace(".jpg", ".json")) # 加载图像 image = Image.open(image_path).convert("RGB") # 加载边界框信息（假设为JSON格式） with open(annotation_path, 'r') as f: annotation = json.load(f) boxes = torch.tensor(annotation['boxes'], dtype=torch.float32) labels = torch.tensor(annotation['labels'], dtype=torch.int64) # 应用图像变换 if self.transform: image = self.transform(image) # 应用边界框变换 if self.target_transform: boxes = self.target_transform(boxes) return image, {"boxes": boxes, "labels": labels} ``` #### 2. 使用DataLoader进行批量加载 ```python from torch.utils.data import DataLoader from torchvision import transforms # 定义图像变换 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # 创建自定义数据集实例 dataset = CustomDataset(root_dir="path/to/dataset", transform=transform) # 创建DataLoader data_loader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=4, collate_fn=lambda x: tuple(zip(*x))) # 遍历DataLoader for images, targets in data_loader: # 处理批量数据 pass ``` 上述代码实现了以下功能： - **自定义Dataset类**：加载图像和边界框数据，并支持图像和边界框的变换[^1]。 - **DataLoader**：批量加载数据，支持数据打乱和多进程处理[^1]。 #### 注意事项 - 数据集目录结构应包括 `images` 和 `annotations` 子目录，分别存储图像文件和对应的边界框信息文件。 - 边界框信息以 JSON 文件形式存储，包含 `boxes` 和 `labels` 字段。 - `collate_fn` 参数用于处理不规则形状的数据（例如不同大小的边界框列表）。 ###