[FIXED] Augmentation using Albumentations in Pytorch OD

Issue

I followed the pytorch tutorial for object detection on the website here. I decided to add more augmentations using albumentations if it would improve my traning. However after calling the __getitem__() method in the dataset class I get this error.

AttributeError                            Traceback (most recent call last)
<ipython-input-54-563a9295c274> in <module>()
----> 1 train_ds.__getitem__(220)

2 frames
<ipython-input-48-0169e540fb13> in __getitem__(self, idx)
     45       }
     46 
---> 47       transformed = self.transforms(**image_data)
     48       img = transformer['image']
     49       target['boxes'] = torch.as_tensor(transformed['bboxes'],dtype=torch.float332)

/usr/local/lib/python3.7/dist-packages/albumentations/core/composition.py in __call__(self, force_apply, **data)
    172             if dual_start_end is not None and idx == dual_start_end[0]:
    173                 for p in self.processors.values():
--> 174                     p.preprocess(data)
    175 
    176             data = t(force_apply=force_apply, **data)

/usr/local/lib/python3.7/dist-packages/albumentations/core/utils.py in preprocess(self, data)
     58         data = self.add_label_fields_to_data(data)
     59 
---> 60         rows, cols = data["image"].shape[:2]
     61         for data_name in self.data_fields:
     62             data[data_name] = self.check_and_convert(data[data_name], rows, cols, direction="to")

AttributeError: 'Image' object has no attribute 'shape'

I have include the augmentation codes I used as well.

    def transform_ds(train):
  if train:
    return A.Compose([
                      A.HorizontalFlip(p=0.2),
                      A.VerticalFlip(p=0.2),
                      A.RandomSizedBBoxSafeCrop(height=450,width=450,erosion_rate=0.2,p=0.3),
                      A.RandomBrightness(limit=(0.2,0.5),p=0.3),
                      A.RandomContrast(limit=(0.2,0.5),p=0.3),
                      A.Rotate(limit=90,p=0.3),
                      A.GaussianBlur(blur_limit=(3,3),p=0.1),
                      ToTensorV2()
    ], bbox_params=A.BboxParams(
        format='pascal_voc',
        min_area=0, 
        min_visibility=0,
        label_fields=['labels']
    ))

  else:
    return A.Compose([ToTensor()])

Solution

Images in PyTorch are loaded via pillow library (PIL.Image.open specifically).

If you look at albumentations docs its transformations required torch.Tensor (or np.ndarray object).

In order to do it, you should place A.ToTensorV2 as a first transformation and use other documentation transforms after that.

Answered By - Szymon Maszke