两点xml

最新推荐文章于 2024-05-17 16:53:11 发布

373023820

最新推荐文章于 2024-05-17 16:53:11 发布

阅读量26

点赞数

文章标签： xml 深度学习目标检测

本文链接：https://blog.csdn.net/qq_53821866/article/details/134382039

版权

import os
import xml.etree.ElementTree as ET

import cv2

classes = ['holothurian','echinus','scallop','starfish','waterweeds']  # 标注时的标签
def convert_xml_to_yolo(xml_path, image_width, image_height):
    tree = ET.parse(xml_path)
    root = tree.getroot()

    yolo_labels = []

    for object_elem in root.findall('object'):
        name = object_elem.find('name').text
        class_index = classes.index(name)
        xmin = int(object_elem.find('bndbox/xmin').text)
        ymin = int(object_elem.find('bndbox/ymin').text)
        xmax = int(object_elem.find('bndbox/xmax').text)
        ymax = int(object_elem.find('bndbox/ymax').text)

        # Convert bounding box to YOLO format (normalized x,y,w,h)
        x = (xmin + xmax) / 2 / image_width
        y = (ymin + ymax) / 2 / image_height
        width = (xmax - xmin) / image_width
        height = (ymax - ymin) / image_height

        yolo_label = f"{class_index} {x} {y} {width} {height}"
        yolo_labels.append(yolo_label)

    return yolo_labels

def convert_dataset_to_yolo(xml_dir, image_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    for xml_file in os.listdir(xml_dir):
        if xml_file.endswith('.xml'):
            xml_path = os.path.join(xml_dir, xml_file)
            image_name = os.path.splitext(xml_file)[0] + '.jpg'
            image_path = os.path.join(image_dir, image_name)

            image_width, image_height = get_image_size(image_path)
            yolo_labels = convert_xml_to_yolo(xml_path, image_width, image_height)

            yolo_file = os.path.join(output_dir, os.path.splitext(xml_file)[0] + '.txt')
            with open(yolo_file, 'w') as f:
                f.write('\n'.join(yolo_labels))

# def get_image_size(image_path):
#     # Add code to get image width and height
#     # You can use libraries like OpenCV or PIL for this task
#     width = ...  # Get image width
#     height = ...  # Get image height
#     return width, height
def get_image_size(image_path):
    image = cv2.imread(image_path)
    height, width, _ = image.shape
    return width, height
# Usage
xml_dir = r'E:\URPC2020\train\box'
image_dir = r'E:\URPC2020\train\image'
output_dir = r'E:\URPC2020\train\qwe'

convert_dataset_to_yolo(xml_dir, image_dir, output_dir)