我在使用 DataLoader 时遇到 KeyError: 'bboxes'，为什么会发生这种情况？

Question

我正在使用 PyTorch 框架进行图像分类，并遇到以下错误：

KeyError: Caught KeyError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
    data = fetcher.fetch(index)
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 51, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 51, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/tmp/ipykernel_32/245518470.py", line 32, in __getitem__
    augmented = self.albu_transform(image=image)
  File "/opt/conda/lib/python3.10/site-packages/albumentations/core/composition.py", line 207, in __call__
    p.preprocess(data)
  File "/opt/conda/lib/python3.10/site-packages/albumentations/core/utils.py", line 83, in preprocess
    data[data_name] = self.check_and_convert(data[data_name], rows, cols, direction="to")
KeyError: 'bboxes

我的自定义数据如下： 提供了图像数据集及其标签、类的长度（7）和变换。 请注意，在变换中，图像增强是由 albumentations 使用的。

import os
import pandas as pd
from PIL import Image
import torch
from torch.utils.data import Dataset
import torchvision.transforms as transforms
import albumentations as A
from albumentations.pytorch import ToTensorV2

class CustomDataset(Dataset):
    def __init__(self, image_paths, label_paths, classes, transform=None):
        self.image_paths = image_paths
        self.label_paths = label_paths
        self.transform = transform
        self.classes = classes
        self.class_to_idx = {class_name: idx for idx, class_name in enumerate(self.classes)}
        self.idx_to_class = {idx: class_name for class_name, idx in self.class_to_idx.items()}

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

    def __getitem__(self, idx):
        # Load image and annotations
        img_path = self.image_paths[idx]
        annot_path = self.label_paths[idx]
        
        img = Image.open('/kaggle/input/dataset-vehicle1/images/'+img_path).convert("RGB")
        w, h = img.size
        
        annot_data = pd.read_csv('/kaggle/input/labels1/labels/'+annot_path, delimiter=" ", header=None).values
        class_data = annot_data[:, 0]
        class_idx = class_data
         
        boxes = []
        labels = []
        for bbox in annot_data:
            class_name_or_idx, center_x_normalized, center_y_normalized, width_normalized, height_normalized = bbox
            center_x_pixel = int(center_x_normalized * w)
            center_y_pixel = int(center_y_normalized * h)
            width_pixel = int(width_normalized * w)
            height_pixel = int(height_normalized * h)
            
            # Convert (center_x, center_y, width, height) to (x_min, y_min, x_max, y_max)
            top_left_x = center_x_pixel - width_pixel // 2
            top_left_y = center_y_pixel - height_pixel // 2
            bottom_right_x = center_x_pixel + width_pixel // 2
            bottom_right_y = center_y_pixel + height_pixel // 2

            boxes.append([top_left_x, top_left_y, bottom_right_x, bottom_right_y])
            labels.append(int(class_idx[0].flatten()))
            
        # Apply transformations
        if self.transform is not None:
            try:
                print(f"Index: {idx}, Boxes: {boxes}")
                augmented = self.transform(image=np.array(img), bboxes=boxes, labels=labels)
                img = augmented['image']
                boxes = augmented['bboxes']
                
            except Exception as e:
                print(f"Error at index {idx}: {e}")

        img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
        target = {
            "boxes": torch.tensor(boxes, dtype=torch.float32),
            "labels": torch.tensor(labels, dtype=torch.int64),
            "image_id": torch.tensor(idx)
        }
        print(f"Index: {idx}")
        print("Image Shape:", np.array(img).shape)
        print("Boxes:", boxes)
        print("Labels:", labels)

        return img, target

定义自定义数据并使用以下数据加载器后

train_dataset = CustomDataset(train_image_paths, train_label_paths, classes=vehicle_classes, transform=albu_transform)
training_loader= DataLoader(train_dataset, batch_size=4,shuffle=True,num_workers= 0,
                    collate_fn=collate_fn,
                    pin_memory=True if device=="cuda" else False)

在训练模型之前，我收到 KeyError: 如上所述：

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

for i, (img, target) in enumerate(train_loader):
    image = image.to(device)
    target['boxes'] = target['boxes'].to(device)
    target['labels'] = target['labels'].to(device)
    target['image_id'] = target['image_id'].to(device)

Answer 1

看来您正在使用 yolo 格式标签。

你应该像下面这样调用albumentations

它的版本，专辑=1.3.1

import albumentations as A
P = 0.2
AAAtransform = A.Compose([
    A.HorizontalFlip(p=P),
    A.RandomGamma(p=P),
], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))

class trainDataset(Dataset):
        image = cv2.imread(imgpath)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        bboxes.append([center_x, center_y, width, height])  # [0, 1] normalized
        labels.append(int(label))

    def __getitem__(self, idx):
        transformed = AAAtransform(image=image, bboxes = bboxes, class_labels = labels, )
        image = transformed['image']
        gtt = transformed['bboxes']
        class_labels = transformed['class_labels']