import os
import cv2
import pathlib
from xml.dom.minidom import Document
# 支持多种格式图片,可以自己添加
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):
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/liu_hui/YM-factory02/YM-factory02/data_door_white/images/test"
src_label_root = "/home/liu_hui/YM-factory02/YM-factory02/runs/detect/exp3/labels"
dst_label_root = "/home/liu_hui/YM-factory02/YM-factory02/runs/detect/exp3/labels/voc"
# txt文件中每列第一个为类别序号,分别对应label名
class_map = {'0': "door_open",'1':"door_close"} #, '1': "person", '2': 'ce'
#class_map = {'0': "badminton, '1': "person", '2': 'ce'"}
yolo2voc(src_img_root, src_label_root, dst_label_root, class_map)