├── LICENSE
├── README.md
├── augment.py
├── augment_with_KPs.py
└── doc
├── BeforeAfter.png
├── ImgaugExample.PNG
└── LabelExample1.png
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2020 Evan
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Image-Augmentation-Examples-for-Machine-Learning
2 |
3 |
4 |
5 |
6 |
7 | The code in this repository shows how to use [imgaug](https://github.com/aleju/imgaug) to create thousands of augmented images for training machine learning models. Image augmentation is a quick way to improve accuracy for an image classification or object detection model without having to manually acquire more training images.
8 |
9 | This repository is provided as a reference and example for my talk at the Embedded Vision Summit 2020 conference, [Practical Image Data Augmentation Methods for Training Deep Learning Object Detection Models](https://embeddedvisionsummit.com/2020/session/practical-image-data-augmentation-methods-for-training-deep-learning-object-detection-models/).
10 |
11 |
12 | ## Requirements
13 |
14 | - Python 3.4+
15 | - OpenCV v4.1.0+
16 | - numpy v1.16.4+
17 | - imgaug (`pip install imgaug` or `conda install imgaug` if using an Anaconda environment)
18 |
19 |
20 | ## Usage
21 | ### Basic Augmentation for Image Classification Models
22 | The augment.py script will go through a folder of images and create a specified number of augmented images from each original image.
23 |
24 | These are the input arguments:
25 |
26 | * `--imgdir` specifies the name of the folder with the images to augment
27 |
28 | * `--imgext` specifies the extension of the images to augment (default is .JPG)
29 |
30 | * `--numaugs` specifies the number of augmented images to create from each original image. (default is 5)
31 |
32 | * `--debug` should only be used for test purposes. It causes each augmented image to be dispalyed to preview what the augmentations will look like. (default is False)
33 |
34 | Here's an example for how to run the script. This will create 5 augmented images from each image in the "Squirrels" folder.
35 |
36 | ```
37 | python augment.py --imgdir=Squirrels
38 | ```
39 |
40 | ### Augmentation with Keypoints or Bounding Boxes for Object Detection Models
41 | The augment_with_KPs.py script will go through a folder of images and create a specified number of augmented images from each original image. Each image must have a corresponding annotation data file in [Pascal VOC format](https://gist.github.com/Prasad9/30900b0ef1375cc7385f4d85135fdb44). A new annotation file will also be created for each augmented image.
42 |
43 | There are several input arguments:
44 |
45 | * `--imgdir` specifies the name of the folder with the images to augment
46 |
47 | * `--imgext` specifies the extension of the images to augment (default is .JPG)
48 |
49 | * `--labels` specifies the name of the label file that contains the names of each class for your model. The label file must be a text file listing the classes, with each class seperated by a newline. See the image below for an example label file for a bird, squirrel, and raccoon detection model.
50 |
51 | * `--numaugs` specifies the number of augmented images to create from each original image. (default is 5)
52 |
53 | * `--debug` should only be used for test purposes. It causes each augmented image to be dispalyed to preview what the augmentations will look like. (default is False)
54 |
55 | Here's an example for how to run the script. In this example, I am pointing the script at a folder called "Squirrels" and creating 4 augmented images from each original image.
56 |
57 | ```
58 | python augment_with_KPs.py --imgdir=Squirrels --labels=BSR_labels.txt --numaugs=4
59 | ```
60 |
61 | Here's a picture of the image folder before and after running augmentation.
62 |
63 |
64 |
65 |
66 |
67 |
68 | As another example, here is how to point the script at a folder called "Test", indicate that the image extension is .png, create 5 augmented images per original image, and display every augmented image as they are being augmented.
69 |
70 | ```
71 | python augment_with_KPs.py --imgdir=Test --imgext=.png --labels=BSR_labels.txt --debug=True
72 | ```
73 |
74 | Here is an example of a label file.
75 |
76 |
77 |
78 |
79 |
80 |
81 |
--------------------------------------------------------------------------------
/augment.py:
--------------------------------------------------------------------------------
1 | #### Image Augmentation
2 | #
3 | # Author: Evan Juras
4 | # Date: 5/15/20
5 | # Description:
6 | # This program takes a set of original images and creates augmented images that can be used
7 | # as additional training data for an image classification model. It loads each original image,
8 | # creates N randomly augmented images, and saves the new images.
9 |
10 | import imgaug as ia
11 | from imgaug import augmenters as iaa
12 | import os
13 | import sys
14 | import argparse
15 | from glob import glob
16 | import cv2
17 |
18 |
19 | ## Define control variables and parse user inputs
20 | parser = argparse.ArgumentParser()
21 | parser.add_argument('--imgdir', help='Folder containing images to augment',
22 | required=True)
23 | parser.add_argument('--imgext', help='File extension of images (for example, .JPG)',
24 | default='.JPG')
25 | parser.add_argument('--numaugs', help='Number of augmented images to create from each original image',
26 | default=5)
27 | parser.add_argument('--debug', help='Displays every augmented image when enabled',
28 | default=False)
29 |
30 | args = parser.parse_args()
31 | IMG_DIR = args.imgdir
32 | if not os.path.isdir(IMG_DIR):
33 | print('%s is not a valid directory.' % IMG_DIR)
34 | sys.exit(1)
35 | IMG_EXTENSION = args.imgext
36 | NUM_AUG_IMAGES = int(args.numaugs)
37 | debug = bool(args.debug)
38 | cwd = os.getcwd()
39 |
40 | #### Define augmentation sequence ####
41 | # This can be tweaked to create a huge variety of image augmentations.
42 | # See https://github.com/aleju/imgaug for a list of augmentation techniques available.
43 | seq1 = iaa.Sequential([
44 | iaa.Fliplr(0.5), # Horizontal flip 50% of images
45 | iaa.Crop(percent=(0, 0.10)), # Crop all images between 0% to 10%
46 | iaa.GaussianBlur(sigma=(0, 1)), # Add slight blur to images
47 | ## iaa.Multiply((0.7, 1.3), per_channel=0.2), # Slightly brighten, darken, or recolor images
48 | ## iaa.Affine(
49 | ## scale={"x": (0.8, 1.2), "y": (0.8,1.2)}, # Resize image
50 | ## translate_percent={"x": (-0.1, 0.1), "y": (-0.1, 0.1)}, # Translate image
51 | ## rotate=(-5, 5), # Rotate image
52 | ## mode=ia.ALL, cval=(0,255) # Filling in extra pixels
53 | ## )
54 | ])
55 |
56 |
57 | #### Start of main program ####
58 |
59 | # Obtain list of images in IMG_DIR directory
60 | img_fns = glob(IMG_DIR + '/*' + IMG_EXTENSION)
61 |
62 | # Go through every image in directory, augment it, and save new image/annotation data
63 | for img_fn in img_fns:
64 |
65 | #---- Load image ----#
66 | img1_bgr = cv2.imread(img_fn) # Load image with OpenCV
67 | img1 = cv2.cvtColor(img1_bgr, cv2.COLOR_BGR2RGB) # Re-color to RGB from BGR
68 |
69 | #---- Augment image N times----#
70 | for i in range(NUM_AUG_IMAGES):
71 | img_aug1 = seq1(images=[img1])[0] # Apply augmentation to image
72 |
73 | #---- Save image ----#
74 | base_fn = img_fn.replace(IMG_EXTENSION,'') # Get image base filename
75 | img_aug_fn = base_fn + ('_aug%d' % (i+1)) + IMG_EXTENSION # Append "aug#" to filename
76 | img_aug_bgr1 = cv2.cvtColor(img_aug1, cv2.COLOR_RGB2BGR) # Re-color to BGR from RGB
77 | cv2.imwrite(img_aug_fn,img_aug_bgr1) # Save image to disk
78 |
79 | # Display original and augmented images (if debug is enabled)
80 | if debug:
81 | cv2.imshow('Original image',img1_bgr)
82 | cv2.imshow('Augmented image',img_aug_bgr1)
83 | cv2.waitKey()
84 |
85 | # If images were displayed, close them at the end of the program
86 | if debug:
87 | cv2.destroyAllWindows()
88 |
--------------------------------------------------------------------------------
/augment_with_KPs.py:
--------------------------------------------------------------------------------
1 | #### Image Augmentation for Object Detection Models
2 | #
3 | # Author: Evan Juras
4 | # Date: 5/15/20
5 | # Description:
6 | # This program takes a set of original images and creates augmented images that can be used
7 | # as additional training data for an object detection model. It loads each original image,
8 | # creates N randomly augmented images, and saves the new images and annotation data. It keeps
9 | # track of bounding boxes that are needed for training object detection models.
10 |
11 | import imgaug as ia
12 | from imgaug import augmenters as iaa
13 | import os
14 | import sys
15 | import argparse
16 | from glob import glob
17 | import xml.etree.ElementTree as ET
18 | import cv2
19 | import numpy as np
20 |
21 |
22 | #### Define control variables and parse user inputs
23 | parser = argparse.ArgumentParser()
24 | parser.add_argument('--imgdir', help='Folder containing images to augment',
25 | required=True)
26 | parser.add_argument('--imgext', help='File extension of images (for example, .JPG)',
27 | default='.JPG')
28 | parser.add_argument('--labels', help='Text file with list of classes',
29 | required=True)
30 | parser.add_argument('--numaugs', help='Number of augmented images to create from each original image',
31 | default=5)
32 | parser.add_argument('--debug', help='Displays every augmented image when enabled',
33 | default=False)
34 |
35 | args = parser.parse_args()
36 | IMG_DIR = args.imgdir
37 | if not os.path.isdir(IMG_DIR):
38 | print('%s is not a valid directory.' % IMG_DIR)
39 | sys.exit(1)
40 | IMG_EXTENSION = args.imgext
41 | LABEL_FN = args.labels
42 | if not os.path.isfile(LABEL_FN):
43 | print('%s is not a valid file path.' % LABEL_FN)
44 | sys.exit(1)
45 | NUM_AUG_IMAGES = int(args.numaugs)
46 | debug = args.debug
47 | cwd = os.getcwd()
48 |
49 |
50 | #### Define augmentation sequence ####
51 | # This can be tweaked to create a huge variety of image augmentations.
52 | # See https://github.com/aleju/imgaug for a list of augmentation techniques available.
53 | seq1 = iaa.Sequential([
54 | iaa.Fliplr(0.5), # Horizontal flip 50% of images
55 | iaa.Crop(percent=(0, 0.20)), # Crop all images between 0% to 20%
56 | #iaa.GaussianBlur(sigma=(0, 1)), # Add slight blur to images
57 | #iaa.Multiply((0.7, 1.3), per_channel=0.2)), # Slightly brighten, darken, or recolor images
58 | ## iaa.Affine(
59 | ## scale={"x": (0.8, 1.2), "y": (0.8,1.2)}, # Resize image
60 | ## translate_percent={"x": (-0.1, 0.1), "y": (-0.1, 0.1)}, # Translate image
61 | ## rotate=(-5, 5), # Rotate image
62 | ## mode=ia.ALL, cval=(0,255) # Filling in extra pixels
63 | ## )
64 | ])
65 |
66 | #### Function definitions ####
67 | # Function for reading annotation data from a XML file
68 | def read_annotation_data(xml_fn):
69 | file = open(xml_fn,'r')
70 | tree = ET.parse(file)
71 | root = tree.getroot()
72 | size = root.find('size')
73 | imw = int(size.find('width').text)
74 | imh = int(size.find('height').text)
75 | objects = []
76 |
77 | for obj in root.iter('object'):
78 | difficult = obj.find('difficult').text
79 | cls = obj.find('name').text
80 | if cls not in classes or int(difficult) == 1:
81 | continue
82 | cls_id = classes.index(cls)
83 | xmlbox = obj.find('bndbox')
84 | xmin = int(xmlbox.find('xmin').text)
85 | xmax = int(xmlbox.find('xmax').text)
86 | ymin = int(xmlbox.find('ymin').text)
87 | ymax = int(xmlbox.find('ymax').text)
88 | bb = [(xmin,ymin),(xmax,ymin),(xmax, ymax),(xmin,ymax)] # Top left, top right, bottom right, bottom left
89 | objects.append([cls,bb])
90 |
91 | return imw, imh, objects
92 |
93 |
94 | # Function for finding bounding box from keypoints
95 | def kps_to_BB(kps,imgW,imgH):
96 | """
97 | Determine imgaug bounding box from imgaug keypoints
98 | """
99 | extend=1 # To make the bounding box a little bit bigger
100 | kpsx=[kp[0] for kp in kps]
101 | xmin=max(0,int(min(kpsx)-extend))
102 | xmax=min(imgW,int(max(kpsx)+extend))
103 | kpsy=[kp[1] for kp in kps]
104 | ymin=max(0,int(min(kpsy)-extend))
105 | ymax=min(imgH,int(max(kpsy)+extend))
106 | if xmin==xmax or ymin==ymax:
107 | return None
108 | else:
109 | #return ia.BoundingBox(x1=xmin,y1=ymin,x2=xmax,y2=ymax)
110 | return [(xmin, ymin),(xmax, ymax)]
111 |
112 | # Define XML annotation format for creating new XML files
113 | xml_body_1="""
114 | {FOLDER}
115 | {FILENAME}
116 | {PATH}
117 |
118 | Unknown
119 |
120 |
121 | {WIDTH}
122 | {HEIGHT}
123 | 3
124 |
125 | """
126 | xml_object="""
138 | """
139 | xml_body_2="""
140 | """
141 |
142 | # Function to create XML files
143 | def create_xml(folder, image_path, xml_path, size, imBBs):
144 |
145 | # Get image size and filename
146 | imH, imW = size[0], size[1]
147 | image_fn = os.path.split(image_path)[-1]
148 |
149 | # Create XML file and write data
150 | with open(xml_path,'w') as f:
151 | f.write(xml_body_1.format(**{'FOLDER':folder, 'FILENAME':image_fn, 'PATH':image_path,
152 | 'WIDTH':imW, 'HEIGHT':imH}))
153 |
154 | for bbox in imBBs:
155 | f.write(xml_object.format(**{'CLASS':bbox[0], 'XMIN':bbox[1][0], 'YMIN':bbox[1][1],
156 | 'XMAX':bbox[1][2], 'YMAX':bbox[1][3]}))
157 |
158 | f.write(xml_body_2)
159 |
160 | return
161 |
162 |
163 | #### Start of main program ####
164 |
165 | # Load classes from labelmap. Labelmap is a text file listing class names, with a newline between each class
166 | with open(LABEL_FN,'r') as file:
167 | labels=file.read().split('\n')
168 | classes = [label for label in labels if label!=''] # Remove blank labels (happens when there are extra lines at the end of the file)
169 |
170 | # Obtain list of images in IMG_DIR directory
171 | img_fns = glob(IMG_DIR + '/*' + IMG_EXTENSION)
172 |
173 | # Go through every image in directory, augment it, and save new image/annotation data
174 | for img_fn in img_fns:
175 |
176 | # Open image, get shape and base filename
177 | original_img = cv2.imread(img_fn)
178 | imgH, imgW, _ = original_img.shape
179 | base_img_fn = os.path.split(img_fn)[-1]
180 | base_fn = base_img_fn.replace(IMG_EXTENSION,'')
181 |
182 | # Read annotation data from image's corresponding XML file
183 | xml_fn = img_fn.replace(IMG_EXTENSION,'.xml')
184 | imgW_xml, imgH_xml, objects = read_annotation_data(xml_fn)
185 | if ((imgW_xml != imgW) or (imgH_xml != imgH)):
186 | print('Warning! Annotation data does not match image data for %s. Skipping image.' % img_fn)
187 | continue
188 | im_kps = []
189 | im_classes = []
190 | num_obj = len(objects)
191 | for obj in objects:
192 | im_classes.append(obj[0])
193 | im_kps.append(obj[1][0]) # Top left corner
194 | im_kps.append(obj[1][1]) # Top right corner
195 | im_kps.append(obj[1][2]) # Bottom right corner
196 | im_kps.append(obj[1][3]) # Bottom left corner
197 |
198 | # Define keypoints on image
199 | ia_kps = [ia.Keypoint(x=p[0], y=p[1]) for p in im_kps]
200 | original_kps = ia.KeypointsOnImage(ia_kps, shape=original_img.shape)
201 |
202 | # Define bounding boxes on original image (not needed for augmentation, but used for displaying later if debug == True)
203 | bboxes = []
204 | for i in range(num_obj):
205 | obj_kps = im_kps[i*4:(i*4+4)]
206 | bboxes.append(kps_to_BB(obj_kps,imgW,imgH))
207 |
208 | # Create new augmented images, and save them in folder with annotation data
209 | for n in range(NUM_AUG_IMAGES):
210 |
211 | # Define new filenames
212 | img_aug_fn = base_fn + ('_aug%d' % (n+1)) + IMG_EXTENSION
213 | img_aug_path = os.path.join(cwd,IMG_DIR,img_aug_fn)
214 | xml_aug_fn = img_aug_fn.replace(IMG_EXTENSION,'.xml')
215 | xml_aug_path = os.path.join(cwd,IMG_DIR,xml_aug_fn)
216 |
217 | # Augment image and keypoints.
218 | # First, copy original image and keypoints
219 | img = np.copy(original_img)
220 | kps = original_kps
221 | # Next, need to make sequence determinstic so it performs the same augmentation on image as it does on the keypoints
222 | seq1_det = seq1.to_deterministic()
223 | # Finally, run the image and keypoints through the augmentation sequence
224 | img_aug = seq1_det.augment_images([img])[0]
225 | kps_aug = seq1_det.augment_keypoints([kps])[0]
226 | imgH_aug, imgW_aug, _ = img_aug.shape
227 |
228 | # Extract augmented keypoints back into a list array, find BBs, and write annotation data to new file
229 | list_kps_aug = [(int(kp.x), int(kp.y)) for kp in kps_aug.keypoints]
230 | bboxes_aug = []
231 | bboxes_aug_data = []
232 |
233 | # Loop over every object, determine bounding boxes for new KPs, and save annotation data
234 | for i in range(num_obj):
235 | obj_aug_kps = list_kps_aug[i*4:(i*4+4)] # Augmented keypoints for each object
236 | obj_bb = kps_to_BB(obj_aug_kps,imgW_aug,imgH_aug) # Augmented bounding boxes for each object
237 | if obj_bb: # Sometimes the bbox coordinates are invalid and obj_bb is empty, so need to check if obj_bb valid
238 | bboxes_aug.append(obj_bb) # List of bounding boxes for each object
239 | xmin = int(obj_bb[0][0])
240 | ymin = int(obj_bb[0][1])
241 | xmax = int(obj_bb[1][0])
242 | ymax = int(obj_bb[1][1])
243 | coords = [xmin, ymin, xmax, ymax]
244 | label = im_classes[i]
245 | bboxes_aug_data.append([label, coords]) # List of bounding box data for each object (class name and box coordinates)
246 |
247 | # Save image and XML files to hard disk
248 | cv2.imwrite(img_aug_path,img_aug)
249 | create_xml(IMG_DIR, img_aug_fn, xml_aug_path, [imgH_aug,imgW_aug], bboxes_aug_data)
250 |
251 | # Display original and augmented images with keypoints (if debug is enabled)
252 | if debug:
253 |
254 | img_show = np.copy(img)
255 | img_aug_show = np.copy(img_aug)
256 |
257 | # Draw keypoints and BBs on normal image
258 | for bb in bboxes:
259 | cv2.rectangle(img_show, bb[0], bb[1], (50,255,50), 5)
260 | for kp in im_kps:
261 | cv2.circle(img_show,kp, 8, (255,255,0), -1)
262 | cv2.circle(img_show,kp, 9, (10,10,10), 2)
263 |
264 | # Draw keypoints and BBs on augmented image
265 | for bb in bboxes_aug:
266 | cv2.rectangle(img_aug_show, bb[0], bb[1], (50,255,50), 5)
267 | for kp in list_kps_aug:
268 | cv2.circle(img_aug_show,kp, 10, (255,255,0), -1)
269 | cv2.circle(img_aug_show,kp, 10, (10,10,10), 2)
270 |
271 | if n == 0:
272 | cv2.imshow('Normal',img_show)
273 | cv2.imshow('Augmented %d' % n,img_aug_show)
274 | cv2.waitKey()
275 |
276 | # If images were displayed, close them at the end of the program
277 | if debug:
278 | cv2.destroyAllWindows()
279 |
--------------------------------------------------------------------------------
/doc/BeforeAfter.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/EdjeElectronics/Image-Augmentation-Examples-for-Machine-Learning/2763a0d990016de878cc5f6ff13f4e7fcb2e408c/doc/BeforeAfter.png
--------------------------------------------------------------------------------
/doc/ImgaugExample.PNG:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/EdjeElectronics/Image-Augmentation-Examples-for-Machine-Learning/2763a0d990016de878cc5f6ff13f4e7fcb2e408c/doc/ImgaugExample.PNG
--------------------------------------------------------------------------------
/doc/LabelExample1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/EdjeElectronics/Image-Augmentation-Examples-for-Machine-Learning/2763a0d990016de878cc5f6ff13f4e7fcb2e408c/doc/LabelExample1.png
--------------------------------------------------------------------------------
5 | 
65 | 
78 |