（19）目标检测算法之模型预测并保存结果到xml中，自动生成标签文件/自动标注

目标检测算法之模型预测并保存结果到xml中，自动生成标签文件/自动标注

• 在做深度学习算法研究中，标图往往是工作量很大且性价比较低的一项工作了，最近在用YOLO，在此参考网上已有程序完善下可执行demo
• 看过一下demo，需要导入utils下的模块，我没有安装成功，在此特别感谢这篇文章：YOLOv5实现半标注—告别大量重复标注工作，本文是在其代码上做了些更改

1. 自动生成标签实现

• 具体过程的话呢需要一些几点：

  • 1. 对少些图像进行标注并训练，保存模型权重文件
  • 2. 封装预测模块，并将预测结果整理成代码中所需要的格式：由[x_min,y_min,x_max,y_max,label[x_min,y_min,x_max,y_max,label]组成的列表
  • 3. 修改测试路径，运行自动检测demo，批量生成标签

• 我用的是YOLO，如果你的是其他模型，自行做出相应修改

• 完整代码如下：

import os
import cv2
from lxml.etree import Element, SubElement, tostring
import detect_yolov5 as detect

classes=['desk','packet','cue']


def create_xml(list_xml,list_images,xml_path):
    """
    创建xml文件，依次写入xml文件必备关键字
    :param list_xml:   txt文件中的box
    :param list_images:   图片信息，xml中需要写入WHC
    :return:
    """
    node_root = Element('annotation')
    node_folder = SubElement(node_root, 'folder')
    node_folder.text = 'Images'
    node_filename = SubElement(node_root, 'filename')
    node_filename.text = str(list_images[3])
    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_width.text = str(list_images[1])
    node_height = SubElement(node_size, 'height')
    node_height.text = str(list_images[0])
    node_depth = SubElement(node_size, 'depth')
    node_depth.text = str(list_images[2])

    if len(list_xml)>=1:        # 循环写入box
        for list_ in list_xml:
            node_object = SubElement(node_root, 'object')
            node_name = SubElement(node_object, 'name')
            # if str(list_[4]) == "person":                # 根据条件筛选需要标注的标签,例如这里只标记person这类，不符合则直接跳过
            #     node_name.text = str(list_[4])
            # else:
            #     continue
            node_name.text = str(list_[4])
            node_difficult = SubElement(node_object, 'difficult')
            node_difficult.text = '0'
            node_bndbox = SubElement(node_object, 'bndbox')
            node_xmin = SubElement(node_bndbox, 'xmin')
            node_xmin.text = str(list_[0])
            node_ymin = SubElement(node_bndbox, 'ymin')
            node_ymin.text = str(list_[1])
            node_xmax = SubElement(node_bndbox, 'xmax')
            node_xmax.text = str(list_[2])
            node_ymax = SubElement(node_bndbox, 'ymax')
            node_ymax.text = str(list_[3])

    xml = tostring(node_root, pretty_print=True)   # 格式化显示，该换行的换行

    file_name = list_images[3].split(".")[0]
    filename = xml_path+"/{}.xml".format(file_name)

    f = open(filename, "wb")
    f.write(xml)
    f.close()


if __name__ == '__main__':

    path = r"./mytrain/images"        # 图片路径
    xml_path = r"mytrain/images"      # xml标注保存路径

    for name in os.listdir(path):
        print(name)
        #xml_name=name.split('.')[0]+".xml"

        if(name.split('.')[-1]=='jpg'):
            image = cv2.imread(os.path.join(path, name))
            list_image = (image.shape[0], image.shape[1], image.shape[2], name)  # 图片的宽高等信息

            result = detect.detect(image)

            xyxy_list = []
            for res in result:
                x_min = res['position'][0]
                y_min = res['position'][1]
                x_max = res['position'][0] + res['position'][2]
                y_max = res['position'][1] + res['position'][3]
                name = res['class']
                _list = [x_min, y_min, x_max, y_max, name]
                xyxy_list.append(_list)

            create_xml(xyxy_list, list_image, xml_path)  # 生成标注的xml文件

2. 批量复制xml

• 此方法针对背景相同且图像中多数待标注对象相同时使用
• 输入待标注图像名称及已标注.xml，即可批量复制并保存成对应图像名字的标签文件.xml
• 完成代码：

#批量赋值图像的标注文件.xml并重命名为图像名字

import os
import shutil

# imgspath="./file3/jpg"
# xmlpath="file3/jpg/005.xml"
imgspath="./imgs"
xmlpath="./imgs/0.xml"

for root,dirs,files in os.walk(imgspath,topdown=True):
    for name in files:
        str=name.split('.')
        if(str[-1] == 'jpg'):
            newname=str[0]+'.xml'
            newxml=os.path.join(root,newname)
            print(newxml)
            if not os.path.exists(newxml):
                shutil.copy(xmlpath,newxml)

3. 读取文件夹图像并批量重命名

针对：

• 有重复名字的数据集
• 名字特别复杂的数据集

代码如下：

'''
读取文件夹路径下图像并且批量重命名
'''

# -*- coding:UTF-8 -*-
import os
import cv2
import random

def traverse_path(file_path):
    files = os.listdir(file_path)
    for fi in files:
        fi_d = os.path.join(file_path, fi)
        if os.path.isdir(fi_d):

            traverse_path(fi_d)
        else:
            img_name=os.path.join(file_path, fi)
            #print(img_name)
            if img_name[-4:]=='tiff':
                #print(img_name)
                imglist.append(img_name)
                txt_file.write(img_name)
                txt_file.write('
')
            elif img_name[-3:]=="jpg":
                # print(img_name)
                imglist.append(img_name)
                txt_file.write(img_name)
                txt_file.write('
')

    # for i in len(lists):
    #     print(lists(i))
    return imglist


if __name__ == '__main__':
    dirname = "./file3"
    txt_path = './list.txt'
    savepath="./png"
    if not os.path.exists(savepath):
        os.makedirs(savepath)
    txt_file = open(txt_path, 'w')
    imglist = []

    traverse_path(dirname)
    print("读取图像数量：", len(imglist))

    #批量重命名保存
    n=0
    t=0
    while t< len(imglist):
        img=cv2.imread(imglist[t])
        # newname=imglist[t].rsplit("\",1)[0] + "\" + str(n)+".png"
        newname = savepath+ "\" + str(n) + ".png"

        #重命名保存
        cv2.imwrite(newname,img)

        print(newname)
        t+=1
        n+=1