YOLOv3使用Pytorch进行数据集的加载的代码

以下是使用 PyTorch 加载数据集的 YOLOv3 代码示例：

import torch
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image

class YOLODataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None):
        self.img_dir = img_dir
        self.transform = transform

        with open(annotations_file, "r") as f:
            self.annotations = list(f)

    def __len__(self):
        return len(self.annotations)

    def __getitem__(self, index):
        annotation = self.annotations[index].strip().split(",")

        img_path = annotation[0]
        bbox = list(map(int, annotation[1:]))

        img = Image.open(img_path).convert("RGB")

        if self.transform:
            img = self.transform(img)

        return img, torch.tensor(bbox)

# 定义数据集的转换
transform = transforms.Compose([
    transforms.Resize((416, 416)),
    transforms.ToTensor(),
])

# 加载数据集
dataset = YOLODataset("annotations.txt", "images/", transform=transform)
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

# 遍历数据集
for images, bboxes in dataloader:
    print(images.shape, bboxes.shape)

在上面的代码中，我们定义了一个 YOLODataset 类，该类继承自 PyTorch 的 Dataset 类。我们在 `__init__` 方法中加载了标注文件和图像目录，并定义了数据集的转换。在 `__getitem__` 方法中，我们首先读取图像，然后将其转换为 PyTorch 张量，并返回图像和边界框的张量。最后，在主函数中，我们使用 DataLoader 类加载数据集，并遍历数据集。