python工具方法 25 txt标注(yolo格式标注)的目标检测文件转voc数据

该代码来自于 卫星应用赛题——海上船舶智能检测比赛中某网友公开的数据处理代码，可以将txt描述（yolo格式标注）的目标检测数据转换为voc格式的数据。本章节对应的数据可视化如以下链接所示python工具方法 22 基于相对值描述的目标检测标注文件的可视化_a486259的博客-CSDN博客

1、创建基本目录

这里假设标注数据（txt和jpg）都存储在train目录下，以下命令的意思是，创建Images、txts两个目录，并将train目录下的jpg文件复制到Images目录下，将txt文件复制到txts下

mkdir Images
mkdir txts
cp -r train/*.jpg Images
cp -r train/*.txt txts

2、格式转换

以下代码默认图片的格式为jpg或png，如果图片实际后缀请自行修改一下X2VOC.convert方法里面的if函数。请按照个人的实际情况修改label_list，和img_h, img_w = 256, 256这两个关键代码。最后数据都会被保存到save_dir目录下。

import cv2
import os
import os.path as osp
import shutil


class X2VOC(object):
    def __init__(self):
        pass

    def convert(self, image_dir, txt_dir, dataset_save_dir):
        """转换。
        Args:
            image_dir (str): 图像文件存放的路径。
            txt_dir (str): 与每张图像对应的txt文件的存放路径。
            dataset_save_dir (str): 转换后数据集存放路径。
        """
        assert osp.exists(image_dir), "The image folder does not exist!"
        assert osp.exists(txt_dir), "The txt folder does not exist!"
        if not osp.exists(dataset_save_dir):
            os.makedirs(dataset_save_dir)
        # Convert the image files.
        new_image_dir = osp.join(dataset_save_dir, "JPEGImages")
        os.makedirs(new_image_dir, exist_ok=True)
        for img_name in os.listdir(image_dir):
            if ".jpg" in img_name or ".png" in img_name:
                shutil.copyfile(
                    osp.join(image_dir, img_name),
                    osp.join(new_image_dir, img_name))
        # Convert the txt files.
        xml_dir = osp.join(dataset_save_dir, "Annotations")
        os.makedirs(xml_dir, exist_ok=True)
        self.txt2xml(new_image_dir, txt_dir, xml_dir)


class Txt2VOC(X2VOC):
    """将使用LabelMe标注的数据集转换为VOC数据集。
    """

    def __init__(self):
        pass

    def txt2xml(self, image_dir, txt_dir, xml_dir):
        import xml.dom.minidom as minidom
        i = 0
        label_list=["类别1","类别2","类别3","类别4","类别5","类别6"]
        img_h, img_w = 256, 256
        for img_name in os.listdir(image_dir):
            img_name_part = osp.splitext(img_name)[0]
            txt_file = osp.join(txt_dir, img_name_part + ".txt")
            i += 1
            if not osp.exists(txt_file):
                os.remove(osp.join(image_dir, img_name))
                continue
            xml_doc = minidom.Document()
            root = xml_doc.createElement("annotation")
            xml_doc.appendChild(root)
            node_folder = xml_doc.createElement("folder")
            node_folder.appendChild(xml_doc.createTextNode("JPEGImages"))
            root.appendChild(node_folder)
            node_filename = xml_doc.createElement("filename")
            node_filename.appendChild(xml_doc.createTextNode(img_name))
            root.appendChild(node_filename)
            with open(txt_file, mode="r", \
                      encoding="utf-8") as j:
                
                node_size = xml_doc.createElement("size")
                node_width = xml_doc.createElement("width")
                node_width.appendChild(xml_doc.createTextNode(str(img_w)))
                node_size.appendChild(node_width)
                node_height = xml_doc.createElement("height")
                node_height.appendChild(xml_doc.createTextNode(str(img_h)))
                node_size.appendChild(node_height)
                node_depth = xml_doc.createElement("depth")
                node_depth.appendChild(xml_doc.createTextNode(str(3)))
                node_size.appendChild(node_depth)
                root.appendChild(node_size)

                txts = j.readlines()
                for line in txts:
                    box = line.split()
                    cls = int(box[0])
                    box = [float(i) for i in box[1:]]
                    x, y, w, h = box[0], box[1], box[2], box[3]
                    xmin = (x - w/2)*img_w
                    xmax = (x + w/2)*img_w
                    ymin = (y - h/2)*img_h
                    ymax = (y + h/2)*img_h
                    label = label_list[cls]
                    node_obj = xml_doc.createElement("object")
                    node_name = xml_doc.createElement("name")
                    node_name.appendChild(xml_doc.createTextNode(label))
                    node_obj.appendChild(node_name)
                    node_diff = xml_doc.createElement("difficult")
                    node_diff.appendChild(xml_doc.createTextNode(str(0)))
                    node_obj.appendChild(node_diff)
                    node_box = xml_doc.createElement("bndbox")
                    node_xmin = xml_doc.createElement("xmin")
                    node_xmin.appendChild(xml_doc.createTextNode(str(xmin)))
                    node_box.appendChild(node_xmin)
                    node_ymin = xml_doc.createElement("ymin")
                    node_ymin.appendChild(xml_doc.createTextNode(str(ymin)))
                    node_box.appendChild(node_ymin)
                    node_xmax = xml_doc.createElement("xmax")
                    node_xmax.appendChild(xml_doc.createTextNode(str(xmax)))
                    node_box.appendChild(node_xmax)
                    node_ymax = xml_doc.createElement("ymax")
                    node_ymax.appendChild(xml_doc.createTextNode(str(ymax)))
                    node_box.appendChild(node_ymax)
                    node_obj.appendChild(node_box)
                    root.appendChild(node_obj)
            with open(osp.join(xml_dir, img_name_part + ".xml"), 'w') as fxml:
                xml_doc.writexml(
                    fxml,
                    indent='\t',
                    addindent='\t',
                    newl='\n',
                    encoding="utf-8")


data_dir = './'
img_dir = osp.join(data_dir, 'Images')
txt_dir = osp.join(data_dir, 'txts')
save_dir="voc_data"
dataset_save_dir = osp.join(data_dir, save_dir)
txt2voc = Txt2VOC()
txt2voc.convert(img_dir, txt_dir, dataset_save_dir)