YOLO格式转VOC格式 支持png jpeg等图片格式_yolov7对输入图片格式的要求jpg和bmp

# -*-  coding = utf-8 -*-

 
import os
import cv2
import pathlib
from xml.dom.minidom import Document
 
# 支持多种格式图片,可以自己添加
# 这么做是为了适应不同文件名，比如 'xxx.xxx xx.jpeg'等
def search_file_name(file_path,img_root):
    file_format = ['.jpg','.JPG','.png','.PNG','.JPEG','.jpeg']
    file_path = pathlib.Path(file_path)
    file_name = file_path.stem
    for endsw in file_format:
        file_judge = file_name + endsw
        if os.path.isfile(os.path.join(img_root,file_judge)):
            return file_name,endsw
    return 'stop',0
 
def yolo2voc(src_img_root, src_label_root, dst_label_root,class_map):
    """
    src_img_root:图片路径（需要获取图片shape）
    src_label_root：输入label路径（YOLO的txt格式）
    dst_label_root：输出label路径（PASCAL VOC的xml格式）
    target：讲txt格式转为xml格式
    """
    not_have_img = []
    if not os.path.exists(dst_label_root):
        os.makedirs(dst_label_root)
    # 遍历所有txt文件
    for i, src_label_name in enumerate(os.listdir(src_label_root)):
        xmlBuilder = Document()
        annotation = xmlBuilder.createElement("annotation")  # 创建annotation标签
        xmlBuilder.appendChild(annotation)
        src_label_path = os.path.join(src_label_root,src_label_name)
        file_name,endsw = search_file_name(src_label_path,src_img_root)
        if file_name == 'stop':
            not_have_img.append(os.path.basename(src_label_path))
            continue
        with open(src_label_path, 'r') as fr:
            txtlines = fr.readlines()
        src_img_name = file_name + endsw
        src_img_path = os.path.join(src_img_root,src_img_name)
        image_np = cv2.imread(src_img_path)
        Pheight, Pwidth, Pdepth = image_np.shape
        folder = xmlBuilder.createElement("folder")  # folder标签
        foldercontent = xmlBuilder.createTextNode("driving_annotation_dataset")
        folder.appendChild(foldercontent)
        annotation.appendChild(folder)  # folder标签结束
 
        filename = xmlBuilder.createElement("filename")  # filename标签
        filenamecontent = xmlBuilder.createTextNode(src_img_name)
        filename.appendChild(filenamecontent)
        annotation.appendChild(filename)  # filename标签结束
 
        size = xmlBuilder.createElement("size")  # size标签
        width = xmlBuilder.createElement("width")  # size子标签width
        widthcontent = xmlBuilder.createTextNode(str(Pwidth))
        width.appendChild(widthcontent)
        size.appendChild(width)  # size子标签width结束
 
        height = xmlBuilder.createElement("height")  # size子标签height
        heightcontent = xmlBuilder.createTextNode(str(Pheight))
        height.appendChild(heightcontent)
        size.appendChild(height)  # size子标签height结束
 
        depth = xmlBuilder.createElement("depth")  # size子标签depth
        depthcontent = xmlBuilder.createTextNode(str(Pdepth))
        depth.appendChild(depthcontent)
        size.appendChild(depth)  # size子标签depth结束
        annotation.appendChild(size)  # size标签结束
        for line in txtlines:
            oneline = line.strip().split(" ")
            object = xmlBuilder.createElement("object")  # object 标签
            picname = xmlBuilder.createElement("name")  # name标签
            namecontent = xmlBuilder.createTextNode(class_map[oneline[0]])
            picname.appendChild(namecontent)
            object.appendChild(picname)  # name标签结束
 
            pose = xmlBuilder.createElement("pose")  # pose标签
            posecontent = xmlBuilder.createTextNode("Unspecified")
            pose.appendChild(posecontent)
            object.appendChild(pose)  # pose标签结束
 
            truncated = xmlBuilder.createElement("truncated")  # truncated标签
            truncatedContent = xmlBuilder.createTextNode("0")
            truncated.appendChild(truncatedContent)
            object.appendChild(truncated)  # truncated标签结束
 
            difficult = xmlBuilder.createElement("difficult")  # difficult标签
            difficultcontent = xmlBuilder.createTextNode("0")
            difficult.appendChild(difficultcontent)
            object.appendChild(difficult)  # difficult标签结束
 
            bndbox = xmlBuilder.createElement("bndbox")  # bndbox标签
            xmin = xmlBuilder.createElement("xmin")  # xmin标签
            mathData = int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth)
            xminContent = xmlBuilder.createTextNode(str(mathData))
            xmin.appendChild(xminContent)
            bndbox.appendChild(xmin)  # xmin标签结束
 
            ymin = xmlBuilder.createElement("ymin")  # ymin标签
            mathData = int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight)
            yminContent = xmlBuilder.createTextNode(str(mathData))
            ymin.appendChild(yminContent)
            bndbox.appendChild(ymin)  # ymin标签结束
 
            xmax = xmlBuilder.createElement("xmax")  # xmax标签
            mathData = int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth)
            xmaxContent = xmlBuilder.createTextNode(str(mathData))
            xmax.appendChild(xmaxContent)
            bndbox.appendChild(xmax)  # xmax标签结束
 
            ymax = xmlBuilder.createElement("ymax")  # ymax标签
            mathData = int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight)
            ymaxContent = xmlBuilder.createTextNode(str(mathData))
            ymax.appendChild(ymaxContent)
            bndbox.appendChild(ymax)  # ymax标签结束
            object.appendChild(bndbox)  # bndbox标签结束
            annotation.appendChild(object)  # object标签结束
        dst_label_path = os.path.join(dst_label_root,file_name + '.xml')
        with open(dst_label_path,'w') as f:
            xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
        print(src_img_path)
    print('not have img:',not_have_img,'\n',not_have_img.__len__())
 
 
if __name__ == "__main__":
    #注意： opencv 不支持中文路径，路径及文件名需要是英文数字符号等
    src_img_root = "/home/fnce"
    src_label_root = "/home/label"
    dst_label_root = "/home/dst_label"
    # txt文件中每列第一个为类别序号，分别对应label名
    class_map = {'0': "cross",'1': "person",'2':'ce'}
    yolo2voc(src_img_root, src_label_root, dst_label_root,class_map)
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137