Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- def label2xml(list_path, category_id, ANN_DIR='data/Annotations/', IMG_DIR='data/images/', SET_DIR='data/ImageSets/',
- train=True):
- set_file_list = [SET_DIR + 'train.txt', SET_DIR + 'minival.txt', SET_DIR + 'testdev.txt', SET_DIR + 'test.txt']
- TYPEPREFIX = 'train' if train else 'val'
- lines = read_lines(list_path)
- line_i = 0
- last_name = ''
- E = None
- img_annotation = None
- image_file = None
- while line_i < len(lines):
- line = lines[line_i]
- infos = line.split()
- image_path = infos[0]
- box = infos[1:]
- imagename = str(int(image_path.split('.')[0].split('/')[-1]))
- print('last_path: %s, image_path: %s' % (last_name, image_path))
- if last_name != imagename:
- if train:
- write_line(set_file_list[0], "%s.jpg %s.xml" % (imagename, imagename))
- else:
- for i in range(1, 4):
- write_line(set_file_list[i], "%s.jpg %s.xml" % (imagename, imagename))
- if last_name != '':
- xml_pretty = etree.tostring(img_annotation, pretty_print=True)
- with open(ANN_DIR + last_name + ".xml", 'wb') as ann_file:
- ann_file.write(xml_pretty)
- image_file = imread(image_path)
- copyfile(image_path, IMG_DIR + '%d.jpg' % int(imagename))
- if path.exists(ANN_DIR + imagename + ".xml"):
- E = objectify.ElementMaker(annotate=False)
- img_annotation = objectify.fromstring(read_content(ANN_DIR + imagename + ".xml"))
- else:
- E = objectify.ElementMaker(annotate=False)
- img_annotation = E.annotation(
- E.folder(TYPEPREFIX),
- E.filename(imagename),
- E.source(
- E.database('coco_cattle'),
- ),
- E.size(
- E.width(image_file.shape[1]),
- E.height(image_file.shape[0]),
- E.depth(3),
- ),
- E.segmented(0)
- )
- last_name = imagename
- objectNode = E.object(
- E.name(str(category_id)),
- E.pose("Unspecified"),
- E.truncated("0"),
- E.difficult("0"),
- E.bndbox(
- E.xmin(str(int(float(box[0]) * image_file.shape[1]))),
- E.ymin(str(int(float(box[1]) * image_file.shape[0]))),
- E.xmax(str(int(float(box[2]) * image_file.shape[1]))),
- E.ymax(str(int(float(box[3]) * image_file.shape[0]))),
- ),
- )
- img_annotation.append(objectNode)
- line_i += 1
Add Comment
Please, Sign In to add comment