在 Pytorch OD 中使用 Albumentations 进行增强

Augmentation using Albumentations in Pytorch OD

我按照网站上的对象检测 pytorch 教程进行操作 here。如果可以改善我的训练,我决定使用 albumentations 添加更多增强功能。但是,在数据集 class 中调用 __getitem__() 方法后,出现此错误。

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'

我也包含了我使用的增强代码。

    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()])

PyTorch 中的图像是通过枕头库加载的(PIL.Image.open)。

如果您查看 albumentations docs 其转换需要 torch.Tensor(或 np.ndarray 对象)。

为此,您应该将 A.ToTensorV2 作为第一个转换,然后使用其他文档转换。