-
Notifications
You must be signed in to change notification settings - Fork 0
/
transform_data.py
97 lines (80 loc) · 4.18 KB
/
transform_data.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
import os
import shutil
import xml.etree.ElementTree as ET
import copy
_copy_suffix = 1
_slash = os.sep
def remove_inside(path):
if os.path.exists(os.path.join(path, "annotations")):
shutil.rmtree(os.path.join(path, "annotations"), ignore_errors=True)
if os.path.exists(os.path.join(path, "annotations", "xmls")):
shutil.rmtree(os.path.join(path, "annotations", "xmls"), ignore_errors=True)
if os.path.exists(os.path.join(path, "images")):
shutil.rmtree(os.path.join(path, "images"), ignore_errors=True)
def transform(input_folder, output_folder, folder_ann, folder_img, display_mode):
count_files = 0
try:
os.mkdir(os.path.join(output_folder, 'annotations'))
os.mkdir(os.path.join(output_folder, 'annotations', 'xmls'))
os.mkdir(os.path.join(output_folder, 'images'))
except OSError:
print("Create directory %s failed. Maybe it already exists." % output_folder)
else:
if display_mode != 0:
print("Directory successfully created %s" % output_folder)
files = os.listdir(input_folder + _slash + folder_ann)
xml_folder = os.path.join(output_folder, 'annotations', 'xmls')
trainval_path = os.path.join(output_folder, 'annotations')
img_folder = os.path.join(output_folder, 'images')
_file = open(os.path.join(trainval_path, 'trainval.txt'), 'w')
for xml_file in files:
if xml_file.endswith('.xml'):
count_files += 1
ann_path_in = os.path.join(input_folder, folder_ann, xml_file)
ann_path_out = os.path.join(xml_folder, xml_file)
file_name = xml_file[:-4]
_file.write(str(file_name) + '\n')
img_path_in = os.path.join(input_folder, folder_img, file_name + '.jpeg')
img_path_out = os.path.join(img_folder, file_name + '.jpeg')
shutil.copyfile(ann_path_in, ann_path_out)
shutil.copyfile(img_path_in, img_path_out)
print("Copying completed! \n Number of files:" + str(count_files))
_file.close()
def move_files(in_path, _debug, xmls_path, images_path):
files = os.listdir(os.path.join(in_path, 'annotations' + _slash + 'xmls'))
for xml_file in files:
if xml_file.endswith('.xml'):
ann_path_in = os.path.join(in_path, 'annotations', 'xmls', xml_file)
ann_path_out = os.path.join(xmls_path, xml_file)
file_name = xml_file[:-4]
if os.path.exists(ann_path_out):
global _copy_suffix
new_filename = file_name + '_' + str(_copy_suffix)
ann_path_out = os.path.join(xmls_path, new_filename + '.xml')
shutil.copyfile(ann_path_in, ann_path_out)
xml_tree = ET.parse(ann_path_out)
new_xml_tree = copy.deepcopy(xml_tree)
filename_img = new_xml_tree.find('filename')
filename_img.text = new_filename + '.jpeg'
new_xml_tree.write(ann_path_out)
img_path_in = os.path.join(in_path, 'images', file_name + '.jpeg')
img_path_out = os.path.join(images_path, new_filename + '.jpeg')
shutil.copyfile(img_path_in, img_path_out)
_copy_suffix += 1
else:
img_path_in = os.path.join(in_path + _slash + 'images' + _slash, file_name + '.jpeg')
img_path_out = os.path.join(images_path, file_name + '.jpeg')
shutil.copyfile(ann_path_in, ann_path_out)
shutil.copyfile(img_path_in, img_path_out)
def delete_excess(in_path):
files = os.listdir(in_path + _slash + 'annotations' + _slash + 'xmls')
for xml_file in files:
xml_tree = ET.parse(os.path.join(in_path, 'annotations', 'xmls', xml_file))
new_xml_tree = copy.deepcopy(xml_tree)
filename_img = new_xml_tree.find('object')
if filename_img is None:
print('Delete file - ' + str(xml_file))
os.remove(os.path.join(in_path, 'annotations', 'xmls', xml_file))
file_name = xml_file[:-4]
img = os.path.join(in_path, 'images', file_name + '.jpeg')
os.remove(img)