Как преобразовать файлы .txt в xml для обнаружения объектов yolo - PullRequest
Новая
       102

Как преобразовать файлы .txt в xml для обнаружения объектов yolo

0 голосов
/ 20 января 2020

Я работаю над обнаружением объектов с помощью yolo (darkflow). У меня есть набор данных и метки для набора данных. Но проблема в том, что мне нужны аннотации в формате XML, но у меня есть в формате .txt. Как преобразовать TXT в xml?

Это образец TXT-файла. 

001.jpg 528 412 162 52  001

Я написал код для преобразования TXT в xml

import os
import xml.etree.cElementTree as ET
from lxml import etree
import cv2

def text_to_xml(file):

    filename = file.split('.')[0]
    extention = file.split('.')[1]

    image = cv2.imread(os.path.join(folder, file))
    height, width, depth = image.shape
    f = open(filename+'.txt', 'r')
    content = f.readlines()[0]
    content = content.split('\t')
    xmin, xmax, ymin, ymax = content[1], content[2], content[3], content[4]
    f.close()

    annotation = ET.Element('annotation')
    ET.SubElement(annotation, 'folder').text = folder
    ET.SubElement(annotation, 'filename').text = filename + '.jpg'
    ET.SubElement(annotation, 'segmented').text = '0'

    size = ET.SubElement(annotation, 'size')
    ET.SubElement(size, 'width').text = str(width)
    ET.SubElement(size, 'height').text = str(height)
    ET.SubElement(size, 'depth').text = str(depth)

    ob = ET.SubElement(annotation, 'object')
    ET.SubElement(ob, 'name').text = 'sign'
    ET.SubElement(ob, 'pose').text = 'Unspecified'
    ET.SubElement(ob, 'truncated').text = '0'
    ET.SubElement(ob, 'difficult').text = '0'

    bbox = ET.SubElement(ob, 'bndbox')
    ET.SubElement(bbox, 'xmin').text = str(xmin)
    ET.SubElement(bbox, 'ymin').text = str(ymin)
    ET.SubElement(bbox, 'xmax').text = str(xmax)
    ET.SubElement(bbox, 'ymax').text = str(ymax)

    xml_str = ET.tostring(annotation)
    root = etree.fromstring(xml_str)
    xml_str = etree.tostring(root, pretty_print=True)

    return xml_str

if __name__ == '__main__':
    folder = './images/'
    for file in os.listdir(folder):
        if file.split('.')[1] == 'jpg':
            text_to_xml(file)

Но когда я тренирую модель с преобразованными файлами xml, модель ничего не обнаруживает. Так что, думаю, я где-то напутал.

...