Plug and play, semantic segmentation data enhancement

At present, there are many libraries for data enhancement, but for different purposes, the methods of data enhancement are different. This time, I wrote a semantic segmentation data enhancement method.

Typically, data augmentation methods include cropping, scaling, flipping, adding noise, modifying contrast, and so on. This time the method is used for random cropping and scaling stitching.

data preparation

Data in VOC format

Under JPEGImages is the original image, in jpg format, and under SegmentationClass is the label, in png format.

 

 I wrote a data enhancement method before, which is used for random splicing, and it is not perfect. This time it is an update to the previous method. The current pictures are randomly selected from the results of the last data enhancement. It can be seen that they are simply zoomed and stitched. There is not much processing, and you can also use pictures and tags that have not been data-enhanced, as long as the format of the tags is VOC format. If it is a label marked by labelme, just check the script on the Internet to convert it into VOC format.

the code

def random_concat_images_and_masks(img_folder, mask_folder,num):

    import os
    import random
    from PIL import Image

    img_names = os.listdir(img_folder)
    selected_img_names = random.sample(img_names, num)

    base_img_name = selected_img_names[0]
    base_img = Image.open(os.path.join(img_folder, base_img_name))
    base_mask_name = os.path.splitext(base_img_name)[0] + '.png'
    base_mask = Image.open(os.path.join(mask_folder, base_mask_name))

    for img_name in selected_img_names[1:]:
        img = Image.open(os.path.join(img_folder, img_name))
        mask_name = os.path.splitext(img_name)[0] + '.png'
        mask = Image.open(os.path.join(mask_folder, mask_name))

        scale = random.uniform(0.5, 1.5)
        img = img.resize((int(img.width * scale), int(img.height * scale)))
        mask = mask.resize((int(mask.width * scale), int(mask.height * scale)))

        base_img.paste(img, (0, 0))
        base_mask.paste(mask, (0, 0))

    return base_img, base_mask

Directly on the code, this function randomly scales the image and its label, then crops it, and stitches it randomly.

How to use:

base_img,base_mask = random_concat_images_and_masks(img_folder,mask_folder,num)

parameter:

img_folder is the path of the image, 
mask_folder is the path of the label, and 
num is randomly selected images.

The role of the method:

This method will randomly select num images and their corresponding labels in the data set, perform random scaling, cropping, splicing, and return the images and labels

demo：

if __name__ == '__main__':
    img_folder = 'E:/project3/unet-pytorch-main/datasets/JPEGImages'
    mask_folder = 'E:/project3/unet-pytorch-main/datasets/SegmentationClass'
    img_aug_folder = 'E:/project3/unet-pytorch-main/datasets/JPEGImages_augmentation'
    mask_aug_folder = 'E:/project3/unet-pytorch-main/datasets/SegmentationClass_augmentation'

    for i in range(3):
        base_img,base_mask = random_concat_images_and_masks(img_folder,mask_folder,3)
        base_img.save(os.path.join(img_aug_folder,f'aug{i}.jpg'))
        base_mask.save(os.path.join(mask_aug_folder,f'aug{i}.png'))

 Cycle three times, and select three random pictures in the data set each time. and save to a new folder

Result: (I cycled three times and generated three cards, here I use one as an example.

 

 

 

 The code is a bit unrestrained hahaha, these small scripts are used by myself, don't be surprised.