参考链接

对里面的txt2xml的代码进行了改变,写成以下代码:

  1. 需要改变第5行代码的labels
  2. 改变65~67行为自己的文件夹地址
import cv2
import os
from tqdm import tqdm

labels = ['nest', 'kite', 'balloon', 'trash']  # 数据集类别名

xml_head = """<annotation>
	<folder>images</folder>
	<filename>{}</filename>
	<path>{}</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>{}</width>
		<height>{}</height>
		<depth>{}</depth>
	</size>
	<segmented>0</segmented>"""

xml_obj = """
	<object>
		<name>{}</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>{}</xmin>
			<ymin>{}</ymin>
			<xmax>{}</xmax>
			<ymax>{}</ymax>
		</bndbox>
	</object>"""

xml_end = """
</annotation>"""


def convert_txt2xml(file, imagepath, txtpath, xmlpath):
    obj = ''

    img = cv2.imread(imagepath)
    img_h, img_w, img_depth = img.shape
    head = xml_head.format(file, imagepath, str(img_w), str(img_h), str(img_depth))
    with open(txtpath, '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(xmlpath, 'w') as f_xml:
        f_xml.write(head + obj + xml_end)


if __name__ == '__main__':
    images_rootdir = r'G:\pycharmprojects\yoloair-iscyy-beta\datasets\A_linepaste_trash_png_2\images'
    labels_rootdir = r'G:\pycharmprojects\yoloair-iscyy-beta\datasets\A_linepaste_trash_png_2\labels'
    save_xmlsrootdir = r'G:\pycharmprojects\yoloair-iscyy-beta\datasets\A_linepaste_trash_png_2\Annotations_xml'  # save_xmlsrootdir

    if not os.path.exists(save_xmlsrootdir):
        os.makedirs(save_xmlsrootdir)

    testset = os.listdir(images_rootdir)  # 我的images_rootdir下又两个子文件夹:train val
    for test in testset:
        images_dir = os.path.join(images_rootdir, test)  # png 图像文件
        labels_dir = os.path.join(labels_rootdir, test)  # txt 标签文件

        filelist = os.listdir(images_dir)
        for file in tqdm(filelist):
            filename = os.path.splitext(file)[0]
            imagepath = os.path.join(images_dir, file)
            txtpath = os.path.join(labels_dir, filename + '.txt')
            xmlpath = os.path.join(save_xmlsrootdir, filename + '.xml')

            convert_txt2xml(file, imagepath, txtpath, xmlpath)
Logo

技术共进,成长同行——讯飞AI开发者社区

更多推荐