정우진

data_generator

1 +import os
2 +import PIL
3 +from PIL import Image, ImageEnhance, ImageFilter
4 +import random
5 +import cv2, argparse
6 +import numpy as np
7 +from xml.etree.ElementTree import *
8 +from pascal_voc_writer import Writer
9 +
10 +
11 +def image_overlay(src, color="#FFFFFF", alpha=0.5):
12 + overlay = Image.new('RGBA', src.size, color)
13 + bw_src = ImageEnhance.Color(src).enhance(0.0)
14 + return Image.blend(bw_src, overlay, alpha)
15 +
16 +def insert_black_mask(img):
17 + black_mask = Image.new('RGBA',img.size, (0, 0, 0))
18 + val = random.randint(100, 150)
19 + black_mask.putalpha(val)
20 +
21 + img.paste(black_mask, (0, 0,img.size[0], img.size[1]), mask = black_mask)
22 + return img
23 +
24 +
25 +def random_bright_contrast(img):
26 + img = ImageEnhance.Contrast(img)
27 + num_contrast = random.uniform(0.7, 1.5)
28 + img = img.enhance(num_contrast)
29 + num_brightness = random.uniform(0.8, 1.0)
30 + img = ImageEnhance.Brightness(img)
31 + img = img.enhance(num_brightness)
32 +
33 + if random.random() < 0.1:
34 + img = img.filter(ImageFilter.BLUR)
35 +
36 + return img
37 +
38 +
39 +
40 +def image_filtering(img, ang_range=1.2, shear_range=1.5, trans_range=1):
41 +
42 + img = np.array(img)
43 +
44 + # Rotation
45 + ang_rot = np.random.uniform(ang_range) - ang_range / 2
46 + rows, cols, ch = img.shape
47 + Rot_M = cv2.getRotationMatrix2D((cols / 2, rows / 2), ang_rot, 0.9)
48 +
49 + # Translation
50 + tr_x = trans_range * np.random.uniform() - trans_range / 2
51 + tr_y = trans_range * np.random.uniform() - trans_range / 2
52 + Trans_M = np.float32([[1, 0, tr_x], [0, 1, tr_y]])
53 +
54 + # Shear
55 + pts1 = np.float32([[5, 5], [20, 5], [5, 20]])
56 +
57 + pt1 = 5 + shear_range * np.random.uniform() - shear_range / 2
58 + pt2 = 20 + shear_range * np.random.uniform() - shear_range / 2
59 + pts2 = np.float32([[pt1, 5], [pt2, pt1], [5, pt2]])
60 + shear_M = cv2.getAffineTransform(pts1, pts2)
61 +
62 + img = cv2.warpAffine(img, Rot_M, (cols, rows))
63 + img = cv2.warpAffine(img, Trans_M, (cols, rows))
64 + img = cv2.warpAffine(img, shear_M, (cols, rows))
65 +
66 + img = Image.fromarray(img)
67 +
68 + return img
69 +
70 +def image_augmentation(img, car_img, car_boxes):
71 + bbox = None
72 + car_box_1 = car_boxes[0]
73 + car_box_2 = car_boxes[1]
74 +
75 +
76 + width_of_plate_holder = car_box_1[2] - car_box_1[0]
77 + wpercent = (width_of_plate_holder / float(img.size[0]))
78 + hsize = int((float(img.size[1]) * float(wpercent)))
79 + img = img.resize((width_of_plate_holder, hsize), PIL.Image.ANTIALIAS)
80 +
81 + img = np.array(img)
82 + img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
83 + w, h, _ = img.shape
84 + pts1 = np.float32([[0, 0], [0, w], [h, 0], [h, w]])
85 +
86 +
87 + h_change = abs(car_box_1[1] - car_box_2[1])
88 + if (abs(car_box_1[0] - car_box_2[0])) < (width_of_plate_holder / 2):
89 + pts2 = np.float32([[0, h_change], [0, w], [h, 0], [h, w - h_change]])
90 + elif (abs(car_box_1[0] - car_box_2[0])) >(width_of_plate_holder / 2):
91 + pts2 = np.float32([[0, 0], [0, w - h_change], [h, h_change], [h, w ]])
92 + else:
93 + pts2 = np.float32([[0, 0], [0, w], [h, 0], [h, w]])
94 +
95 + M = cv2.getPerspectiveTransform(pts1, pts2)
96 + img = cv2.warpPerspective(img, M, (h, w), flags = cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT, borderValue = [0, 0, 0, 0])
97 + img = Image.fromarray(img)
98 +
99 + width, height = img.size
100 + width_prop = car_boxes[0][0]
101 + height_prop = car_boxes[0][1]
102 + car_img.paste(img, (width_prop, height_prop, width + width_prop, height + height_prop), mask = img)
103 +
104 + bbox = [width_prop, height_prop, width + width_prop, height + height_prop]
105 +
106 + return car_img, bbox
107 +
108 +
109 +class ImageGenerator:
110 + def __init__(self, components_path, detection_images_path,
111 + detection_annotations_path, recognition_data_path, cars_path, background_path):
112 +
113 + # paths inintialization
114 + self.components_path = components_path
115 + self.detection_images_path = detection_images_path
116 + self.detection_annotations_path = detection_annotations_path
117 + self.recognition_data_path = recognition_data_path
118 + self.cars_path = cars_path
119 + self.background_path = background_path
120 +
121 +
122 +
123 + # plates initializtion
124 + plate_files_path = 'plates_type_'
125 + self.plate_types = ["blank", "1", "2", "3", "4", "5", "6"]
126 + # plate type 1
127 + self.Plate_1_images, self.Plate_1_names = self.init_component(
128 + components_path + plate_files_path + self.plate_types[1] + '/')
129 + # plate type 2
130 + self.Plate_2_images, self.Plate_2_names = self.init_component(
131 + components_path + plate_files_path + self.plate_types[2] + '/')
132 + # plate type 3
133 + self.Plate_3_images, self.Plate_3_names = self.init_component(
134 + components_path + plate_files_path + self.plate_types[3] + '/')
135 + # plate type 4
136 + self.Plate_4_images, self.Plate_4_names = self.init_component(
137 + components_path + plate_files_path + self.plate_types[4] + '/')
138 + # plate type 5
139 + self.Plate_5_images, self.Plate_5_names = self.init_component(
140 + components_path + plate_files_path + self.plate_types[5] + '/')
141 + # plate type 6
142 + self.Plate_6_images, self.Plate_6_names = self.init_component(
143 + components_path + plate_files_path + self.plate_types[6] + '/')
144 +
145 + # components init paths
146 + chars_path = 'char_'
147 + nums_path = 'num_'
148 + regions_path = 'regions_'
149 + # folders dictionary
150 + chars_dic = {"black":"black", "white":"white"}
151 + nums_dic = {"white":"white", "black":"black"}
152 + regions_dic = {"region_black":"black", "region_white":"white", "regions_type_4":"type_4"}
153 +
154 + # black chars
155 + self.black_chars_images, self.black_chars_names = self.init_component(
156 + components_path + chars_path + chars_dic["black"] + '/')
157 + # white chars
158 + self.white_chars_images, self.white_chars_names = self.init_component(
159 + components_path + chars_path + chars_dic["white"] + '/')
160 +
161 + # white numbers
162 + self.white_numbers_images, self.white_numbers_names = self.init_component(
163 + components_path + nums_path + nums_dic["white"] + '/')
164 + # black numbers
165 + self.black_numbers_images, self.black_numbers_names = self.init_component(
166 + components_path + nums_path + nums_dic["black"] + '/')
167 +
168 + # white regions
169 + self.white_regions_images, self.white_regions_names = self.init_component(
170 + components_path + regions_path + regions_dic["region_white"] + '/')
171 + # black regions
172 + self.black_regions_images, self.black_regions_names = self.init_component(
173 + components_path + regions_path + regions_dic["region_black"] + '/')
174 + # type 4 regions
175 + self.type_4_regions_images, self.type_4_regions_names = self.init_component(
176 + components_path + regions_path + regions_dic["regions_type_4"] + '/')
177 +
178 +
179 + self.type_1_components = {"Plate_images":self.Plate_1_images,
180 + "Plate_resize":(520, 110),
181 + "Chars":[self.black_chars_images, self.black_chars_names],
182 + "Numbers":[self.black_numbers_images, self.black_numbers_names],
183 + "Order":[{"type":"Numbers", "place":(50,15), "resize":(50,80)},
184 + {"type":"Numbers", "place":(105,15), "resize":(50,80)},
185 + {"type":"Chars", "place":(160,20), "resize":(50,70)},
186 + {"type":"Numbers", "place":(255,15), "resize":(50,80)},
187 + {"type":"Numbers", "place":(310,15), "resize":(50,80)},
188 + {"type":"Numbers", "place":(365,15), "resize":(50,80)},
189 + {"type":"Numbers", "place":(420,15), "resize":(50,80)}]}
190 +
191 + self.type_2_components = {"Plate_images":self.Plate_2_images,
192 + "Plate_resize":(335, 155),
193 + "Chars":[self.black_chars_images, self.black_chars_names],
194 + "Numbers":[self.black_numbers_images, self.black_numbers_names],
195 + "Order":[{"type":"Numbers", "place":(10,50), "resize":(40,80)},
196 + {"type":"Numbers", "place":(55,50), "resize":(40,80)},
197 + {"type":"Chars", "place":(100,60), "resize":(40,70)},
198 + {"type":"Numbers", "place":(150,50), "resize":(40,80)},
199 + {"type":"Numbers", "place":(195,50), "resize":(40,80)},
200 + {"type":"Numbers", "place":(240,50), "resize":(40,80)},
201 + {"type":"Numbers", "place":(285,50), "resize":(40,80)}]}
202 +
203 + self.type_3_components = {"Plate_images":self.Plate_3_images,
204 + "Plate_resize":(335, 170),
205 + "Chars":[self.black_chars_images, self.black_chars_names],
206 + "Numbers":[self.black_numbers_images, self.black_numbers_names],
207 + "Regions":[self.black_regions_images, self.black_regions_names],
208 + "Order":[{"type":"Regions", "place":(80,10), "resize":(80,50)},
209 + {"type":"Numbers", "place":(180,10), "resize":(30,50)},
210 + {"type":"Numbers", "place":(215,10), "resize":(30,50)},
211 + {"type":"Chars", "place":(15,70), "resize":(60,70)},
212 + {"type":"Numbers", "place":(80,70), "resize":(50,90)},
213 + {"type":"Numbers", "place":(140,70), "resize":(50,90)},
214 + {"type":"Numbers", "place":(200,70), "resize":(50,90)},
215 + {"type":"Numbers", "place":(260,70), "resize":(50,90)}]}
216 +
217 +
218 + self.type_4_components = {"Plate_images":self.Plate_4_images,
219 + "Plate_resize":(520, 110),
220 + "Chars":[self.black_chars_images, self.black_chars_names],
221 + "Numbers":[self.black_numbers_images, self.black_numbers_names],
222 + "Regions":[ self.type_4_regions_images, self.type_4_regions_names],
223 + "Order":[{"type":"Regions", "place":(30,15), "resize":(50,80)},
224 + {"type":"Numbers", "place":(90,15), "resize":(50,80)},
225 + {"type":"Numbers", "place":(145,15), "resize":(50,80)},
226 + {"type":"Chars", "place":(200,20), "resize":(50,70)},
227 + {"type":"Numbers", "place":(270,15), "resize":(50,80)},
228 + {"type":"Numbers", "place":(325,15), "resize":(50,80)},
229 + {"type":"Numbers", "place":(380,15), "resize":(50,80)},
230 + {"type":"Numbers", "place":(435,15), "resize":(50,80)}]}
231 +
232 + self.type_5_components = {"Plate_images":self.Plate_5_images,
233 + "Plate_resize":(440, 220),
234 + "Chars":[self.white_chars_images, self.white_chars_names],
235 + "Numbers":[self.white_numbers_images, self.white_numbers_names],
236 + "Order":[{"type":"Numbers", "place":(115,15), "resize":(65,55)},
237 + {"type":"Numbers", "place":(185,15), "resize":(65,55)},
238 + {"type":"Chars", "place":(255,10), "resize":(65,55)},
239 + {"type":"Numbers", "place":(15,80), "resize":(95,125)},
240 + {"type":"Numbers", "place":(120,80), "resize":(95,125)},
241 + {"type":"Numbers", "place":(225,80), "resize":(95,125)},
242 + {"type":"Numbers", "place":(330,80), "resize":(95,125)}]}
243 +
244 + self.type_6_components = {"Plate_images":self.Plate_6_images,
245 + "Plate_resize":(335, 170),
246 + "Chars":[self.white_chars_images, self.white_chars_names],
247 + "Numbers":[self.white_numbers_images, self.white_numbers_names],
248 + "Regions":[self.white_regions_images, self.white_regions_names],
249 + "Order":[
250 + {"type":"Regions", "place":(90,10), "resize":(80,50)},
251 + {"type":"Numbers", "place":(170,10), "resize":(35,50)},
252 + {"type":"Numbers", "place":(210,10), "resize":(35,50)},
253 + {"type":"Chars", "place":(15,65), "resize":(60,65)},
254 + {"type":"Numbers", "place":(95,65), "resize":(50,95)},
255 + {"type":"Numbers", "place":(150,65), "resize":(50,95)},
256 + {"type":"Numbers", "place":(205,65), "resize":(50,95)},
257 + {"type":"Numbers", "place":(260,65), "resize":(50,95)}]}
258 +
259 +
260 +
261 + self.components_to_types = {"1":self.type_1_components,
262 + "2":self.type_2_components,
263 + "3":self.type_3_components,
264 + "4":self.type_4_components,
265 + "5":self.type_5_components,
266 + "6":self.type_6_components}
267 +
268 + self.cars_images, self.cars_list, self.cars_boxes_dic = self.init_car_data(cars_path)
269 + self.backgrounds, self.background_list = self.init_component(background_path)
270 +
271 +
272 + def init_component(self, files_dir):
273 + # a bit initialization
274 + components = list()
275 + components_list = list()
276 + files_path = os.listdir(files_dir)
277 + files_path = [f for f in os.listdir(files_dir) if f.endswith('.jpg') or f.endswith('.png')]
278 +
279 + for file in files_path:
280 + component_path = files_dir + file
281 + component = Image.open(component_path)
282 + components.append(component)
283 + components_list.append(file)
284 + return components, components_list
285 +
286 + def build_data(self, desired_number, desired_types, desired_dataset, count):
287 +
288 + for iteration in range(desired_number):
289 +
290 + # generate desired plate
291 + plate_type = random.choice(self.plate_types)
292 + while plate_type not in desired_types:
293 + plate_type = random.choice(self.plate_types)
294 +
295 + generated_plate, label = self.build_plate(plate_type)
296 + if desired_dataset == 'recognition':
297 + desired_height = 110
298 +
299 + if random.random() < 0.9:
300 + generated_plate = image_filtering(generated_plate)
301 + self.save_plate(generated_plate, label, desired_height)
302 + try:
303 +
304 + print(str(count) + " / " + str(desired_number))
305 + count += 1
306 +
307 + except:
308 + print("Fuck! Save error, probably same combination")
309 +
310 + continue
311 +
312 +
313 + full_img, bbox = self.generate_full_image(desired_dataset, generated_plate)
314 +
315 + try:
316 + self.save_img_annotation(full_img, bbox, label, plate_type, count)
317 + print(str(count) + " / " + str(desired_number))
318 + count += 1
319 + except:
320 + print('Could not create or save')
321 +
322 + def build_plate(self, plate_type):
323 + generated_plate = None
324 + components = self.components_to_types[plate_type]
325 + plates = components["Plate_images"]
326 + # choose plate
327 + Plate_img = random.choice(plates)
328 + Plate_img = Plate_img.resize(components["Plate_resize"], PIL.Image.ANTIALIAS)
329 + Generated_Plate, Plate_label = self.place_components(Plate_img, components)
330 +
331 + Generated_Plate = Generated_Plate.convert('RGB')
332 + Generated_Plate = random_bright_contrast(Generated_Plate)
333 + return Generated_Plate, Plate_label
334 +
335 + def place_components(self, Plate_img, components):
336 + label = str()
337 + for component in components["Order"]:
338 + index = random.choice(range(len(components[component["type"]][0])))
339 + component_img = components[component["type"]][0][index]
340 +
341 + component_label = components[component["type"]][1][index][:-4]
342 + resize_ratio = component["resize"]
343 + resized_component_img = component_img.resize(resize_ratio, PIL.Image.ANTIALIAS)
344 + xmin,ymin = component["place"]
345 + xmax, ymax = xmin + resized_component_img.size[0], ymin + resized_component_img.size[1]
346 + place_holder = (xmin, ymin, xmax, ymax)
347 + Plate_img.paste(resized_component_img, place_holder, mask = resized_component_img)
348 + label += component_label
349 + return Plate_img, label
350 +
351 + def save_plate(self, plate, label, desired_height):
352 + wpercent = (desired_height / float(plate.size[1]))
353 + hsize = int((float(plate.size[0]) * float(wpercent)))
354 + plate = plate.resize((hsize, desired_height), PIL.Image.ANTIALIAS)
355 +
356 + plate = plate.convert('RGB')
357 + plate = plate.convert('L')
358 +
359 + plate.save(self.recognition_data_path + label + '.jpg')
360 +
361 + def init_car_data(self, cars_path):
362 + cars_xml = [f for f in os.listdir(cars_path) if f.endswith('.xml')]
363 + cars_imgs, cars_list = self.init_component(cars_path)
364 + car_boxes_dic = {}
365 + for car_xml in cars_xml:
366 + car_boxes = []
367 + car_box1 = []
368 + car_box2 = []
369 + node = parse(cars_path + car_xml)
370 + elems = node.findall('object')
371 +
372 + for item in (['xmin', 'ymin', 'xmax', 'ymax']):
373 + car_box1.append(int(int(elems[0].find('bndbox').find(item).text)))
374 + car_box2.append(int(int(elems[1].find('bndbox').find(item).text)))
375 + car_boxes.append(car_box1)
376 + car_boxes.append(car_box2)
377 + car_boxes_dic[car_xml[:-3] + "png"] = car_boxes
378 + return cars_imgs, cars_list, car_boxes_dic
379 + def generate_full_image(self, desired_dataset, plate):
380 +
381 + if random.random() < 0.3:
382 + plate = insert_black_mask(plate)
383 +
384 +
385 + full_img, bbox = None, None
386 + car = random.choice(self.cars_list)
387 + print(car)
388 + car_img = Image.open(self.cars_path + car)
389 + car_boxes = self.cars_boxes_dic[car]
390 + background = random.choice(self.background_list)
391 + background_img = Image.open(self.background_path + background)
392 +
393 + car_with_plate_img, bbox = image_augmentation(plate, car_img, car_boxes)
394 + full_img, bbox = self.insert_car_to_background(background_img, car_with_plate_img, desired_dataset, bbox)
395 +
396 + return full_img, bbox
397 +
398 + def insert_car_to_background(self, background_img, car_with_plate_img, desired_dataset, bbox):
399 +
400 + background_img = background_img.resize((1600, 1200), PIL.Image.ANTIALIAS)
401 + img_width, img_height = background_img.size
402 +
403 + background_car_holder_width = None
404 +
405 + if desired_dataset == 'parking':
406 + background_car_holder_width = random.randint(int(img_width * 0.7), int(img_width * 0.8))
407 + width_prop = random.randint(int(img_width * 0.05), int(img_width * 0.4))
408 + height_prop = random.randint(int(img_height * 0.1), int(img_height * 0.3))
409 + elif desired_dataset == 'cctv':
410 + width_prop = random.randint(int(img_width * 0.1), int(img_width * 0.4))
411 + height_prop = random.randint(int(img_height * 0.1), int(img_height * 0.25))
412 + background_car_holder_width = random.randint(int(img_width * 0.3 * height_prop * 4/1200),
413 + int(img_width * 0.45 * height_prop * 4/1200))
414 + else:
415 + print('Error!')
416 + return
417 +
418 + wpercent = (background_car_holder_width / float(car_with_plate_img.size[0]))
419 + hsize = int((float(car_with_plate_img.size[1]) * float(wpercent)))
420 + resized_car_with_plate_img = car_with_plate_img.resize((background_car_holder_width, hsize), PIL.Image.ANTIALIAS)
421 +
422 + # take bbox
423 + bbox[0] = int(bbox[0] * resized_car_with_plate_img.size[0] / car_with_plate_img.size[0])
424 + bbox[1] = int(bbox[1] * resized_car_with_plate_img.size[1] / car_with_plate_img.size[1])
425 + bbox[2] = int(bbox[2] * resized_car_with_plate_img.size[0] / car_with_plate_img.size[0])
426 + bbox[3] = int(bbox[3] * resized_car_with_plate_img.size[1] / car_with_plate_img.size[1])
427 +
428 + # insert car in background
429 + width, height = resized_car_with_plate_img.size
430 +
431 +
432 + background_img.paste(resized_car_with_plate_img, (width_prop, height_prop, width + width_prop, height + height_prop),
433 + mask = resized_car_with_plate_img)
434 + background_img = background_img.convert('L')
435 +
436 + bbox[0] += width_prop
437 + bbox[1] += height_prop
438 + bbox[2] += width_prop
439 + bbox[3] += height_prop
440 +
441 + #background_img.save('test.jpg')
442 +
443 + return background_img, bbox
444 +
445 + def save_img_annotation(self, full_img, bbox, label, plate_type, count):
446 + saved = False
447 +
448 +
449 + #self.detection_images_path = detection_images_path
450 + #self.detection_annotations_path = detection_annotations_path
451 +
452 +
453 + try:
454 + # annotation file making
455 + file_name = str(count) + '_P' + plate_type + '_img.jpg'
456 + img_save_path = self.detection_images_path + file_name
457 + full_img.save(img_save_path)
458 +
459 +
460 + writer = Writer(img_save_path, full_img.size[0], full_img.size[1])
461 + writer.addObject('P' + plate_type, bbox[0], bbox[1], bbox[2], bbox[3])
462 + writer.save(self.detection_annotations_path + file_name[:-3] +'xml')
463 + saved = True
464 +
465 + except:
466 + print("Could not save")
467 +
468 +
469 + return saved
470 +
471 +components_path = 'components_2/'
472 +detection_images_path = '../dataset/yolo'
473 +detection_annotations_path = 'detection_annotations/'
474 +recognition_data_path = 'recognition_data/'
475 +
476 +# choose right folder for desired type
477 +
478 +# for parking ->
479 +# cars_path = 'cars/'
480 +# background_path = 'background/'
481 +
482 +# for cctv ->
483 +# cars_path = 'cctv_cars/'
484 +# background_path = 'cctv_background/'
485 +
486 +
487 +cars_path = 'cars/'
488 +background_path = 'background/'
489 +
490 +
491 +desired_number = 10000
492 +desired_types = ['1', '2', '3', '4', '5', '6']
493 +# cctv parking recognition provide choose one to generate and store things
494 +desired_dataset = 'parking'
495 +count = 1
496 +
497 +
498 +
499 +generator = ImageGenerator(components_path = components_path,
500 + detection_images_path = detection_images_path,
501 + detection_annotations_path = detection_annotations_path,
502 + recognition_data_path = recognition_data_path,
503 + cars_path = cars_path,
504 + background_path = background_path)
505 +
506 +generator.build_data(desired_number = desired_number,
507 + desired_types = desired_types,
508 + desired_dataset = desired_dataset,
509 + count = count)
...\ No newline at end of file ...\ No newline at end of file
1 +import os
2 +import PIL
3 +from PIL import Image, ImageEnhance, ImageFilter
4 +import random
5 +import cv2, argparse
6 +import pandas as pd
7 +import numpy as np
8 +
9 +# from xml.etree.ElementTree import parse
10 +# from pascal_voc_writer import Writer
11 +
12 +
13 +def image_overlay(src, color="#FFFFFF", alpha=0.5):
14 + overlay = Image.new("RGBA", src.size, color)
15 + bw_src = ImageEnhance.Color(src).enhance(0.0)
16 + return Image.blend(bw_src, overlay, alpha)
17 +
18 +
19 +def insert_black_mask(img):
20 + black_mask = Image.new("RGBA", img.size, (0, 0, 0))
21 + val = random.randint(100, 150)
22 + black_mask.putalpha(val)
23 +
24 + img.paste(black_mask, (0, 0, img.size[0], img.size[1]), mask=black_mask)
25 + return img
26 +
27 +
28 +def random_bright_contrast(img):
29 + img = ImageEnhance.Contrast(img)
30 + num_contrast = random.uniform(0.7, 1.5)
31 + img = img.enhance(num_contrast)
32 + num_brightness = random.uniform(0.8, 1.0)
33 + img = ImageEnhance.Brightness(img)
34 + img = img.enhance(num_brightness)
35 +
36 + if random.random() < 0.1:
37 + img = img.filter(ImageFilter.BLUR)
38 +
39 + return img
40 +
41 +
42 +def image_filtering(img, ang_range=1.2, shear_range=1.5, trans_range=1):
43 +
44 + img = np.array(img)
45 +
46 + # Rotation
47 + ang_rot = np.random.uniform(ang_range) - ang_range / 2
48 + rows, cols, ch = img.shape
49 + Rot_M = cv2.getRotationMatrix2D((cols / 2, rows / 2), ang_rot, 0.9)
50 +
51 + # Translation
52 + tr_x = trans_range * np.random.uniform() - trans_range / 2
53 + tr_y = trans_range * np.random.uniform() - trans_range / 2
54 + Trans_M = np.float32([[1, 0, tr_x], [0, 1, tr_y]])
55 +
56 + # Shear
57 + pts1 = np.float32([[5, 5], [20, 5], [5, 20]])
58 +
59 + pt1 = 5 + shear_range * np.random.uniform() - shear_range / 2
60 + pt2 = 20 + shear_range * np.random.uniform() - shear_range / 2
61 + pts2 = np.float32([[pt1, 5], [pt2, pt1], [5, pt2]])
62 + shear_M = cv2.getAffineTransform(pts1, pts2)
63 +
64 + img = cv2.warpAffine(img, Rot_M, (cols, rows))
65 + img = cv2.warpAffine(img, Trans_M, (cols, rows))
66 + img = cv2.warpAffine(img, shear_M, (cols, rows))
67 +
68 + img = Image.fromarray(img)
69 +
70 + return img
71 +
72 +
73 +def image_augmentation(img, car_img, car_boxes):
74 + bbox = None
75 + car_box_1 = car_boxes[0]
76 + car_box_2 = car_boxes[1]
77 +
78 + width_of_plate_holder = car_box_1[2] - car_box_1[0]
79 + wpercent = width_of_plate_holder / float(img.size[0])
80 + hsize = int((float(img.size[1]) * float(wpercent)))
81 + img = img.resize((width_of_plate_holder, hsize), PIL.Image.ANTIALIAS)
82 +
83 + img = np.array(img)
84 + img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
85 + w, h, _ = img.shape
86 + pts1 = np.float32([[0, 0], [0, w], [h, 0], [h, w]])
87 +
88 + h_change = abs(car_box_1[1] - car_box_2[1])
89 + if (abs(car_box_1[0] - car_box_2[0])) < (width_of_plate_holder / 2):
90 + pts2 = np.float32([[0, h_change], [0, w], [h, 0], [h, w - h_change]])
91 + elif (abs(car_box_1[0] - car_box_2[0])) > (width_of_plate_holder / 2):
92 + pts2 = np.float32([[0, 0], [0, w - h_change], [h, h_change], [h, w]])
93 + else:
94 + pts2 = np.float32([[0, 0], [0, w], [h, 0], [h, w]])
95 +
96 + M = cv2.getPerspectiveTransform(pts1, pts2)
97 + img = cv2.warpPerspective(
98 + img,
99 + M,
100 + (h, w),
101 + flags=cv2.INTER_CUBIC,
102 + borderMode=cv2.BORDER_CONSTANT,
103 + borderValue=[0, 0, 0, 0],
104 + )
105 + img = Image.fromarray(img)
106 +
107 + width, height = img.size
108 + width_prop = car_boxes[0][0]
109 + height_prop = car_boxes[0][1]
110 + car_img.paste(
111 + img,
112 + (width_prop, height_prop, width + width_prop, height + height_prop),
113 + mask=img,
114 + )
115 +
116 + bbox = [width_prop, height_prop, width + width_prop, height + height_prop]
117 +
118 + return car_img, bbox
119 +
120 +
121 +class ImageGenerator:
122 + def __init__(
123 + self,
124 + components_path,
125 + plate_images_path,
126 + detection_annotations_path,
127 + recognition_data_path,
128 + cars_path,
129 + background_path,
130 + ):
131 +
132 + # paths inintialization
133 + self.components_path = components_path
134 + self.plate_images_path = plate_images_path
135 + self.detection_annotations_path = detection_annotations_path
136 + self.recognition_data_path = recognition_data_path
137 + self.cars_path = cars_path
138 + self.background_path = background_path
139 +
140 + # plates initializtion
141 + plate_files_path = "plates_type_"
142 + self.plate_types = ["blank", "1", "2", "3", "4", "5", "6"]
143 + # plate type 1
144 + self.Plate_1_images, self.Plate_1_names = self.init_component(
145 + components_path + plate_files_path + self.plate_types[1] + "/"
146 + )
147 + # plate type 2
148 + self.Plate_2_images, self.Plate_2_names = self.init_component(
149 + components_path + plate_files_path + self.plate_types[2] + "/"
150 + )
151 + # plate type 3
152 + self.Plate_3_images, self.Plate_3_names = self.init_component(
153 + components_path + plate_files_path + self.plate_types[3] + "/"
154 + )
155 + # plate type 4
156 + self.Plate_4_images, self.Plate_4_names = self.init_component(
157 + components_path + plate_files_path + self.plate_types[4] + "/"
158 + )
159 + # plate type 5
160 + self.Plate_5_images, self.Plate_5_names = self.init_component(
161 + components_path + plate_files_path + self.plate_types[5] + "/"
162 + )
163 + # plate type 6
164 + self.Plate_6_images, self.Plate_6_names = self.init_component(
165 + components_path + plate_files_path + self.plate_types[6] + "/"
166 + )
167 +
168 + # components init paths
169 + chars_path = "char_"
170 + nums_path = "num_"
171 + regions_path = "regions_"
172 + # folders dictionary
173 + chars_dic = {"black": "black", "white": "white"}
174 + nums_dic = {"white": "white", "black": "black"}
175 + regions_dic = {
176 + "region_black": "black",
177 + "region_white": "white",
178 + "regions_type_4": "type_4",
179 + }
180 +
181 + # black chars
182 + self.black_chars_images, self.black_chars_names = self.init_component(
183 + components_path + chars_path + chars_dic["black"] + "/"
184 + )
185 + # white chars
186 + self.white_chars_images, self.white_chars_names = self.init_component(
187 + components_path + chars_path + chars_dic["white"] + "/"
188 + )
189 +
190 + # white numbers
191 + self.white_numbers_images, self.white_numbers_names = self.init_component(
192 + components_path + nums_path + nums_dic["white"] + "/"
193 + )
194 + # black numbers
195 + self.black_numbers_images, self.black_numbers_names = self.init_component(
196 + components_path + nums_path + nums_dic["black"] + "/"
197 + )
198 +
199 + # white regions
200 + self.white_regions_images, self.white_regions_names = self.init_component(
201 + components_path + regions_path + regions_dic["region_white"] + "/"
202 + )
203 + # black regions
204 + self.black_regions_images, self.black_regions_names = self.init_component(
205 + components_path + regions_path + regions_dic["region_black"] + "/"
206 + )
207 + # type 4 regions
208 + self.type_4_regions_images, self.type_4_regions_names = self.init_component(
209 + components_path + regions_path + regions_dic["regions_type_4"] + "/"
210 + )
211 +
212 + self.type_1_components = {
213 + "Plate_images": self.Plate_1_images,
214 + "Plate_resize": (520, 110),
215 + "Chars": [self.black_chars_images, self.black_chars_names],
216 + "Numbers": [self.black_numbers_images, self.black_numbers_names],
217 + "Order": [
218 + {"type": "Numbers", "place": (50, 15), "resize": (50, 80)},
219 + {"type": "Numbers", "place": (105, 15), "resize": (50, 80)},
220 + {"type": "Chars", "place": (160, 20), "resize": (50, 70)},
221 + {"type": "Numbers", "place": (255, 15), "resize": (50, 80)},
222 + {"type": "Numbers", "place": (310, 15), "resize": (50, 80)},
223 + {"type": "Numbers", "place": (365, 15), "resize": (50, 80)},
224 + {"type": "Numbers", "place": (420, 15), "resize": (50, 80)},
225 + ],
226 + }
227 +
228 + self.type_2_components = {
229 + "Plate_images": self.Plate_2_images,
230 + "Plate_resize": (335, 155),
231 + "Chars": [self.black_chars_images, self.black_chars_names],
232 + "Numbers": [self.black_numbers_images, self.black_numbers_names],
233 + "Order": [
234 + {"type": "Numbers", "place": (10, 50), "resize": (40, 80)},
235 + {"type": "Numbers", "place": (55, 50), "resize": (40, 80)},
236 + {"type": "Chars", "place": (100, 60), "resize": (40, 70)},
237 + {"type": "Numbers", "place": (150, 50), "resize": (40, 80)},
238 + {"type": "Numbers", "place": (195, 50), "resize": (40, 80)},
239 + {"type": "Numbers", "place": (240, 50), "resize": (40, 80)},
240 + {"type": "Numbers", "place": (285, 50), "resize": (40, 80)},
241 + ],
242 + }
243 +
244 + self.type_3_components = {
245 + "Plate_images": self.Plate_3_images,
246 + "Plate_resize": (335, 170),
247 + "Chars": [self.black_chars_images, self.black_chars_names],
248 + "Numbers": [self.black_numbers_images, self.black_numbers_names],
249 + "Regions": [self.black_regions_images, self.black_regions_names],
250 + "Order": [
251 + {"type": "Regions", "place": (80, 10), "resize": (80, 50)},
252 + {"type": "Numbers", "place": (180, 10), "resize": (30, 50)},
253 + {"type": "Numbers", "place": (215, 10), "resize": (30, 50)},
254 + {"type": "Chars", "place": (15, 70), "resize": (60, 70)},
255 + {"type": "Numbers", "place": (80, 70), "resize": (50, 90)},
256 + {"type": "Numbers", "place": (140, 70), "resize": (50, 90)},
257 + {"type": "Numbers", "place": (200, 70), "resize": (50, 90)},
258 + {"type": "Numbers", "place": (260, 70), "resize": (50, 90)},
259 + ],
260 + }
261 +
262 + self.type_4_components = {
263 + "Plate_images": self.Plate_4_images,
264 + "Plate_resize": (520, 110),
265 + "Chars": [self.black_chars_images, self.black_chars_names],
266 + "Numbers": [self.black_numbers_images, self.black_numbers_names],
267 + "Regions": [self.type_4_regions_images, self.type_4_regions_names],
268 + "Order": [
269 + {"type": "Regions", "place": (30, 15), "resize": (50, 80)},
270 + {"type": "Numbers", "place": (90, 15), "resize": (50, 80)},
271 + {"type": "Numbers", "place": (145, 15), "resize": (50, 80)},
272 + {"type": "Chars", "place": (200, 20), "resize": (50, 70)},
273 + {"type": "Numbers", "place": (270, 15), "resize": (50, 80)},
274 + {"type": "Numbers", "place": (325, 15), "resize": (50, 80)},
275 + {"type": "Numbers", "place": (380, 15), "resize": (50, 80)},
276 + {"type": "Numbers", "place": (435, 15), "resize": (50, 80)},
277 + ],
278 + }
279 +
280 + self.type_5_components = {
281 + "Plate_images": self.Plate_5_images,
282 + "Plate_resize": (440, 220),
283 + "Chars": [self.white_chars_images, self.white_chars_names],
284 + "Numbers": [self.white_numbers_images, self.white_numbers_names],
285 + "Order": [
286 + {"type": "Numbers", "place": (115, 15), "resize": (65, 55)},
287 + {"type": "Numbers", "place": (185, 15), "resize": (65, 55)},
288 + {"type": "Chars", "place": (255, 10), "resize": (65, 55)},
289 + {"type": "Numbers", "place": (15, 80), "resize": (95, 125)},
290 + {"type": "Numbers", "place": (120, 80), "resize": (95, 125)},
291 + {"type": "Numbers", "place": (225, 80), "resize": (95, 125)},
292 + {"type": "Numbers", "place": (330, 80), "resize": (95, 125)},
293 + ],
294 + }
295 +
296 + self.type_6_components = {
297 + "Plate_images": self.Plate_6_images,
298 + "Plate_resize": (335, 170),
299 + "Chars": [self.white_chars_images, self.white_chars_names],
300 + "Numbers": [self.white_numbers_images, self.white_numbers_names],
301 + "Regions": [self.white_regions_images, self.white_regions_names],
302 + "Order": [
303 + {"type": "Regions", "place": (90, 10), "resize": (80, 50)},
304 + {"type": "Numbers", "place": (170, 10), "resize": (35, 50)},
305 + {"type": "Numbers", "place": (210, 10), "resize": (35, 50)},
306 + {"type": "Chars", "place": (15, 65), "resize": (60, 65)},
307 + {"type": "Numbers", "place": (95, 65), "resize": (50, 95)},
308 + {"type": "Numbers", "place": (150, 65), "resize": (50, 95)},
309 + {"type": "Numbers", "place": (205, 65), "resize": (50, 95)},
310 + {"type": "Numbers", "place": (260, 65), "resize": (50, 95)},
311 + ],
312 + }
313 +
314 + self.components_to_types = {
315 + "1": self.type_1_components,
316 + "2": self.type_2_components,
317 + "3": self.type_3_components,
318 + "4": self.type_4_components,
319 + "5": self.type_5_components,
320 + "6": self.type_6_components,
321 + }
322 +
323 + ####################################################################################
324 + # self.cars_images, self.cars_list, self.cars_boxes_dic = self.init_car_data(
325 + # cars_path
326 + # )
327 + # self.backgrounds, self.background_list = self.init_component(background_path)
328 +
329 + def init_component(self, files_dir):
330 + # a bit initialization
331 + components = list()
332 + components_list = list()
333 + files_path = os.listdir(files_dir)
334 + files_path = [
335 + f for f in os.listdir(files_dir) if f.endswith(".jpg") or f.endswith(".png")
336 + ]
337 +
338 + for file in files_path:
339 + component_path = files_dir + file
340 + component = Image.open(component_path)
341 + components.append(component)
342 + components_list.append(file)
343 + return components, components_list
344 +
345 + def build_data(self, desired_number, desired_types, desired_dataset, count):
346 + # def build_data(self, desired_number, desired_types, count):
347 +
348 + for iteration in range(desired_number):
349 +
350 + # generate desired plate
351 + plate_type = random.choice(self.plate_types)
352 + while plate_type not in desired_types:
353 + plate_type = random.choice(self.plate_types)
354 +
355 + generated_plate, label = self.build_plate(plate_type)
356 +
357 + # parking??
358 + # if desired_dataset == "recognition":
359 + # desired_height = 110
360 +
361 + # if random.random() < 0.9:
362 + # generated_plate = image_filtering(generated_plate)
363 + # self.save_plate(generated_plate, label, desired_height)
364 + # try:
365 +
366 + # print(str(count) + " / " + str(desired_number))
367 + # count += 1
368 +
369 + # except:
370 + # print("Fuck! Save error, probably same combination")
371 +
372 + # continue
373 +
374 + ################################################################################
375 + # full_img, bbox = self.generate_full_image(desired_dataset, generated_plate)
376 +
377 + # try:
378 + # # self.save_img_annotation(full_img, bbox, label, plate_type, count)
379 + # self.save_img_annotation(generated_plate, label, plate_type, count)
380 + # print(str(count) + " / " + str(desired_number))
381 + # count += 1
382 +
383 + # except:
384 + # print("Could not create or save")
385 +
386 + self.save_img_annotation(generated_plate, label, plate_type, count)
387 + print(str(count) + " / " + str(desired_number))
388 + count += 1
389 +
390 + def build_plate(self, plate_type):
391 + generated_plate = None
392 + components = self.components_to_types[plate_type]
393 + plates = components["Plate_images"]
394 + # choose plate
395 + Plate_img = random.choice(plates)
396 + Plate_img = Plate_img.resize(components["Plate_resize"], PIL.Image.ANTIALIAS)
397 + Generated_Plate, Plate_label = self.place_components(Plate_img, components)
398 +
399 + Generated_Plate = Generated_Plate.convert("RGB")
400 + Generated_Plate = random_bright_contrast(Generated_Plate)
401 + return Generated_Plate, Plate_label
402 +
403 + def place_components(self, Plate_img, components):
404 + label = str()
405 + for component in components["Order"]:
406 + index = random.choice(range(len(components[component["type"]][0])))
407 + component_img = components[component["type"]][0][index]
408 +
409 + component_label = components[component["type"]][1][index][:-4]
410 + resize_ratio = component["resize"]
411 + resized_component_img = component_img.resize(
412 + resize_ratio, PIL.Image.ANTIALIAS
413 + )
414 + xmin, ymin = component["place"]
415 + xmax, ymax = (
416 + xmin + resized_component_img.size[0],
417 + ymin + resized_component_img.size[1],
418 + )
419 + place_holder = (xmin, ymin, xmax, ymax)
420 + Plate_img.paste(
421 + resized_component_img, place_holder, mask=resized_component_img
422 + )
423 + label += component_label
424 + return Plate_img, label
425 +
426 + # def save_img_annotation(self, full_img, bbox, label, plate_type, count): #
427 + def save_img_annotation(self, plate_img, label, plate_type, count): #
428 + saved = False
429 +
430 + # print("label param : ", label)
431 +
432 + # self.plate_images_path = plate_images_path
433 + # self.detection_annotations_path = detection_annotations_path
434 +
435 + # try:
436 + # # annotation file making
437 + # file_name = str(count) + "_P" + plate_type + "_img.jpg"
438 + # img_save_path = self.plate_images_path + file_name
439 + # plate_img.save(img_save_path)
440 +
441 + # annotations = []
442 + # vectors = self.components_to_types[plate_type]["Order"]
443 + # for v in vectors:
444 + # x1, y1 = v["place"]
445 + # x2 = v["place"][0] + v["resize"][0]
446 + # y2 = v["place"][1] + v["resize"][1]
447 + # annotations.append([img_save_path, x1, y1, x2, y2])
448 + # saved = True
449 +
450 + # annotations = np.array(annotations)
451 + # label = label.split("")
452 + # if plate_type in ["4","5","6"]:
453 + # col = [label[0] + label[1]] + label[2:]
454 + # else:
455 + # col = label
456 +
457 + # print(col)
458 + # np.insert(annotations, -1, col, axis=1)
459 +
460 + # print(annotations)
461 +
462 + # pd.DataFrame(
463 + # annotations, columns=["image path", "x1", "y1", "x2", "y2", "label"]
464 + # )
465 + # pd.to_csv("annotations.csv", mode="a", header=False, index=True)
466 +
467 + # except:
468 + # print("Could not save")
469 +
470 + file_name = str(count) + "_P" + plate_type + "_img.jpg"
471 + img_save_path = self.plate_images_path + file_name
472 + plate_img.save(img_save_path)
473 +
474 + annotations = []
475 + vectors = self.components_to_types[plate_type]["Order"]
476 + width, height = self.components_to_types[plate_type]["Plate_resize"]
477 + for v in vectors:
478 + x1, y1 = v["place"]
479 + x2 = v["place"][0] + v["resize"][0]
480 + y2 = v["place"][1] + v["resize"][1]
481 + annotations.append([file_name, width, height, x1, y1, x2, y2])
482 + saved = True
483 +
484 + annotations = np.array(annotations)
485 + label = list(label)
486 + # print("label : ", label)
487 +
488 + if plate_type in ["3", "4", "6"]:
489 + col = label[:3] + [label[3] + label[4]] + label[5:]
490 + else:
491 + col = label[:2] + [label[2] + label[3]] + label[4:]
492 +
493 + # print(col)
494 + # np.insert(annotations, 5, col, axis=1)
495 +
496 + # print(annotations)
497 +
498 + df = pd.DataFrame(
499 + # annotations, columns=["image path", "x1", "y1", "x2", "y2", "label"]
500 + annotations,
501 + columns=["image path","width","height", "x1", "y1", "x2", "y2"],
502 + )
503 + df["label"] = col
504 +
505 + print(df)
506 +
507 + df.to_csv(self.plate_images_path+"annotations.csv", mode="a", header=False, index=False)
508 +
509 + return saved
510 +
511 +
512 +components_path = "components_2/"
513 +plate_images_path = "../dataset/test/test1/"
514 +detection_annotations_path = "detection_annotations/"
515 +recognition_data_path = "recognition_data/"
516 +
517 +# choose right folder for desired type
518 +
519 +# for parking ->
520 +# cars_path = 'cars/'
521 +# background_path = 'background/'
522 +
523 +# for cctv ->
524 +# cars_path = 'cctv_cars/'
525 +# background_path = 'cctv_background/'
526 +
527 +
528 +cars_path = "cars/"
529 +background_path = "background/"
530 +
531 +
532 +desired_number = 500
533 +desired_types = ["1", "2", "3", "4", "5", "6"]
534 +# cctv parking recognition provide choose one to generate and store things
535 +desired_dataset = "parking"
536 +count = 1
537 +
538 +
539 +generator = ImageGenerator(
540 + components_path=components_path,
541 + plate_images_path=plate_images_path,
542 + detection_annotations_path=detection_annotations_path,
543 + recognition_data_path=recognition_data_path,
544 + cars_path=cars_path,
545 + background_path=background_path,
546 +)
547 +
548 +generator.build_data(
549 + desired_number=desired_number,
550 + desired_types=desired_types,
551 + desired_dataset=desired_dataset,
552 + count=count,
553 +)
voc2coco @ 30a272cc
1 +Subproject commit 30a272ccb3778bade716d0fb1556df47ef8225f4