目标检测 -- 基于训练好的YOLO深度学习网络模型实现自动辅助标注

在实际工程中，我们一般会先用少量的数据集进行人工标注，然后训练一个简单的深度学习网络模型，然后基于该网络模型进行预测，对剩余大量测试图片进行初步标注，然后基于机器标注之后的结果进行人工微调标注框，这样即省时又省力。这篇blog主要就讲讲基于训练好的YOLO v5模型实现自动辅助标注。

1. 人工标注少量数据集
2. 基于上述标注好的数据集，使用YOLO V5框架进行训练
3. 基于训练好的网络权重文件进行对未标注的数据进行预测，得到txt预测结果文件
4. 代码实现txt2xml
5. 对生成的xml进行人工微调

针对上述总体思路，前两个不比具体展开，网上很多对yolo v5训练的教程，而且很简单。对于第三步，这里我使用的是官方的YOLO V5 pytorch v6版本，在detect.py中有一个参数parser.add_argument('--save-txt', action='store_true', default=True, help='save results to *.txt')，这里加上default=True(貌似不用设置，官方提供的代码就是默认保存txt，以防万一，最好加上吧~)。
至此，我们已经得到了基于训练好的模型对待标注数据集的初步测试。接下来就是实现txt2xml，可以直接按照我的方法。

代码：

# .txt-->.xml
# ! /usr/bin/python
# -*- coding:UTF-8 -*-
import os
import cv2

def txt_to_xml(txt_path, img_path, xml_path):
    # 字典对标签中的类别进行转换
    dict = {'0': "person"
            }
    files = os.listdir(txt_path)
    pre_img_name = ''
    for i, name in enumerate(files):
        if name == "desktop.ini":
            continue
        print(name)
        txtFile = open(txt_path + name)
        txtList = txtFile.readlines()
        img_name = name.split(".")[0]
        pic = cv2.imread(img_path + img_name + ".jpg")
        Pheight, Pwidth, Pdepth = pic.shape
        for row in txtList:
            oneline = row.strip().split(" ")
            if img_name != pre_img_name:
                xml_file = open((xml_path + img_name + '.xml'), 'w')
                xml_file.write('n')
                xml_file.write('    VOC2007n')
                xml_file.write('    ' + img_name + '.jpg' + 'n')
                xml_file.write('    n')
                xml_file.write('        orgaquantn')
                xml_file.write('        organoidsn')
                xml_file.write('    n')
                xml_file.write('    n')
                xml_file.write('        ' + str(Pwidth) + 'n')
                xml_file.write('        ' + str(Pheight) + 'n')
                xml_file.write('        ' + str(Pdepth) + 'n')
                xml_file.write('    n')
                xml_file.write('    n')
                xml_file.close()
                pre_img_name = img_name 
            else:  
                xml_file = open((xml_path + img_name + '.xml'), 'a')
                xml_file.write('    n')
                xml_file.close()

        xml_file1 = open((xml_path + pre_img_name + '.xml'), 'a')
        xml_file1.write('')
        xml_file1.close()
    print("Done !")


# 修改成自己的文件夹 注意文件夹最后要加上/
txt_to_xml(r"", #txt路径
           r"", #图片路径
           r"") #生成的xml路径

这里，需要修改四个地方：txt路径、图片路径、生成的xml路径、类别（在函数第一行）。运行之后，会得到对应的xml文件。然后再进行人工微调就完成了自动标注。

1.上述代码有一个小bug，就是如果图片名字中有.号，会报错。
2.我这里是使用YOLO V5 pytorch v6版本，其实所有的模型都是可以这么做的，只要保证测试的时候txt能够保存就能进行人工标注。