The world’s leading publication for data science, AI, and ML professionals.

Image Augmentation With Domain-Related Artifacts

With Examples Using Hair, Insect & Needle Augmentation

Image by OpticalNomad from Unsplash
Image by OpticalNomad from Unsplash

Introduction

When training deep neural network models on images, augmenting training examples (images) can allow the model to generalize better by training across more images artificially generated by the augmentation. Commonly used augmentations include horizontal and vertical flips/shifts, randomized rotations at an angle and direction (clockwise/counter-clockwise), brightness, saturation, contrast and zoom augmentations.

An extremely popular image augmentation library in Python is albumentations (https://albumentations.ai/), which makes augmenting images easy with their intuitive functions and excellent documentation. It can be used alongside popular deep learning frameworks such as PyTorch and TensorFlow as well.

Domain-Related Artifact Augmentation

Intuitions

The idea behind this approach of augmenting image data with artifacts that can be found in real-life scenarios came from the motivation to mimic and generalize the model across images that it could encounter in reality. For example, augmentations like snow or raindrops are artifacts that should not be found in x-ray images, but chest tubes and pacemakers are artifacts that can be found within x-ray images.

Where The Idea Stemmed From

I first changed upon the approach from Roman’s (@ nroman on Kaggle) hair augmentation for the SIIM-ISIC Melanoma Classification competition. Details of his approach can be found here: https://www.kaggle.com/c/siim-isic-melanoma-classification/discussion/159176. A snippet of the augmentation is as follows:

Original image (above left) and hair augmented image (above right), taken from https://www.kaggle.com/c/siim-isic-melanoma-classification/discussion/159176
Original image (above left) and hair augmented image (above right), taken from https://www.kaggle.com/c/siim-isic-melanoma-classification/discussion/159176

My team did use his augmentations in our model training, which helped improve the cross-validation (CV) scores for most of our models. I’d say this form of augmentations may have played a pivotal role in our final placing! Since then, the idea of using hair (or artifact in general) to augment image data stood close in subsequent competitions that I took part it, and applied it where I could. In particular, the approach was generalized and applied in the Global Wheat Detection, Cassava Leaf Disease Classification and the RANZCR CLiP – Catheter and Line Position Challenge competitions.

Insect Augmentation

As the title suggests, this approach involves augmenting images with insects. The right domain for this artifact can be a nature kind of setting in the data, where insects are commonly found up and about, in the air or on a surface. In this example, bees were used as the insect of choice when augmenting leaf images in the Cassava and Global Wheat detection competitions. The following is an example of how the augmented image will look like:

Augmented image with bees flying around the leaf, taken from https://www.kaggle.com/khoongweihao/insect-augmentation-et-al/notebook
Augmented image with bees flying around the leaf, taken from https://www.kaggle.com/khoongweihao/insect-augmentation-et-al/notebook

We can also, use the masked form of the artifacts, leading to black spots in the image (similar to CoarseDropout in Albumentations), i.e. bees without color and in black:

Augmented image with dark/black bees flying around the leaf, taken from https://www.kaggle.com/khoongweihao/insect-augmentation-et-al/notebook
Augmented image with dark/black bees flying around the leaf, taken from https://www.kaggle.com/khoongweihao/insect-augmentation-et-al/notebook

The following code written in the style of Albumentations allows the artifact augmentation to be used alongside other augmentations from the Albumentations library with ease:

from albumentations.core.transforms_interface import ImageOnlyTransform

class InsectAugmentation(ImageOnlyTransform):
    """
    Impose an image of a insect to the target image
    -----------------------------------------------
    Args:
        insects (int): maximum number of insects to impose
        insects_folder (str): path to the folder with insects images
    """

    def __init__(self, insects=2, dark_insect=False, always_apply=False, p=0.5):
        super().__init__(always_apply, p)
        self.insects = insects
        self.dark_insect = dark_insect
        self.insects_folder = "/kaggle/input/bee-augmentation/"

    def apply(self, image, **kwargs):
        """
        Args:
            image (PIL Image): Image to draw insects on.

        Returns:
            PIL Image: Image with drawn insects.
        """
        n_insects = random.randint(1, self.insects) # for this example I put 1 instead of 0 to illustrate the augmentation

        if not n_insects:
            return image

        height, width, _ = image.shape  # target image width and height
        insects_images = [im for im in os.listdir(self.insects_folder) if 'png' in im]

        for _ in range(n_insects):
            insect = cv2.cvtColor(cv2.imread(os.path.join(self.insects_folder, random.choice(insects_images))), cv2.COLOR_BGR2RGB)
            insect = cv2.flip(insect, random.choice([-1, 0, 1]))
            insect = cv2.rotate(insect, random.choice([0, 1, 2]))

            h_height, h_width, _ = insect.shape  # insect image width and height
            roi_ho = random.randint(0, image.shape[0] - insect.shape[0])
            roi_wo = random.randint(0, image.shape[1] - insect.shape[1])
            roi = image[roi_ho:roi_ho + h_height, roi_wo:roi_wo + h_width]

            # Creating a mask and inverse mask 
            img2gray = cv2.cvtColor(insect, cv2.COLOR_BGR2GRAY)
            ret, mask = cv2.threshold(img2gray, 10, 255, cv2.THRESH_BINARY)
            mask_inv = cv2.bitwise_not(mask)

            # Now black-out the area of insect in ROI
            img_bg = cv2.bitwise_and(roi, roi, mask=mask_inv)

            # Take only region of insect from insect image.
            if self.dark_insect:
                img_bg = cv2.bitwise_and(roi, roi, mask=mask_inv)
                insect_fg = cv2.bitwise_and(img_bg, img_bg, mask=mask)
            else:
                insect_fg = cv2.bitwise_and(insect, insect, mask=mask)

            # Put insect in ROI and modify the target image
            dst = cv2.add(img_bg, insect_fg, dtype=cv2.CV_64F)

            image[roi_ho:roi_ho + h_height, roi_wo:roi_wo + h_width] = dst

        return image

Should you wish to use the black version of the artifact, set dark_insect to True. An example implementation can be found in my Kaggle notebook here: https://www.kaggle.com/khoongweihao/insect-augmentation-with-efficientdet-d6/notebook.

Needle Augmentation

In this approach, needles are used to augment the images, which can be x-ray images or a table top with sewing kits for example. The following is an example of how the augmented image will look like:

Augmented image with needles on the left side of the x-ray, taken from https://www.kaggle.com/khoongweihao/x-ray-needle-augmentation-et-al/notebook
Augmented image with needles on the left side of the x-ray, taken from https://www.kaggle.com/khoongweihao/x-ray-needle-augmentation-et-al/notebook

Similarly, we can use instead the black version of the needle artifacts, leading to the following augmented image:

Augmented image with dark/black needles on the both sides of the x-ray, taken from https://www.kaggle.com/khoongweihao/x-ray-needle-augmentation-et-al/notebook
Augmented image with dark/black needles on the both sides of the x-ray, taken from https://www.kaggle.com/khoongweihao/x-ray-needle-augmentation-et-al/notebook

A snippet of the code that serves as the augmentation module for the above is as follows:

def NeedleAugmentation(image, n_needles=2, dark_needles=False, p=0.5, needle_folder='../input/xray-needle-augmentation'):
    aug_prob = random.random()
    if aug_prob < p:
        height, width, _ = image.shape  # target image width and height
        needle_images = [im for im in os.listdir(needle_folder) if 'png' in im]

        for _ in range(1, n_needles):
            needle = cv2.cvtColor(cv2.imread(os.path.join(needle_folder, random.choice(needle_images))), cv2.COLOR_BGR2RGB)
            needle = cv2.flip(needle, random.choice([-1, 0, 1]))
            needle = cv2.rotate(needle, random.choice([0, 1, 2]))

            h_height, h_width, _ = needle.shape  # needle image width and height
            roi_ho = random.randint(0, abs(image.shape[0] - needle.shape[0]))
            roi_wo = random.randint(0, abs(image.shape[1] - needle.shape[1]))
            roi = image[roi_ho:roi_ho + h_height, roi_wo:roi_wo + h_width]

            # Creating a mask and inverse mask 
            img2gray = cv2.cvtColor(needle, cv2.COLOR_BGR2GRAY)
            ret, mask = cv2.threshold(img2gray, 10, 255, cv2.THRESH_BINARY)
            mask_inv = cv2.bitwise_not(mask)

            # Now black-out the area of needle in ROI
            img_bg = cv2.bitwise_and(roi, roi, mask=mask_inv)

            # Take only region of insect from insect image.
            if dark_needles:
                img_bg = cv2.bitwise_and(roi, roi, mask=mask_inv)
                needle_fg = cv2.bitwise_and(img_bg, img_bg, mask=mask)
            else:
                needle_fg = cv2.bitwise_and(needle, needle, mask=mask)

            # Put needle in ROI and modify the target image
            dst = cv2.add(img_bg, needle_fg, dtype=cv2.CV_64F)

            image[roi_ho:roi_ho + h_height, roi_wo:roi_wo + h_width] = dst

    return image

Note that the above is not in the Albumentations format and cannot be applied directly with the usual Albumentations augmentations. Some tweaking has to be made so that it is in the same format as in the insect/bee augmentation above. But changes should be minor!

Similarly, should you wish to use the black version of the artifact, set dark_needles to True. An example implementation can be found in my Kaggle notebook here: https://www.kaggle.com/khoongweihao/x-ray-needle-augmentation-et-al/notebook.

Experimental Results

In general, local CV results improved, mostly marginally (e.g. 0.001–0.003). But there are cases where using this approach of artifact augmentation ‘fails’ during training. An example can be found in the Global Wheat Detection competition where the task involved detecting wheat heads, i.e. an object detection task. Using bee augmentation with the original colors of the bees resulted in training validation losses being sporadic with huge fluctuations despite extensive hyperparameter tuning. Though using the augmentations did improve CV, it can be said it was indeed a lucky shot. Using artifact augmentations with retention of only black pixels (like coarse dropout augmentation) was proven to be stable across domains of application. In particular, the boost in CV was substantial and consistent as well. As of date it has yet to be found the reason why the bee augmentation led to such sporadic training results between epochs, but a hypothesis is that the bees had a color close to some wheat heads, thus ‘confusing’ the detection algorithm which then captures both the wheat head and the nearest bees within the same bounding box. This was observed in some bounding box predictions, but there were not enough observed cases to say with certainty that the hypothesis was true. In any case, one should also consider whether the image properties (color) of the artifact has a distribution close to the target (e.g. wheat head) as well. Needle augmentation on the other hand, proved to be relatively stable for both kinds of augmentation (original artifact and its black/dark version). In that example, it is possible that the target of the prediction though similar in color distribution, has a distinct characteristic (e.g. chest tubes look vastly different from tiny needles) so that the classification algorithm did not get confused as to whether the needle was the right target or not.


Related Articles