snappy_annotator/voc_save_load.py at master · wvuvl/snappy_annotator · GitHub

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
"""
Methods for saving annotations to Pascal VOC format
"""

import os
import colored
import xml.etree.ElementTree as et
from xml.dom import minidom

import json

error_msg = colored.fg("red") + colored.attr("bold")


# Takes annotation and other data for current image and translates into a Pascal VOC-formatted .xml file.
def save_to_voc_xml(filename, folder, path, database, dims, annotations, labels, file_extension, observation_rank):

    xml = et.Element('annotation')
    fold = et.SubElement(xml, 'folder')
    fold.text = folder
    et.SubElement(xml, 'filename').text = filename
    et.SubElement(xml, 'path').text = path
    src = et.SubElement(xml, 'source')
    et.SubElement(src, 'database').text = database
    sz = et.SubElement(xml, 'size')
    et.SubElement(sz, 'width').text = str(dims[1])
    et.SubElement(sz, 'height').text = str(dims[0])
    et.SubElement(sz, 'depth').text = str(dims[2])

    for i in range(0, int(len(annotations) / 2)):
        obj = et.SubElement(xml, 'object')
        et.SubElement(obj, 'name').text = labels[i]
        et.SubElement(obj, 'pose').text = 'Unspecified'
        et.SubElement(obj, 'truncated').text = '0'
        et.SubElement(obj, 'difficult').text = '0'
        et.SubElement(obj, 'observation_rank').text = str(observation_rank)
        bndbox = et.SubElement(obj, 'bndbox')
        et.SubElement(bndbox, 'xmin').text = str(annotations[i * 2][0])
        et.SubElement(bndbox, 'ymin').text = str(annotations[i * 2][1])
        et.SubElement(bndbox, 'xmax').text = str(annotations[i * 2 + 1][0])
        et.SubElement(bndbox, 'ymax').text = str(annotations[i * 2 + 1][1])

    rough_string = et.tostring(xml, 'utf-8')
    reparsed = minidom.parseString(rough_string)
    pretty_string = reparsed.toprettyxml(indent="\t")

    p = filename.find('.')
    with open(os.path.join(folder, filename[:p] + file_extension), 'w') as x:
        x.writelines(pretty_string)


# Reads in an xml file, pulls all important information and returns it to program in usable data format
# NOTE: For metadata, currently only reads in path, image dimensions, and database,
# and for each annotation, only reads label name and bounding box dimensions.
# NOTE: This method will return the path and database metadata, but currently does nothing
# with them. When saving, this information will be overwritten by the save function
def load_from_voc_xml(file_pth, filename, file_extension):
    path = ''
    database = ''
    annotation = []
    labels = []
    xml_dims = ()
    filename = os.path.join(file_pth, str(filename[:filename.find('.')]) + file_extension)

    if os.path.exists(filename):
        tree = et.parse(filename)
        root = tree.getroot()
        if root.find('path') is not None:
            path = root.find('path').text
        src = root.find('source')
        if src is not None and src.find('database') is not None:
            database = src.find('database').text
        if root.find('size') is not None:
            size = root.find('size')
            width = int(size.find('width').text)
            height = int(size.find('height').text)
            depth = int(size.find('depth').text)
            xml_dims = (width, height, depth)
        for object in root.iter('object'):
            labels.append(object.find('name').text)
            annotation.append(
                (int(object.find('bndbox').find('xmin').text), int(object.find('bndbox').find('ymin').text)))
            annotation.append(
                (int(object.find('bndbox').find('xmax').text), int(object.find('bndbox').find('ymax').text)))

    return path, database, xml_dims, annotation, labels