Utils/labelme_converter_utils.py at master · RT-Rakesh/Utils · 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
# -*- coding: utf-8 -*-
"""
Created on Wen Dec 12 2018
@author: Rt-Rakesh

This script is a utility to convert data to and from the format accepted by the LabelMe tool.
Usage: Could be used by calling in jupyter notbooks.
"""

import pandas as pd
import xml.etree.cElementTree as ET
from tqdm import tqdm
import os


def convert_csv_2_labelme_xml(csv_path, xml_folder):
    """
    This Function writes the annoations in a xml format compatible with label me toolself.
    Args:
    1.csv path: --str The path to the csv, it should have the data in the following format (path,xmin,ymin,xmax,ymax,label).
    2.xml_folder: -- str The output folder path to be mentioned here.

    Returns:
    The xml corresponding to the input annotaions are written in the path mentioned.
    """

    data = pd.read_csv(csv_path)
    path_list = list(set(data.path))
    for imgs in path_list:
        temp_data = data[data['path'] == imgs].reset_index(drop=True)
        annotation = ET.Element("annotation")
        ET.SubElement(annotation, "filename").text = str(imgs)
        ET.SubElement(annotation, "folder").text = "roboreader_examples"
        source = ET.SubElement(annotation, "source")
        ET.SubElement(source, "sourceImage").text = "The MIT-CSAIL database of objects and scenes"
        ET.SubElement(source, "sourceAnnotation").text = "LabelMe Webtool"
        for i, rows in tqdm(temp_data.iterrows()):
            obj = ET.SubElement(annotation, "object")
            ET.SubElement(obj, "name").text = str(rows['label'])
            ET.SubElement(obj, "deleted").text = "0"
            ET.SubElement(obj, "verified").text = "0"
            ET.SubElement(obj, "occluded").text = "no"
            ET.SubElement(obj, "attributes")
            parts = ET.SubElement(obj, "parts")
            ET.SubElement(parts, "hasparts")
            ET.SubElement(parts, "ispartof")
            ET.SubElement(obj, "date").text = "31-Dec-2099 11:59:59"
            ET.SubElement(obj, "id").text = str(int(i)+1)
            ET.SubElement(obj, "type").text = "bounding_box"
            polygon = ET.SubElement(obj, "polygon")
            ET.SubElement(polygon, "username").text = "anonymous"
            pt = ET.SubElement(polygon, "pt")
            ET.SubElement(pt, "x").text = str(rows['xmin'])
            ET.SubElement(pt, "y").text = str(rows['ymin'])
            pt = ET.SubElement(polygon, "pt")
            ET.SubElement(pt, "x").text = str(rows['xmax'])
            ET.SubElement(pt, "y").text = str(rows['ymin'])
            pt = ET.SubElement(polygon, "pt")
            ET.SubElement(pt, "x").text = str(rows['xmax'])
            ET.SubElement(pt, "y").text = str(rows['ymax'])
            pt = ET.SubElement(polygon, "pt")
            ET.SubElement(pt, "x").text = str(rows['xmin'])
            ET.SubElement(pt, "y").text = str(rows['ymax'])
        tree = ET.ElementTree(annotation)
        tree.write(os.path.join(xml_folder, imgs.split('.')[0]+".xml"))