Showing
4 changed files
with
1063 additions
and
0 deletions
data_generator/Full_generator.py
0 → 100644
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 |
data_generator/My_generator.py
0 → 100644
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 | +) |
File moved
-
Please register or login to post a comment