自定义分割数据集调用时触发TypeError：get_train_augs()收到意外关键字参数'image'

Hey, let's fix that annoying TypeError you're running into! The root cause is easy to spot—here's what's going wrong and how to fix it:

The Problem #

When you create your dataset instances:

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

You're passing the function itself (get_train_augs) instead of the augmentation pipeline object that the function returns. Your dataset's __getitem__ method tries to call this function with image and mask arguments, but get_train_augs() doesn't accept any parameters—hence the "unexpected keyword argument" error.

Quickest Fix #

Just add parentheses to call the functions when passing them to your dataset. This way you're passing the Albumentations Compose object that knows how to handle image and mask:

trainset = SegmentationDataset(train_df, get_train_augs())  # Call the function to get the pipeline
validset = SegmentationDataset(valid_df, get_valid_augs())

Alternative Fix (If You Want to Keep Passing Functions) #

If you need to dynamically generate augmentations (e.g., for parameterized pipelines), you can modify your dataset class to handle function inputs. Update either the __init__ or __getitem__ method:

Option 1: Resolve in __init__ #

class SegmentationDataset(Dataset):
    def __init__(self, df, augmentations=None):
        self.df = df
        # If augmentations is a function, call it to get the pipeline
        self.augmentations = augmentations() if callable(augmentations) else augmentations
    # ... rest of your class code stays the same

Option 2: Resolve in __getitem__ #

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)
    mask = np.expand_dims(mask, axis = -1)
    
    if self.augmentations is not None:
        # Check if it's a function, call it to get the pipeline first
        aug_pipeline = self.augmentations() if callable(self.augmentations) else self.augmentations
        data = aug_pipeline(image = image, mask = mask)
        image = data['image']
        mask = data['mask']
    
    # ... rest of your preprocessing code stays the same

Why This Works #

Your get_train_augs() function returns an A.Compose instance—this is the object that accepts image and mask arguments and applies the augmentations. By passing the function without parentheses, you're telling the dataset to call the function itself, which doesn't expect those arguments. Adding parentheses runs the function upfront and passes the correct pipeline object.

内容的提问来源于stack exchange，提问作者Edward Chome