-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy patha.py
131 lines (105 loc) · 5.46 KB
/
a.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
#------------------------------------------------------------------------------------------------
#
# By: Jatin Kumar Mandav
#
# Generating XML Files for Darkflow from .txt files generated by OpenLabeling
# USAGE: generate_xml.py [-h] [--format {yolo,voc}] [--bbox_txt BBOX_TXT]
# [--img_dir IMG_DIR] [--class_list CLASS_LIST]
# [--save_dir SAVE_DIR]
#
# Generate XML Files from .txt file
#
# optional arguments:
# -h, --help show this help message and exit
# --format {yolo,voc} Bounding box format | Deafult: yolo
# --bbox_txt BBOX_TXT Path to bbox_txt dir generated by OpenLabeling | Deafult: bbox_txt/
# --img_dir IMG_DIR Path to Images folder | Default: images/
# --class_list CLASS_LIST Path to class_list.txt file | Default: class_list.txt
# --save_dir SAVE_DIR Where to save XML Files. | Default: annotations_xml/
#
#------------------------------------------------------------------------------------------------
from lxml import etree
import xml.etree.cElementTree as ET
import cv2
import os
import argparse
def write_xml(saveDir, imagefolder, imagename, imgWidth, imgHeight, depth, bounding_boxes, pose="Unspecified"):
annotation = ET.Element("annotaion")
ET.SubElement(annotation, 'folder').text = str(imagefolder)
ET.SubElement(annotation, 'filename').text = str(imagename)
size = ET.SubElement(annotation, 'size')
ET.SubElement(size, 'width').text = str(imgWidth)
ET.SubElement(size, 'height').text = str(imgHeight)
ET.SubElement(size, 'depth').text = str(depth)
ET.SubElement(annotation, 'segmented').text = '0'
for box in bounding_boxes:
obj = ET.SubElement(annotation, 'object')
ET.SubElement(obj, 'name').text = str(label_name)
ET.SubElement(obj, 'pose').text = str(pose)
ET.SubElement(obj, 'truncated').text = '0'
ET.SubElement(obj, 'difficult').text = '0'
bbox = ET.SubElement(obj, 'bndbox')
ET.SubElement(bbox, 'xmin').text = str(bboxx1)
ET.SubElement(bbox, 'ymin').text = str(bboxy1)
ET.SubElement(bbox, 'xmax').text = str(bboxx2)
ET.SubElement(bbox, 'ymax').text = str(bboxy2)
xml_str = ET.tostring(annotation)
root = etree.fromstring(xml_str)
xml_str = etree.tostring(root, pretty_print=True)
save_path = os.path.join(saveDir, os.path.splitext(imagename)[0] + ".xml")
with open(save_path, 'wb') as temp_xml:
temp_xml.write(xml_str)
def yolo_to_x_y(x_center, y_center, x_width, y_height, width, height):
x_center *= width
y_center *= height
x_width *= width
y_height *= height
x_width /= 2.0
y_height /= 2.0
return int(x_center - x_width), int(y_center - y_height), int(x_center + x_width), int(y_center + y_height)
parser = argparse.ArgumentParser(description="Generate XML Files from .txt file")
parser.add_argument('--format', default='yolo', type=str, choices=['yolo', 'voc'], help="Bounding box format | Deafult: yolo")
parser.add_argument('--bbox_txt', default='bbox_txt/', type=str, help='Path to bbox_txt dir generated by OpenLabeling | Default: bbox_txt/')
parser.add_argument('--img_dir', default='images/', type=str, help='Path to Images folder | Default: images/')
parser.add_argument('--class_list', default='class_list.txt', type=str, help='Path to class_list.txt file | Default: class_list.txt')
parser.add_argument('--save_dir', default='annotations_xml/', type=str, help='Where to save XML Files | Default: annotations_xml/')
args = parser.parse_args()
bbox_txt = args.bbox_txt
imagefolder = args.img_dir
saveDir = args.save_dir
format = args.format
if not os.path.exists(saveDir):
os.makedirs(saveDir)
with open(args.class_list, 'r') as f:
labels = f.readlines()
for i in range(len(labels)):
labels[i] = labels[i].replace('\n', '')
for imagename in os.listdir(imagefolder):
img = cv2.imread(os.path.join(imagefolder, imagename))
imgHeight, imgWidth, depth = img.shape
with open(os.path.join(bbox_txt, os.path.splitext(imagename)[0] + '.txt')) as f:
content = f.readlines()
bounding_boxes = []
for line in content:
values_str = line.split()
if format == 'yolo':
class_index, x_center, y_center, x_width, y_height = map(float, values_str)
class_index = int(class_index)
bboxx1, bboxy1, bboxx2, bboxy2 = yolo_to_x_y(x_center, y_center, x_width, y_height, imgWidth, imgHeight)
if x_center == int(x_center):
error = ("You selected the 'yolo' format but your labels "
"seem to be in a different format. Consider "
"removing your old label files.")
raise Exception(textwrap.fill(error, 70))
else:
try:
x1, y1, x2, y2, class_index = map(int, values_str)
except ValueError:
error = ("You selected the 'voc' format but your labels "
"seem to be in a different format. Consider "
"removing your old label files.")
raise Exception(textwrap.fill(error, 70))
bboxx1, bboxy1, bboxx2, bboxy2 = x1-1, y1-1, x2-1, y2-1
label_name = labels[class_index]
bounding_boxes.append([label_name, bboxx1, bboxy1, bboxx2, bboxy2])
write_xml(saveDir, imagefolder, imagename, imgWidth, imgHeight, depth, bounding_boxes, pose="Unspecified")