yolov5中获取xml文件中的数据存入到txt中,object有几个就填写几行

yolov5中获取xml文件中的数据存入到txt中,object有几个就填写几行

import xml.etree.ElementTree as ET
from pathlib import Path

def save_label_file(img_name, img_size, img_boxes):
    label_filename = img_name.split('.')[0] + ".txt"
    with open(label_filename, 'w') as f:
        for box in img_boxes:
            cls_name = box[0]
            x1 = box[1] / img_size[0]
            y1 = box[2] / img_size[1]
            x2 = box[3] / img_size[0]
            y2 = box[4] / img_size[1]
            line = f"{cls_name} {x1} {y1} {x2} {y2}\n"
            f.write(line)

def get_xml_data(img_xml_file):
    tree = ET.parse(str(img_xml_file))
    root = tree.getroot()

    img_name = root.find("filename").text
    img_size = root.find("size")
    img_w = int(img_size.find("width").text)
    img_h = int(img_size.find("height").text)

    img_boxes = []
    for obj in root.findall("object"):
        cls_name = obj.find("name").text
        bbox = obj.find("bndbox")
        x1 = int(bbox.find("xmin").text)
        y1 = int(bbox.find("ymin").text)
        x2 = int(bbox.find("xmax").text)
        y2 = int(bbox.find("ymax").text)
        img_boxes.append([cls_name, x1, y1, x2, y2])

    save_label_file(img_name, [img_w, img_h], img_boxes)

# 调用示例
get_xml_data('path_to_xml_file.xml')

原始数据展示:

import xml.etree.cElementTree as ET

def get_xml_data(img_xml_file):
    dom = parse(str(img_xml_file))
    xml_root = dom.documentElement
    img_name = xml_root.getElementsByTagName("filename")[0].childNodes[0].data
    img_size = xml_root.getElementsByTagName("size")[0]
    img_w = img_size.getElementsByTagName("width")[0].childNodes[0].data
    img_h = img_size.getElementsByTagName("height")[0].childNodes[0].data
    img_c = img_size.getElementsByTagName("depth")[0].childNodes[0].data
    img_boxes = []
    for obj in xml_root.getElementsByTagName("object"):
        cls_name = obj.getElementsByTagName("name")[0].childNodes[0].data
        bbox = obj.getElementsByTagName("bndbox")[0]
        x1 = int(bbox.getElementsByTagName("xmin")[0].childNodes[0].data)
        y1 = int(bbox.getElementsByTagName("ymin")[0].childNodes[0].data)
        x2 = int(bbox.getElementsByTagName("xmax")[0].childNodes[0].data)
        y2 = int(bbox.getElementsByTagName("ymax")[0].childNodes[0].data)
        img_boxes.append([cls_name, x1, y1, x2, y2])
    
    save_label_file(img_name, [img_w, img_h], img_boxes)
import xml.etree.ElementTree as ET
from pathlib import Path

def get_xml_data(img_xml_file):
    tree = ET.parse(str(img_xml_file))
    root = tree.getroot()
    
    img_name = root.find("filename").text
    img_size = root.find("size")
    img_w = int(img_size.find("width").text)
    img_h = int(img_size.find("height").text)
    img_c = int(img_size.find("depth").text)
    
    img_boxes = []
    for obj in root.findall("object"):
        cls_name = obj.find("name").text
        bbox = obj.find("bndbox")
        x1 = int(bbox.find("xmin").text)
        y1 = int(bbox.find("ymin").text)
        x2 = int(bbox.find("xmax").text)
        y2 = int(bbox.find("ymax").text)
        img_boxes.append([cls_name, x1, y1, x2, y2])
    
    save_label_file(img_name, [img_w, img_h], img_boxes)

你可能感兴趣的:(python,YOLO,xml,前端)