[目标检测模型辅助数据标注]YOLOv5检测图片并将结果转为voc的xml格式

1.调用YOLOv5检测模型对图片文件夹执行检测并保存txt和conf

python detect.py --weight runs\train\exp3\weights\best.pt --source [待检测的图片文件夹]--save-txt --save-conf

2.将检测结果转为voc的xml格式

import os 
import copy
from xml.dom import minidom
from tqdm import tqdm
import cv2
import xmltodict

# 图片\YOLO txt\xml对应的文件夹地址
base_dir = "[上面的待检测的图片文件夹]"
img_dir = os.path.join(base_dir,"img")
txt_dir = os.path.join(base_dir,"labels")
xml_dir =  os.path.join(base_dir,"xml")

class_name_list = "YOLO项目中data.yaml的class_names"

obj_base = {
      'name':'name',
      'bndbox':{
          'xmin':1,
          'ymin':1,
          'xmax':1,
          'ymax':1
      }
  }


xml_base = {
  'annotation': {
      'folder':'img',
      'filename':'',
      'size':{
          'width':1,
          'height':1,
          'depth':3
      },
      'object':[]
  }
}

img_name_list = os.listdir(img_dir)
pbar = tqdm(total=len(img_name_list))
pbar.set_description("YOLOTXT2VOC:")
for i in range(len(img_name_list)):
    # 拼接对应文件的地址
    img_name = img_name_list[i]
    img_path = os.path.join(img_dir,img_name)
    txt_path = os.path.join(txt_dir,img_name.split(".")[0]+".txt")
    xml_path = os.path.join(xml_dir,img_name.split(".")[0]+".xml")
    
     # 初始化xml文件对象
    xml_tmp = copy.deepcopy(xml_base)
    xml_tmp['annotation']['filename'] = img_name
    
    
    # 读取图片对应的宽高信息
    img = cv2.imread(img_path)
    img_height,img_width = img.shape[:2]
    
    # 读取txt文件内容
    with open(txt_path,'r') as f:
        content = f.read()
        
    # 逐行解析txt内容
    for line in content.split("\n"):
        if not line:continue
        data_item_list = line.split(" ")
        
        # 跳过类别置信度小于0.5的
        conf = float(data_item_list[5])
        if(conf<0.5):continue
        
        # 单行txt转为xml中对应的obj
        obj_tmp = copy.deepcopy(obj_base)
        obj_tmp["name"] = class_name_list[int(data_item_list[0])]
        x_center = int(float(data_item_list[1])*img_width)
        y_center = int(float(data_item_list[2])*img_height)
        width = int(float(data_item_list[3])*img_width)
        height = int(float(data_item_list[4])*img_height)
        obj_tmp["bndbox"]["xmin"] = int(x_center-width/2.0)
        obj_tmp["bndbox"]["ymin"] = int(y_center-height/2.0)
        obj_tmp["bndbox"]["xmax"] = width + obj_tmp["bndbox"]["xmin"]
        obj_tmp["bndbox"]["ymax"] = height + obj_tmp["bndbox"]["ymin"]
        
        xml_tmp['annotation']['object'].append(obj_tmp)
    
    # 转为xml并写入对应文件
    xmlstr = xmltodict.unparse(xml_tmp)
    xml = minidom.parseString(xmlstr)
    xml_pretty_str = xml.toprettyxml()
    with open(xml_path,"w") as f:
        f.write(xml_pretty_str)
    
    pbar.update(1)
pbar.close()