数据集在如下地址下载,yolov5的github下载的有个别数据对不上:

链接: https://pan.baidu.com/s/1taUkSaWoEh_26uwb8m7Ebg 提取码: 4h6p 
--来自百度网盘超级会员v5的分享

处理脚本如下:

import os
import xml.etree.ElementTree as ET
from PIL import Image

# 检测结果中的类别列表
classes = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
           "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
           "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
           "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
           "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
           "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
           "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
           "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
           "hair drier", "toothbrush"]


def convert_txt_to_xml(txt_file):
    # 读取txt文件中每个检测结果,并创建XML标签
    root = ET.Element("annotation")
    img_folder = "images"
    # 获取图像文件路径和宽度、高度、深度
    img_path = os.path.join(img_folder, os.path.splitext(os.path.basename(txt_file))[0] + ".jpg")
    img = Image.open(img_path)
    width, height = img.size
    depth = 3
    path = os.path.abspath(img_path)

    # 添加各种标签信息
    ET.SubElement(root, "folder").text = img_folder
    ET.SubElement(root, "filename").text = os.path.basename(img_path)
    ET.SubElement(root, "path").text = path
    source_tag = ET.SubElement(root, "source")
    ET.SubElement(source_tag, "database").text = "Unknown"
    size_tag = ET.SubElement(root, "size")
    ET.SubElement(size_tag, "width").text = str(width)
    ET.SubElement(size_tag, "height").text = str(height)
    ET.SubElement(size_tag, "depth").text = str(depth)
    ET.SubElement(root, "segmented").text = "0"

    with open(txt_file) as f:
        for line in f:
            parts = line.strip().split()
            class_id = int(parts[0])
            x_center = float(parts[1])
            y_center = float(parts[2])
            w = float(parts[3])
            h = float(parts[4])

            x_min = (x_center - (w / 2)) * width
            y_min = (y_center - (h / 2)) * height
            x_max = (x_center + (w / 2)) * width
            y_max = (y_center + (h / 2)) * height

            object_tag = ET.SubElement(root, "object")
            ET.SubElement(object_tag, "name").text = classes[class_id]
            ET.SubElement(object_tag, "pose").text = "Unspecified"
            ET.SubElement(object_tag, "truncated").text = "1"
            ET.SubElement(object_tag, "difficult").text = "0"
            bndbox_tag = ET.SubElement(object_tag, "bndbox")

            ET.SubElement(bndbox_tag, "xmin").text = str(int(x_min))
            ET.SubElement(bndbox_tag, "ymin").text = str(int(y_min))
            ET.SubElement(bndbox_tag, "xmax").text = str(int(x_max))
            ET.SubElement(bndbox_tag, "ymax").text = str(int(y_max))

    # 创建XML文件
    xml_dir = "annotations"
    if not os.path.exists(xml_dir):
        os.makedirs(xml_dir)
    xml_file = os.path.join(xml_dir, os.path.splitext(os.path.basename(txt_file))[0] + ".xml")
    tree = ET.ElementTree(root)
    tree.write(xml_file)


if __name__ == "__main__":
    # 设置目录路径和文件扩展名
    dir_path = "txt"
    ext = ".txt"

    # 遍历目录中的文件,并对扩展名为.txt的文件进行转换
    for file_name in os.listdir(dir_path):
        if file_name.endswith(ext):
            txt_file = os.path.join(dir_path, file_name)
            convert_txt_to_xml(txt_file)
07-25 16:32