PyTorch加载多JSON标注图像及自定义数据集微调问题咨询

嘿，我来帮你搞定这个问题！你已经写好了自定义的lDataSet，接下来只需要把它替换掉官方示例里的ImageFolder，再做一些小调整就能完美适配啦，具体步骤如下：

1. 修正并完善自定义Dataset #

先把你Dataset里的小bug和缺失的部分补上，比如数据转换、标注处理，还有适配数据增强：

import os
import json
import cv2
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

class lDataSet(Dataset):
    def __init__(self, path_to_imgs, path_to_json, transform=None):
        self.path_to_imgs = path_to_imgs
        self.path_to_json = path_to_json
        # 确保只获取jpg文件，避免其他干扰文件
        self.img_ids = [f.split('.')[0] for f in os.listdir(path_to_imgs) if f.endswith('.jpg')]
        self.transform = transform

    def __getitem__(self, idx):
        img_id = self.img_ids[idx]
        # 读取图像并转成RGB（cv2默认是BGR）
        img = cv2.imread(os.path.join(self.path_to_imgs, f"{img_id}.jpg"))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # 加载并处理COCO格式的标注
        with open(os.path.join(self.path_to_json, f"{img_id}.json")) as f:
            annot = json.load(f)
        
        # 提取需要的标注信息（比如bbox、类别id，根据你的任务调整）
        bboxes = []
        labels = []
        for obj in annot['annotations']:
            bboxes.append(obj['bbox'])
            labels.append(obj['category_id'])
        
        # 转成Tensor格式
        bboxes = torch.tensor(bboxes, dtype=torch.float32)
        labels = torch.tensor(labels, dtype=torch.long)
        
        # 应用数据增强
        if self.transform:
            # 注意：如果用torchvision的transform，需要先把numpy转成PIL图像
            img = transforms.ToPILImage()(img)
            img = self.transform(img)
        
        return img, {'bboxes': bboxes, 'labels': labels}

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

这里主要做了这些优化：

• 过滤出仅jpg格式的图像，避免读取无效文件
• 把cv2读取的BGR图像转成RGB，符合PyTorch的默认格式
• 解析COCO标注里的bbox和类别，并转成Tensor
• 支持传入数据增强的transform，和官方示例对齐

2. 替换官方示例的数据集创建部分 #

官方示例里用ImageFolder创建数据集，现在换成你的lDataSet：

# 定义数据增强（和官方示例保持一致）
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

# 定义数据根目录
data_dir = "/home"

# 创建训练和验证数据集（替换ImageFolder为自定义Dataset）
image_datasets = {
    'train': lDataSet(
        path_to_imgs=os.path.join(data_dir, 'training/images'),
        path_to_json=os.path.join(data_dir, 'training/json_annotations'),
        transform=data_transforms['train']
    ),
    'val': lDataSet(
        path_to_imgs=os.path.join(data_dir, 'validation/images'),
        path_to_json=os.path.join(data_dir, 'validation/json_annotations'),
        transform=data_transforms['val']
    )
}

3. 创建DataLoader（和官方示例基本一致） #

这一步只需要加个小适配，因为每张图的目标数量可能不一样：

batch_size = 8
dataloaders_dict = {
    x: DataLoader(
        image_datasets[x], 
        batch_size=batch_size, 
        shuffle=True, 
        num_workers=4,
        # 处理变长标注：默认collate函数会报错，用这个lambda把图像和标注分别打包成列表
        collate_fn=lambda batch: tuple(zip(*batch))
    ) for x in ['train', 'val']
}

4. 在训练循环中使用 #

现在你可以像官方示例那样在训练循环里迭代DataLoader了，比如：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 假设你已经定义好了适配多目标任务的模型（比如基于Faster R-CNN的模型）
model = ... # 初始化你的模型
model = model.to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练循环示例
for epoch in range(10):
    print(f'Epoch {epoch}/{9}')
    print('-' * 10)

    for phase in ['train', 'val']:
        if phase == 'train':
            model.train()
        else:
            model.eval()
        
        running_loss = 0.0
        for imgs, annotations in dataloaders_dict[phase]:
            imgs = [img.to(device) for img in imgs]
            # 把标注转移到设备上
            annotations = [{k: v.to(device) for k, v in ann.items()} for ann in annotations]
            
            with torch.set_grad_enabled(phase == 'train'):
                outputs = model(imgs, annotations)
                loss = sum(loss_val for loss_val in outputs.values())
                
                if phase == 'train':
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
            
            running_loss += loss.item() * len(imgs)
        
        epoch_loss = running_loss / len(image_datasets[phase])
        print(f'{phase} Loss: {epoch_loss:.4f}')

这样就完美把你的自定义数据集和官方的微调流程结合起来啦，核心就是用自己的lDataSet替换掉ImageFolder，再根据多目标任务的特性处理变长标注~

内容的提问来源于stack exchange，提问作者bibinwilson