VOC数据和YOLO数据的转换(txt-xml and xml-txt)

好久没写博客 找写博客的连接找了好久。。。

还是得记录代码,很多代码用过之后很难找到了就。

下面代码是yolo数据格式转voc数据格式,即txt转xml(copy一位老哥的代码)

import cv2
import os

xml_head = '''<annotation>
    <folder>VOC2007</folder>
    <!--文件名-->
    <filename>{}</filename>.
    <source>
        <database>The VOC2007 Database</database>
        <annotation>PASCAL VOC2007</annotation>
        <image>flickr</image>
        <flickrid>325991873</flickrid>
    </source>
    <owner>
        <flickrid>null</flickrid>
        <name>null</name>
    </owner>    
    <size>
        <width>{}</width>
        <height>{}</height>
        <depth>{}</depth>
    </size>
    <segmented>0</segmented>
    '''
xml_obj = '''
    <object>        
        <name>{}</name>
        <pose>Rear</pose>
        <!--是否被裁减,0表示完整,1表示不完整-->
        <truncated>0</truncated>
        <!--是否容易识别,0表示容易,1表示困难-->
        <difficult>0</difficult>
        <!--bounding box的四个坐标-->
        <bndbox>
            <xmin>{}</xmin>
            <ymin>{}</ymin>
            <xmax>{}</xmax>
            <ymax>{}</ymax>
        </bndbox>
    </object>
    '''
xml_end = '''
</annotation>'''

labels = ['A', 'B', 'C']  # label for datasets

cnt = 0
for yolo_file in os.listdir('labels'):
    jpg = 'images/' + yolo_file.replace('.txt', '.jpg')  # image path
    txt = 'labels/' + yolo_file  # yolo label txt path
    xml_path = 'Annotations/' + yolo_file.replace('.txt', '.xml')  # xml save path
    print(jpg)
    print(txt)
    obj = ''

    img = cv2.imread(jpg)
    img_h, img_w = img.shape[0], img.shape[1]
    head = xml_head.format(str(jpg), str(img_w), str(img_h),str(3))
    with open(txt, 'r') as f:
        for line in f.readlines():
            yolo_datas = line.strip().split(' ')
            label = int(float(yolo_datas[0].strip()))
            center_x = round(float(str(yolo_datas[1]).strip()) * img_w)
            center_y = round(float(str(yolo_datas[2]).strip()) * img_h)
            bbox_width = round(float(str(yolo_datas[3]).strip()) * img_w)
            bbox_height = round(float(str(yolo_datas[4]).strip()) * img_h)

            xmin = str(int(center_x - bbox_width / 2))
            ymin = str(int(center_y - bbox_height / 2))
            xmax = str(int(center_x + bbox_width / 2))
            ymax = str(int(center_y + bbox_height / 2))

            obj += xml_obj.format(labels[label], xmin, ymin, xmax, ymax)
    with open(xml_path, 'w') as f_xml:
        f_xml.write(head + obj + xml_end)
    cnt += 1
    print(cnt)

下面代码是voc数据格式转yolo数据格式,即xml转txt

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

classes = ["smoke", "white","red"]


def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def convert_annotation(rootpath,xmlname):
    xmlpath = rootpath + '/labels'
    xmlfile = os.path.join(xmlpath,xmlname)
    with open(xmlfile, "r") as in_file:
      txtname = xmlname[:-4]+'.txt'
      txtpath = rootpath + '/labelYOLOs'
      if not os.path.exists(txtpath):
        os.makedirs(txtpath)
      txtfile = os.path.join(txtpath,txtname)
      with open(txtfile, "w+") as out_file:
        tree=ET.parse(in_file)
        root = tree.getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)
        out_file.truncate()
        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in classes or int(difficult)==1:
                continue
            cls_id = classes.index(cls)
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
            bb = convert((w,h), b)
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


if __name__ == "__main__":
    rootpath='/home/share/liubo/darknet/yanhuo/' # 根目录
    xmlpath=rootpath+'/labels/'
    list=os.listdir(xmlpath)
    for i in range(0,len(list)) :
        path = os.path.join(xmlpath,list[i])
        if ('.xml' in path)or('.XML' in path):
            convert_annotation(rootpath,list[i])
            print('done', i)
        else:
            print('not xml file',i)

猜你喜欢

转载自blog.csdn.net/qq_33193309/article/details/109055629