<jupyter_start><jupyter_text>Data augmentation for detectron2 models<jupyter_code>#export
from drone_detector.imports import *
from drone_detector.utils import *

import detectron2
from detectron2.data import transforms as T
# export

def build_aug_transforms(cfg:detectron2.config.CfgNode, flip_horiz:bool=True, flip_vert:bool=False, max_rotate:int=10,
                         brightness_limits:Tuple[int,int]=(0.8,1.4), contrast_limits:Tuple[int,int]=(0.8,1.4), 
                         saturation_limits:Tuple[int,int]=(0.8,1.4), p_lighting:float=0.75
                         ) -> detectron2.data.transforms.AugmentationList:
    "Build a list of detectron2 augmentations"
    augs = []
    augs.append(T.ResizeShortestEdge(cfg.INPUT.MIN_SIZE_TRAIN, 
                                     cfg.INPUT.MAX_SIZE_TRAIN, 
                                     cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING))
    if flip_horiz: augs.append(T.RandomFlip(prob=0.5, horizontal=True, vertical=False))
    if flip_vert: augs.append(T.RandomFlip(prob=0.5, horizontal=False, vertical=True))
    if max_rotate: augs.append(T.RandomRotation(angle=[-max_rotate, max_rotate], expand=False))
    if brightness_limits: augs.append(T.RandomApply(prob=p_lighting, tfm_or_aug=T.RandomBrightness(*brightness_limits)))
    if contrast_limits: augs.append(T.RandomApply(prob=p_lighting, tfm_or_aug=T.RandomContrast(*contrast_limits)))
    if saturation_limits: augs.append(T.RandomApply(prob=p_lighting, tfm_or_aug=T.RandomSaturation(*saturation_limits)))
    
    return augs<jupyter_output><empty_output>