So after running through the code flow and documentation, I found out that Each Augmentation class is dependent on the Transform class which is inherited from the fvcore library
Also the same dependency is also defined vaguely in this highlighted documentation block
Then looking at all the aspects above and then going through this github PR merge fight thread, I sense that I'm on the right track so following PILColorTransform augmentation, I have made this code which works perfectly for custom augmentation
1. CUSTOM Augmentations
class GenericWrapperTransform(Transform):
"""
Generic wrapper for any transform (for color transform only. You can give functionality to apply_coods, apply_segmentation too)
"""
def __init__(self, custom_function:Callable):
"""
Args:
custom_function (Callable): operation to be applied to the image which takes in an ndarray and returns an ndarray.
"""
if not callable(custom_function):
raise ValueError("'custom_function' should be callable")
super().__init__()
self._set_attributes(locals())
def apply_image(self, img):
'''
apply transformation to image array based on the `custom_function`
'''
return self.custom_function(img)
def apply_coords(self, coords):
'''
Apply transformations to Bounding Box Coordinates. Currently is won't do anything but we can change this based on our use case
'''
return coords
def inverse(self):
return NoOpTransform()
def apply_segmentation(self, segmentation):
'''
Apply transformations to segmentation. currently is won't do anything but we can change this based on our use case
'''
return segmentation
class CustomAug(Augmentation):
"""
Given a probability and a custom function, return a GenericWrapperTransform object whose `apply_image` will be called to perform augmentation
"""
def __init__(self, custom_function, prob=1.0):
"""
Args:
custom_op: Operation to use. Must be a function takes an ndarray and returns an ndarray
prob (float): probability of applying the function
"""
super().__init__()
self._init(locals())
def get_transform(self, image):
'''
Based on probability, choose whether you want to apply the given function or not
'''
do = self._rand_range() < self.prob
if do:
return GenericWrapperTransform(self.custom_function)
else:
return NoOpTransform() # it returns a Transform which just returns the original Image array only
def white(image):
return np.ones(image.shape, dtype = np.uint8)*255 # returns white Image
def black(image):
return np.zeros(image.shape, dtype=np.uint8) # returns black image
def rand(image):
return np.random.randint(0,256,image.shape, dtype = np.uint8) # returns random image
def default(image):
return image # returns original image
2. Implementing OneOf like functions
I looked into the implementation of AugmentationList and built my own code to generate transformations randomly. Setting k=1 mimics OneOf like functionality.
class KRandomAugmentationList(Augmentation):
"""
Select and Apply "K" augmentations in "RANDOM" order with "Every" __call__ method invoke
"""
def __init__(self, augs, k:int = -1):
"""
Args:
augs: list of [Augmentation or Transform]
k: Number of augment to use from the given list in range [1,len_augs]. If None, use all. If it is -1, generate K randomly between [1,len_augs]
"""
super().__init__()
self.max_range = len(augs)
self.k = k
self.augs = augs # set augs to use as fixed if we have to use same augs everytime
def _setup_augs(self, augs, k:int):
'''
Setup the argument list. Generates the list of argument to use from the given list
args:
augs: list of [Augmentation or Transform])
k: Number of augment to use from the given list in range [1,len_augs]. If False, use all. If it is -1, generate K randomly between [1,len_augs]
'''
if k == -1: # Generate a random number
k = np.random.randint(1,len(augs)+1)
elif k is None: # use all
k = self.max_range
temp = np.random.choice(augs,k,replace=False) # get k augments randomly
return [_transform_to_aug(x) for x in temp]
def __call__(self, aug_input) -> Transform:
tfms = []
for x in self._setup_augs(self.augs, self.k): # generate auguments to use randomly on the fly
print(x)
tfm = x(aug_input)
tfms.append(tfm)
return TransformList(tfms)
def __repr__(self):
msgs = [str(x) for x in self.augs]
return "AugmentationList[{}]".format(", ".join(msgs))
__str__ = __repr__
3. Putting Everything together
from detectron2.data import transforms as T
import numpy as np
from PIL import Image
augs = KRandomAugmentationList(
[
# my custom augs
CustomAug(white),
CustomAug(black),
CustomAug(default),
CustomAug(rand),
# augs from Detectron
T.RandomBrightness(0.4, 1.6),
T.RandomSaturation(0.4, 1.6),
T.RandomContrast(0.4,1.6),
T.RandomCrop("absolute", (640, 640)),
T.RandomFlip(prob=0.5),
],
k = -1)
# Calling the below block multiple times will give you different combinations
# of Augmentations everytime depending on the value of `k`
image = np.array(Image.open("my_image.png")) # RGB image array
input_ = T.AugInput(image, boxes=None, sem_seg = None) # boxes and segments are optional
transform = augs(input_) # Apply the augmentation
image_transformed = input_.image # augmented image
Image.fromarray(image_transformed) # show RGB image
