Custom transform for augmenting image using mask in detectron2

[ravijo]

In the process of data augmentation, I am trying to modify the image based on the corresponding mask. I need to get the background from the image, which requires knowing the foreground (mask) in advance.

Therefore, I am looking for a Transform that can provide image and mask as input to my function. For example, previously, I used ColorTransform, which takes a callable and provides an image as input to this callable. See an example below:

def apply_blur(img):
  blur_img = cv2.GaussianBlur(img, (3, 3), 0)
  return blur_img

class MyTrainer(DefaultTrainer):
  @classmethod
  def build_train_loader(cls, cfg):
    augmentations = [
        T.RandomFlip(),
        T.ColorTransform(apply_blur),
    ]
    mapper = DatasetMapper(cfg, is_train=True, augmentations=augmentations)
    return build_detection_train_loader(cfg, mapper=mapper)

I found AugInput, but it expects image, mask, etc. as input. Instead, is there something like the following exists?

def my_transform(img, mask):
  # do something on image using mask and return image
  return img

  ...
  def build_train_loader(cls, cfg):
    augmentations = [
        T.RandomFlip(),
        T.CustomTransform(my_transform),
    ]
  ...

Thank you very much

[ravijo]

Any hints, please?