参考链接
对里面的txt2xml的代码进行了改变,写成以下代码:
- 需要改变第5行代码的
labels
- 改变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)