# -*- 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)
YOLO格式转VOC格式 支持png jpeg等图片格式
最新推荐文章于 2024-03-29 17:35:16 发布