How to implement OneOf ?

[Optimox]

Hello,

Thank you very much for this repo, it's very useful.

I wanted to add a new compostion OneOf where you can give a list of augmentations and only one of them is executed according to a random weight.

I tried this but it does not seem to work, could you help me please?

class OneOf:
    def __init__(self, transforms, p=1., weights=[], targets=[['image'],['mask']]):
        self.transforms = transforms
        self.weights = weights
        self.targets = targets
        self.p = p

    def __call__(self, force_apply=False, **data):
        need_to_run = force_apply or random.random() < self.p
        which_to_run = random.choices(range(len(self.transforms)),
                                        weights=self.weights)[0]
        transforms = self.transforms[which_to_run] if need_to_run else self.get_always_apply_transforms()
        data = transforms(force_apply, self.targets, **data)

        return data