TypeError: got an unexpected keyword argument 'image'-CodePudding

I am getting TypeError: get_train_augs() got an unexpected keyword argument 'image', I have my augmentation functions as follows

Augmentation functions

def get_train_augs():
  return A.Compose([
                    A.Resize(IMAGE_SIZE,IMAGE_SIZE),
                    A.HorizontalFlip(p = 0.5),
                    A.VerticalFlip(p = 0.5),
  ])

def get_valid_augs():
  return A.Compose([
                    A.Resize(IMAGE_SIZE,IMAGE_SIZE),
                    
  ])

Custom segmentation dataset class

class SegmentationDataset(Dataset):
  def __init__(self, df, augmentations=None):
    self.df = df
    self.augmentations = augmentations
  
  def __len__(self):
    return len(self.df)

  def __getitem__(self,idx):
    row = self.df.iloc[idx]

    image_path = row.images
    mask_path = row.masks

    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # (h, w, c)
    mask = np.expand_dims(mask, axis = -1)

    if self.augmentations is not None:
      data = self.augmentations(image = image, mask = mask)

      image = data['image']
      mask = data['mask']

    # (h, w, c) -> (c, h, w)
    image = np.transpose(image,(2,0,1)).astype(np.float32)
    mask = np.transpose(mask,(2,0,1)).astype(np.float32)

    image = torch.Tensor(image)/255.0
    mask = torch.round(torch.Tensor(mask)/255.0)

    return image, mask

when I call trainset like this I am getting an error:

trainset = SegmentationDataset(train_df, get_train_augs)
validset = SegmentationDataset(valid_df, get_valid_augs)

calling a random index

idx = 3
image, mask = trainset[idx]

The error I am getting is:

TypeError                                 Traceback (most recent call last)
<ipython-input-28-9b83781b7e3d> in <module>()
      1 idx = 3
      2 
----> 3 image, mask = trainset[idx]
      4 
      5 helper.show(image, mask)

<ipython-input-25-39872478644d> in __getitem__(self, idx)
     20 
     21     if self.augmentations is not None:
---> 22       data = self.augmentations(image = image, mask = mask)
     23 
     24       image = data['image']

TypeError: get_train_augs() got an unexpected keyword argument 'image'

CodePudding user response：

try

trainset = SegmentationDataset(train_df, get_train_augs())
validset = SegmentationDataset(valid_df, get_valid_augs())