0

I am working on image classification by using the PyTorch framework and struggling with the following error:

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

My custom data is as follows: the image dataset along with its labels, the length of the classes which is 7 and the transform are provided. please note that in the transform image augmentation is utilized by 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

after defining the custom data and using the following data loader

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)

I am getting the KeyError: as mentioned above before training the model:

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)
mkrieger1
  • 19,194
  • 5
  • 54
  • 65
  • You did not by any chance mean to write `target = {"bboxes": torch.tensor(... }` instead of `target = {"boxes": torch.tensor(... }`, or `target['bboxes'] = target['bboxes'].to(device)` instead of `target['boxes'] = target['boxes'].to(device)`? – mkrieger1 Aug 29 '23 at 00:55
  • Post the code where you define the transform. It looks like your transform does not expect bboxes, or else perhaps does not return bboxes? The error comes in the transform application line, – DerekG Aug 29 '23 at 01:51
  • from albumentations.pytorch import ToTensorV2 # Define your transformations as before, but now include bbox_params albu_transform = A.Compose([ A.HorizontalFlip(), A.ShiftScaleRotate(rotate_limit=15,value=0, border_mode=cv2.BORDER_CONSTANT), A.CLAHE(), A.GaussNoise(), A.LongestMaxSize(max_size=416), A.PadIfNeeded(min_height=416, min_width=416, border_mode=0), ], A.BboxParams(format='pascal_voc', label_fields=['labels'])) – ibrahim ahmethan Aug 29 '23 at 09:48

1 Answers1

0

it seems like that you are using yolo format label.

you should call albumentations like below

it's version, 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']
Jiu_Zou
  • 463
  • 1
  • 4