Skip to content

Commit 6a16328

Browse files
author
Eli Selkin
committed
allowed selective generation of annotation label types
1 parent 5699835 commit 6a16328

File tree

1 file changed

+24
-24
lines changed

1 file changed

+24
-24
lines changed

scripts/Construct_XML_Annotations_from_COCO_JSON.py

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -125,22 +125,23 @@
125125
E.segmented(0)
126126
)
127127
for row in image[1].iterrows():
128-
objectNode = E.object(
129-
E.name(str(row[1]['category_id'])),
130-
E.pose("Unspecified"),
131-
E.truncated("0"),
132-
E.difficult("0"),
133-
E.bndbox(
134-
E.xmin(
135-
str(int(round(row[1]['bbox'][0]))) if ROUND_COORDS else str(int(row[1]['bbox'][0]))),
136-
E.ymin(
137-
str(int(round(row[1]['bbox'][1]))) if ROUND_COORDS else str(int(row[1]['bbox'][1]))),
138-
E.xmax(str(int(round(row[1]['bbox'][0] + row[1]['bbox'][2])))
139-
if ROUND_COORDS else str(int(row[1]['bbox'][0] + row[1]['bbox'][2]))),
140-
E.ymax(str(int(round(row[1]['bbox'][1] + row[1]['bbox'][3])))
141-
if ROUND_COORDS else str(int(row[1]['bbox'][1] + row[1]['bbox'][3]))),
142-
),
143-
)
128+
if row[1]['category_id'] in label_ids:
129+
objectNode = E.object(
130+
E.name(str(row[1]['category_id'])),
131+
E.pose("Unspecified"),
132+
E.truncated("0"),
133+
E.difficult("0"),
134+
E.bndbox(
135+
E.xmin(
136+
str(int(round(row[1]['bbox'][0]))) if ROUND_COORDS else str(int(row[1]['bbox'][0]))),
137+
E.ymin(
138+
str(int(round(row[1]['bbox'][1]))) if ROUND_COORDS else str(int(row[1]['bbox'][1]))),
139+
E.xmax(str(int(round(row[1]['bbox'][0] + row[1]['bbox'][2])))
140+
if ROUND_COORDS else str(int(row[1]['bbox'][0] + row[1]['bbox'][2]))),
141+
E.ymax(str(int(round(row[1]['bbox'][1] + row[1]['bbox'][3])))
142+
if ROUND_COORDS else str(int(row[1]['bbox'][1] + row[1]['bbox'][3]))),
143+
),
144+
)
144145
img_annotation.append(objectNode)
145146
xml_pretty = etree.tostring(img_annotation, pretty_print=True)
146147
with open(ANN_DIR + imagename + ".xml", 'wb') as ann_file:
@@ -151,12 +152,11 @@
151152
if TRAIN_SPLIT and TRAIN_SPLIT < 1.0:
152153
TRAIN_DIR = 'Images/train_minusminival%d/' % (YEAR)
153154
TRAIN_FULL_DIR = '%s%s' % (ROOT_COCO, TRAIN_DIR)
154-
TRAIN_ANN_DIR = 'Annotations/train_minusminival%d/' % (
155-
ROOT_COCO, YEAR)
155+
TRAIN_ANN_DIR = 'Annotations/train_minusminival%d/' % (YEAR)
156156
TRAIN_FULL_ANN_DIR = '%s%s' % (ROOT_COCO, TRAIN_ANN_DIR)
157-
MINIVAL_DIR = 'Images/minival%d/' % (ROOT_COCO, YEAR)
157+
MINIVAL_DIR = 'Images/minival%d/' % (YEAR)
158158
MINIVAL_FULL_DIR = '%s%s' % (ROOT_COCO, MINIVAL_DIR)
159-
MINIVAL_ANN_DIR = 'Annotations/minival%d/' % (ROOT_COCO, YEAR)
159+
MINIVAL_ANN_DIR = 'Annotations/minival%d/' % (YEAR)
160160
MINIVAL_FULL_ANN_DIR = '%s%s' % (ROOT_COCO, MINIVAL_ANN_DIR)
161161
if not path.isdir(TRAIN_DIR):
162162
mkdir(TRAIN_FULL_DIR)
@@ -173,7 +173,7 @@
173173
SELECTED_FOR_MINIVAL.append(int(RANDOM_IDX))
174174
SELECTED_FOR_TRAIN = sorted(
175175
list(set(list(range(len(IMAGE_NAMES)))).difference(SELECTED_FOR_MINIVAL)))
176-
with open('train.txt', 'w') as train_file:
176+
with open('train2.txt', 'w') as train_file:
177177
for idx in SELECTED_FOR_TRAIN:
178178
if COPY_FILES:
179179
copyfile(
@@ -191,7 +191,7 @@
191191
TRAIN_FULL_ANN_DIR + IMAGE_NAMES[idx] + ".xml")
192192
train_file.write(
193193
"/" + TRAIN_DIR + IMAGE_NAMES[idx] + ".jpg /" + TRAIN_ANN_DIR + IMAGE_NAMES[idx] + ".xml\n")
194-
with open('minival.txt', 'w')as minival_file:
194+
with open('minival2.txt', 'w')as minival_file:
195195
for idx in SELECTED_FOR_MINIVAL:
196196
if COPY_FILES:
197197
copyfile(
@@ -214,7 +214,7 @@
214214
minival_file.write(
215215
"/" + MINIVAL_DIR + IMAGE_NAMES[idx] + ".jpg /" + MINIVAL_ANN_DIR + IMAGE_NAMES[idx] + ".xml\n")
216216
else:
217-
with open('train.txt', 'w') as train_file:
217+
with open('train2.txt', 'w') as train_file:
218218
IMG_RELATIVE = '/Images/train%d/' % (YEAR)
219219
TRAIN_ANN_DIR = 'Annotations/train%d/' % (YEAR)
220220
TRAIN_FULL_ANN_DIR = "%s%s" % (ROOT_COCO, TRAIN_ANN_DIR)
@@ -224,7 +224,7 @@
224224
train_file.write(
225225
IMG_RELATIVE + IMAGE_NAMES[i] + ".jpg /" + TRAIN_ANN_DIR + IMAGE_NAMES[i] + ".xml\n")
226226
else:
227-
with open('val.txt', 'w') as val_file:
227+
with open('val2.txt', 'w') as val_file:
228228
IMG_RELATIVE = '/Images/val%d/' % (YEAR)
229229
VAL_ANN_DIR = 'Annotations/val%d/' % (YEAR)
230230
VALL_FULL_ANN_DIR = '%s%s' % (ROOT_COCO, VAL_ANN_DIR)

0 commit comments

Comments
 (0)