-
Notifications
You must be signed in to change notification settings - Fork 0
/
voc_helpers.py
436 lines (354 loc) · 16.2 KB
/
voc_helpers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
"""
Helper functions for PASCAL VOC Data Set Tools (see prepare_dataset.py)
MIT License
Copyright (c) 2023 @10XTMY, molmez.io
"""
import os
import sys
import traceback
import random
import string
import secrets
import shutil
import tqdm
import datetime
import xml.etree.ElementTree as ElementTree
from PIL import Image
def get_time_date():
timestamp_now = datetime.datetime.now()
time_stamp = datetime.timedelta(minutes=timestamp_now.minute,
seconds=timestamp_now.second,
hours=timestamp_now.hour)
date = timestamp_now.date()
return time_stamp, date
def write_to_txt_file(filename, set_to_write, existing):
"""
Appends data to existing txt file or writes data to a new one
:param filename:name of txt file to write
:param set_to_write: set of data to write to txt file
:param existing: does the file already exist?
"""
mode = 'a' if existing else 'w'
try:
with open(filename, mode) as file:
for line in set_to_write:
to_write = line.rsplit(".", 1)[0]
file.write(f'{to_write}\n')
except IOError as err:
raise IOError(f"Error writing to {filename}: {err}") from err
else:
print(f'{filename} saved')
def get_random_alphanum(length):
"""
Generates a random alphanumerical string
:param length: desired string length to be returned
:return: a random alphanumerical string of the specified length.
"""
letters_and_digits = string.ascii_letters + string.digits
return ''.join(secrets.choice(letters_and_digits) for _ in range(length))
def partition_list(list_to_split, percent, shuffle):
"""
:param list_to_split: provide a list to be partitioned
:param percent: desired percentage of partition
:param shuffle: True to shuffle the list before splitting it
:return: two lists, the first being the percentage of the original
"""
if not 0 <= percent <= 100:
raise ValueError("Percentage must be between 0 and 100")
if shuffle:
random.shuffle(list_to_split)
split_index = round(len(list_to_split) * (percent / 100))
return list_to_split[:split_index], list_to_split[split_index:]
def split_list_in_half(list_to_split):
"""
:param list_to_split: input list
:return: two lists, input list split in half
"""
middle = len(list_to_split) // 2
return list_to_split[:middle], list_to_split[middle:]
def get_new_file_name(existing_names, length=15, max_attempts=100000):
"""
:param existing_names: list of existing file names to avoid conflict
:param length: desired length of new file name
:param max_attempts: maximum number of attempts before giving up
:return: new filename
"""
for _ in range(max_attempts):
filename = get_random_alphanum(length)
if filename not in existing_names:
return filename
raise RuntimeError(f"Failed to generate a unique filename after {max_attempts} attempts.")
def count_xml_labels(xml_dir, lbl_list):
"""
counts the amount of given labels in all xml files within provided directory
:param xml_dir: xml directory to scan
:param lbl_list: a list of label names to count
:return: dictionary of labels and their counts
"""
label_counts = dict.fromkeys(lbl_list, 0)
for dir_path, _, files in os.walk(xml_dir):
for xml_file in files:
if not xml_file.lower().endswith('.xml'):
continue
xml_path = os.path.join(dir_path, xml_file)
try:
tree = ElementTree.parse(xml_path)
except ElementTree.ParseError as e:
raise ElementTree.ParseError(f'error parsing {xml_file}: {e}')
for name in tree.getroot().iter('name'):
if name.text in label_counts:
label_counts[name.text] += 1
return label_counts
def remove_object_from_xml_files(xml_directory, object_names_set):
"""
removes objects from xml files within a provided directory
:param xml_directory: directory path for xml files
:param object_names_set: a set of object names to be removed
"""
for dir_path, _, files in os.walk(xml_directory):
for xml_file in tqdm.tqdm(files):
if not xml_file.lower().endswith('.xml'):
continue
xml_path = os.path.join(dir_path, xml_file)
try:
tree = ElementTree.parse(xml_path)
tree_root = tree.getroot()
# collect items to be removed to avoid modifying the list while iterating over it
objects_to_remove = [obj for obj in tree_root.findall("object")
if obj.find('name') is not None and obj.find('name').text in object_names_set]
for obj in objects_to_remove:
tree_root.remove(obj)
if objects_to_remove:
tree.write(xml_path)
except ElementTree.ParseError as e:
raise ElementTree.ParseError(f"Error parsing {xml_file}: {e}")
except IOError as e:
raise IOError(f"Error writing changes to {xml_file}: {e}")
def fix_missing_xml_object_name(xml_directory, label_text):
"""
scans xml file for objects with no name and names them according
to the label_text
:param xml_directory: directory of xml files
:param label_text: label to replace object name with
"""
for dir_path, _, files in os.walk(xml_directory):
for xml_file in tqdm.tqdm(files):
if not xml_file.lower().endswith('.xml'):
continue
xml_path = os.path.join(dir_path, xml_file)
try:
tree = ElementTree.parse(xml_path)
objects = tree.findall("object")
updated = False
for obj in objects:
class_name = obj.find('name')
if class_name is not None and class_name.text is None:
class_name.text = label_text
updated = True
if updated:
tree.write(xml_path)
except ElementTree.ParseError as e:
raise ElementTree.ParseError(f"Error parsing {xml_file}: {e}")
except Exception as e:
raise Exception(f"Error processing {xml_file}: {e}")
def replace_xml_file_information(xml_file, replace_dict):
"""
overwrites existing xml file with new values
:param xml_file: xml file to edit
:param replace_dict: dictionary of xml tags and the value to place in them
"""
try:
tree = ElementTree.parse(xml_file)
root = tree.getroot()
for key, value in replace_dict.items():
elements = root.findall(key)
for elem in elements:
elem.text = value
tree.write(xml_file)
except ElementTree.ParseError as e:
raise ElementTree.ParseError(f'error parsing {xml_file}: {e}')
except Exception as e:
raise Exception(f'error processing {xml_file}: {e}')
def generate_negative_data_set(existing_names, negative_images, negative_output_dir, xml_template):
"""
generates new filename for both xml and jpg
copies images to output folder -> copies xml template to output -> edits xml file to match image
:param existing_names: list of existing file names
:param negative_images: folder containing negative image files (jpg)
:param negative_output_dir: folder to output the negative data set (jpg & xml)
:param xml_template: negative xml template to use
"""
print('Generating negative data set...')
for dir_path, _, files in os.walk(negative_images):
for image_file in tqdm.tqdm(files):
file_ext = image_file.lower().split('.')[-1]
if file_ext not in ['jpg', 'jpeg']:
continue
try:
# process image
img_path = os.path.join(dir_path, image_file)
img = Image.open(img_path)
width, height = img.size
except IOError as e:
raise IOError(f'error opening {image_file}: {e}')
try:
# generate new filename and prepare file paths
filename = get_new_file_name(existing_names)
existing_names.add(filename)
image_name = f'{filename}.{file_ext}'
image_out_path = os.path.join(negative_output_dir, image_name)
xml_file = f'{filename}.xml'
xml_path = os.path.join(negative_output_dir, xml_file)
# copy image and XML template
shutil.copy(img_path, image_out_path)
shutil.copy(xml_template, xml_path)
except IOError as e:
raise IOError(f'error copying {image_file} to {negative_output_dir}: {e}')
try:
# replace information in XML
replace_dict = {
'filename': image_name,
'path': image_name,
'width': str(width),
'height': str(height),
'xmax': str(width),
'ymax': str(height)
}
try:
replace_xml_file_information(xml_path, replace_dict)
except Exception as e:
raise Exception(f'error processing {xml_file}: {e}')
except ElementTree.ParseError as e:
raise ElementTree.ParseError(f'error parsing {xml_file}: {e}')
print(f'..Negative data set generated in {negative_output_dir} directory')
def collect_current_data_set(data_set_path, jpg_dir, annotations_dir):
"""
copies existing data set to appropriate output directories, making a note of file names
:param data_set_path: directory containing existing data set's jpg and xml files
:param jpg_dir: output jpg directory
:param annotations_dir: output annotations directory
:return: list of unique file names in the data set that has been collected
"""
print('Collecting current data set...')
existing_names = set()
for dir_path, _, files in os.walk(data_set_path):
for file in tqdm.tqdm(files):
file_ext = file.lower().split('.')[-1] # last element of the split is the file extension
if file_ext not in ['jpg', 'jpeg']:
continue
filename, _ = os.path.splitext(file)
existing_names.add(filename)
jpg_path = os.path.join(dir_path, file)
xml_file = f'{filename}.xml'
xml_path = os.path.join(dir_path, xml_file)
try:
shutil.copy(jpg_path, os.path.join(jpg_dir, file))
shutil.copy(xml_path, os.path.join(annotations_dir, xml_file))
except IOError as e:
raise IOError(f'error copying {file} to {jpg_dir} or {annotations_dir}: {e}')
print('..Finished collecting current data set')
return existing_names
def generate_txt_files(jpg_dir, txt_dir, split_percentage):
"""
scans image directory and generates text file lists for the data set
:param jpg_dir: directory containing jpgs to be listed
:param txt_dir: output txt directory
:param split_percentage: split percentage for test and validation sets (recommended 20)
"""
print('Generating txt files...')
try:
jpeg_names = [file for file in os.listdir(jpg_dir) if file.lower().endswith('.jpg')]
except OSError as e:
raise OSError(f"error listing {jpg_dir}: {e}")
# split the list into training list and split_set list
try:
split_set, train_set = partition_list(jpeg_names, split_percentage, shuffle=True)
print(f'Split test count: {len(split_set)}')
print(f'Train set count: {len(train_set)}')
except ValueError as e:
raise ValueError(f"error partitioning list: {e}")
# split the split_set list into validation and test lists
test_set, val_set = split_list_in_half(split_set)
print(f'Test set count: {len(test_set)}')
print(f'Validation set count: {len(val_set)}')
print('Writing txt files...')
file_mappings = {
'train.txt': (train_set, False),
'test.txt': (test_set, False),
'val.txt': (val_set, False),
'trainval.txt': (train_set, False), # first write with train_set
'trainval.txt': (val_set, True) # then append with val_set
}
for filename, (set_to_write, is_existing) in file_mappings.items():
try:
write_to_txt_file(os.path.join(txt_dir, filename), set_to_write, existing=is_existing)
except IOError as e:
raise IOError(f"Error writing to {filename}: {e}")
print(f'..All txt files saved in: {txt_dir}')
def inject_negative_data_set(negatives_dir, jpg_dir, annotations_dir, txt_dir, val_test_percentage):
"""
takes negative data set, copies to output folders, generates text files for new data set
:param negatives_dir: directory containing negative jpg and xml files
:param jpg_dir: output jpg directory
:param annotations_dir: output annotations directory
:param txt_dir: output txt directory
:param val_test_percentage: split percentage for test and validation sets (recommended 20)
"""
print('Injecting negative data set...')
for dirpath, _, files in os.walk(negatives_dir):
for file in files:
file_ext = file.lower().split('.')[-1]
if file_ext in ['jpg', 'jpeg']:
destination = os.path.join(jpg_dir, file)
elif file_ext == 'xml':
destination = os.path.join(annotations_dir, file)
else:
print(f'File is neither .jpg, .jpeg nor .xml: {file}')
continue
source_path = os.path.join(dirpath, file)
try:
shutil.copy(source_path, destination)
except IOError as e:
raise IOError(f'error copying {file} to {destination}: {e}')
try:
generate_txt_files(jpg_dir, txt_dir, val_test_percentage)
except (OSError, ValueError, IOError) as e:
print(f'error generating txt files: {e}')
traceback.print_tb(e.__traceback__)
sys.exit(1)
print('..Finished injecting negative data set into current data set')
def prepare_voc(input_directory, image_directory, annotations_directory, existing_names):
"""
takes current data set and renames all files unique, converts any pngs into jpgs
saves new files into VOC output directories
:param input_directory: current data set
:param image_directory: output directory for images
:param annotations_directory: output directory for annotations
:param existing_names: set containing current filenames
"""
for dir_path, _, files in os.walk(input_directory):
for file_name in tqdm.tqdm(files):
filename_no_ext, file_ext = os.path.splitext(file_name)
file_ext = file_ext.lower()
if file_ext == '.png':
new_filename = get_new_file_name(existing_names)
jpg_path = os.path.join(image_directory, f'{new_filename}.jpg')
xml_path = os.path.join(annotations_directory, f'{new_filename}.xml')
xml_file_name = f'{filename_no_ext}.xml'
xml_input_path = os.path.join(dir_path, xml_file_name)
try:
# convert PNG to JPG
with Image.open(os.path.join(dir_path, file_name)) as img:
rgb_img = img.convert('RGB')
rgb_img.save(jpg_path)
# update and copy XML file if it exists
if xml_file_name in files:
replace_dict = {'filename': f'{new_filename}.jpg', 'path': f'{new_filename}.jpg'}
try:
replace_xml_file_information(xml_input_path, replace_dict)
except (ElementTree.ParseError, Exception) as e:
raise Exception(f"Error replacing xml file information {xml_file_name}: {e}")
shutil.copy(xml_input_path, xml_path)
existing_names.add(new_filename)
except IOError as e:
raise IOError(f"Error processing file {file_name}: {e}")