def get_augmentation()

in src/sagemaker_defect_detection/transforms.py [0:0]

• 21 lines of code
• 2 McCabe index (conditional complexity)

def get_augmentation(split: str) -> Callable:
    """
    Obtains proper image augmentation in train split for detection task.
    We have splitted transformations done for detection task into augmentation and preprocessing
    for clarity

    Parameters
    ----------
    split : str
        train or else

    Returns
    -------
    Callable
        Image augmentation function
    """
    if split == "train":
        return albu.Compose(
            [
                albu.Resize(IMAGE_RESIZE_HEIGHT, IMAGE_RESIZE_WIDTH, always_apply=True),
                albu.RandomCrop(IMAGE_HEIGHT, IMAGE_WIDTH, always_apply=True),
                albu.RandomRotate90(p=PROBABILITY),
                albu.HorizontalFlip(p=PROBABILITY),
                albu.RandomBrightness(p=PROBABILITY),
            ],
            bbox_params=albu.BboxParams(
                format="pascal_voc",
                label_fields=["labels"],
                min_visibility=0.2,
            ),
        )
    else:
        return albu.Compose(
            [albu.Resize(IMAGE_HEIGHT, IMAGE_WIDTH)],
            bbox_params=albu.BboxParams(format="pascal_voc", label_fields=["labels"]),
        )