A fast method of data set enhancement for deep learning

When we prepare data sets before deep learning training , In particular, marked data sets , It takes a lot of manpower and material resources to label , After that, we can also directly enhance the data set , namely Dataset enhancements .

One 、 Common data set enhancement methods

frequently-used Data to enhance There are methods ：

1. translation （Shift） Transformation ： For the original image in the image plane in some way （ Determine the step size of translation in a predefined or random way 、 Scope and direction ） Translate .

2. Flip （Flip） Transformation ： Flip the original image vertically or horizontally .

3. Random cutting （Random Crop）： Randomly define the region of interest to crop the image , It is equivalent to increasing random disturbance .

4. Noise disturbance （Noise）： Randomly add Gaussian noise or salt and pepper noise to the image .

5. Contrast conversion （Contrast）： Change the image contrast , Equivalent to the HSV In the space , Maintain tone components H unchanged , And change the brightness component V And saturation S, Used to simulate the illumination change of real environment .

6. Zoom transform （Zoom）： Reduce or enlarge the image at the set scale .

7. Scale transformation （Scale）： It is similar to scaling transformation , However, the object of scale transformation is the image content rather than the image itself （ You can refer to SIFT Feature extraction method ）, Build image pyramids to get different sizes 、 Blurred image .

Two 、 Use the five modification points of the code

These methods can be quickly enhanced with a piece of code , You only need to change five places in the code to use it directly , All in main In the function ：

（1）IMG_DIR ： The folder path of the original dataset pictures

（2）XML_DIR： original xml File folder path

（3）AUG_XML_DIR ： Data enhanced image saving path

（4）AUG_IMG_DIR： Data enhanced xml Save path of file

（5）AUGLOOP ： How many times each picture is enhanced （ What I set myself is 5）

The specific modification location in the code is shown below ：

notes ： Before using, you need to confirm the original picture and xml Is the file in the folder a A corresponding , For example, there is a picture folder abc.jpg, be xml There must be a folder abc.xml file . The tutorial I provide here is mainly used to enhance VOC Data set in format .

3、 ... and 、 Code

See below for the specific code ：

import xml.etree.ElementTree as ET
import pickle
import os
from os import getcwd
import numpy as np
from PIL import Image
import shutil
import matplotlib.pyplot as plt

import imgaug as ia
from imgaug import augmenters as iaa


ia.seed(1)


def read_xml_annotation(root, image_id):
    in_file = open(os.path.join(root, image_id))
    tree = ET.parse(in_file)
    root = tree.getroot()
    bndboxlist = []

    for object in root.findall('object'):  #  find root All under node country node 
        bndbox = object.find('bndbox')  #  Sub nodes, sub nodes rank Value 

        xmin = int(bndbox.find('xmin').text)
        xmax = int(bndbox.find('xmax').text)
        ymin = int(bndbox.find('ymin').text)
        ymax = int(bndbox.find('ymax').text)
        # print(xmin,ymin,xmax,ymax)
        bndboxlist.append([xmin, ymin, xmax, ymax])
        # print(bndboxlist)

    bndbox = root.find('object').find('bndbox')
    return bndboxlist


# (506.0000, 330.0000, 528.0000, 348.0000) -> (520.4747, 381.5080, 540.5596, 398.6603)
def change_xml_annotation(root, image_id, new_target):
    new_xmin = new_target[0]
    new_ymin = new_target[1]
    new_xmax = new_target[2]
    new_ymax = new_target[3]

    in_file = open(os.path.join(root, str(image_id) + '.xml'))  #  here root It has two meanings 
    tree = ET.parse(in_file)
    xmlroot = tree.getroot()
    object = xmlroot.find('object')
    bndbox = object.find('bndbox')
    xmin = bndbox.find('xmin')
    xmin.text = str(new_xmin)
    ymin = bndbox.find('ymin')
    ymin.text = str(new_ymin)
    xmax = bndbox.find('xmax')
    xmax.text = str(new_xmax)
    ymax = bndbox.find('ymax')
    ymax.text = str(new_ymax)
    tree.write(os.path.join(root, str("%06d" % (str(id) + '.xml'))))


def change_xml_list_annotation(root, image_id, new_target, saveroot, id,img_name):
    in_file = open(os.path.join(root, str(image_id) + '.xml'))  #  here root It has two meanings 
    tree = ET.parse(in_file)
    elem = tree.find('filename')
    elem.text = (img_name + str("_%06d" % int(id)) + '.jpg')
    xmlroot = tree.getroot()
    index = 0

    for object in xmlroot.findall('object'):  #  find root All under node country node 
        bndbox = object.find('bndbox')  #  Sub nodes, sub nodes rank Value 

        # xmin = int(bndbox.find('xmin').text)
        # xmax = int(bndbox.find('xmax').text)
        # ymin = int(bndbox.find('ymin').text)
        # ymax = int(bndbox.find('ymax').text)

        new_xmin = new_target[index][0]
        new_ymin = new_target[index][1]
        new_xmax = new_target[index][2]
        new_ymax = new_target[index][3]

        xmin = bndbox.find('xmin')
        xmin.text = str(new_xmin)
        ymin = bndbox.find('ymin')
        ymin.text = str(new_ymin)
        xmax = bndbox.find('xmax')
        xmax.text = str(new_xmax)
        ymax = bndbox.find('ymax')
        ymax.text = str(new_ymax)

        index = index + 1

    tree.write(os.path.join(saveroot, img_name + str("_%06d" % int(id)) + '.xml'))


def mkdir(path):
    #  Remove the first space 
    path = path.strip()
    #  Remove the tail  \  Symbol 
    path = path.rstrip("\\")
    #  Determine if the path exists 
    #  There is  True
    #  non-existent  False
    isExists = os.path.exists(path)
    #  Judge the result 
    if not isExists:
        #  Create a directory if it doesn't exist 
        #  Create directory manipulation functions 
        os.makedirs(path)
        print(path + '  Create success ')
        return True
    else:
        #  Do not create if directory exists , And prompt that the directory already exists 
        print(path + '  directory already exists ')
        return False


if __name__ == "__main__":

    IMG_DIR = "Images/photo_1"                  ###  The path of the original dataset image 
    XML_DIR = "Images/Annotations"              ###  original xml Path to file 

# =============================================================================
# AUG_XML_DIR = "./Annotations" #  Store enhanced XML Folder path 
# =============================================================================
    AUG_XML_DIR = "Images/Data Enhancement/Annotations Enhancement"              ###  Data enhanced xml Save path of file 
    try:
        shutil.rmtree(AUG_XML_DIR)
    except FileNotFoundError as e:
        a = 1
    mkdir(AUG_XML_DIR)

# =============================================================================
# AUG_IMG_DIR = "./JPEGImages" #  Store the enhanced image folder path 
# =============================================================================
    AUG_IMG_DIR = "Images/Data Enhancement/photo Enhancement"  ###  Save path of image after data enhancement 
    try:
        shutil.rmtree(AUG_IMG_DIR)
    except FileNotFoundError as e:
        a = 1
    mkdir(AUG_IMG_DIR)

    AUGLOOP = 5  #  Number of enhancements per image 

    boxes_img_aug_list = []
    new_bndbox = []
    new_bndbox_list = []

    #  Image enhancement 
    seq = iaa.Sequential([
        iaa.Flipud(0.5),  # vertically flip 20% of all images
        iaa.Fliplr(0.5),  #  Mirror image 
        iaa.Multiply((1.2, 1.5)),  # change brightness, doesn't affect BBs
        iaa.GaussianBlur(sigma=(0, 3.0)),  # iaa.GaussianBlur(0.5),
        iaa.Affine(
            translate_px={
    "x": 15, "y": 15},
            scale=(0.8, 0.95),
            rotate=(-30, 30)
        )  # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs
    ])

    for root, sub_folders, files in os.walk(XML_DIR):

        for name in files:
            print(name)
            bndbox = read_xml_annotation(XML_DIR, name)
            shutil.copy(os.path.join(XML_DIR, name), AUG_XML_DIR)
            shutil.copy(os.path.join(IMG_DIR, name[:-4] + '.jpg'), AUG_IMG_DIR)

            for epoch in range(AUGLOOP):
                seq_det = seq.to_deterministic()  #  Keep the coordinates and images synchronized , Not random 
                #  Read the picture 
                img = Image.open(os.path.join(IMG_DIR, name[:-4] + '.jpg'))
                # sp = img.size
                img = np.asarray(img)
                # bndbox  Coordinate enhancement 
                for i in range(len(bndbox)):
                    bbs = ia.BoundingBoxesOnImage([
                        ia.BoundingBox(x1=bndbox[i][0], y1=bndbox[i][1], x2=bndbox[i][2], y2=bndbox[i][3]),
                    ], shape=img.shape)

                    bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
                    boxes_img_aug_list.append(bbs_aug)

                    # new_bndbox_list:[[x1,y1,x2,y2],...[],[]]
                    n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1)))
                    n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1)))
                    n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2)))
                    n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2)))
                    if n_x1 == 1 and n_x1 == n_x2:
                        n_x2 += 1
                    if n_y1 == 1 and n_y2 == n_y1:
                        n_y2 += 1
                    if n_x1 >= n_x2 or n_y1 >= n_y2:
                        print('error', name)
                    new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
                #  Store the changed picture 
                image_aug = seq_det.augment_images([img])[0]
                path = os.path.join(AUG_IMG_DIR,
                                    name[:-4] + str( "_%06d" % (epoch + 1)) + '.jpg')
                image_auged = bbs.draw_on_image(image_aug, thickness=0)
                Image.fromarray(image_auged).save(path)

                #  Store the changed XML
                change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR,
                                           epoch + 1,name[:-4])
                print( name[:-4] + str( "_%06d" % (epoch + 1)) + '.jpg')
                new_bndbox_list = []

Four 、 Compared with the original data set after enhancement

1、 Raw data set

The original data set and the original corresponding xml file , The original data set contains 26 Zhang image ：

2、 Enhanced dataset

Data set after data enhancement and data set after enhancement xml file , The enhanced pictures are 156 Zhang , Corresponding xml Document has 156 individual ：

The above is the depth of learning VOC Format data set to realize the method of data set enhancement , I hope this method can help you , Support a lot , thank you .