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

目标检测 – 基于训练好的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，可以直接按照我的方法。

2.1 实现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('<annotation>\n')
                xml_file.write('    <folder>VOC2007</folder>\n')
                xml_file.write('    <filename>' + img_name + '.jpg' + '</filename>\n')
                xml_file.write('    <source>\n')
                xml_file.write('        <database>orgaquant</database>\n')
                xml_file.write('        <annotation>organoids</annotation>\n')
                xml_file.write('    </source>\n')
                xml_file.write('    <size>\n')
                xml_file.write('        <width>' + str(Pwidth) + '</width>\n')
                xml_file.write('        <height>' + str(Pheight) + '</height>\n')
                xml_file.write('        <depth>' + str(Pdepth) + '</depth>\n')
                xml_file.write('    </size>\n')
                xml_file.write('    <object>\n')
                xml_file.write('        <name>' + dict[oneline[0]] + '</name>\n')
                xml_file.write('        <bndbox>\n')
                xml_file.write('            <xmin>' + str(
                    int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth)) + '</xmin>\n')
                xml_file.write('            <ymin>' + str(
                    int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight)) + '</ymin>\n')
                xml_file.write('            <xmax>' + str(
                    int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth)) + '</xmax>\n')
                xml_file.write('            <ymax>' + str(
                    int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight)) + '</ymax>\n')
                xml_file.write('        </bndbox>\n')
                xml_file.write('    </object>\n')
                xml_file.close()
                pre_img_name = img_name 
            else:  
                xml_file = open((xml_path + img_name + '.xml'), 'a')
                xml_file.write('    <object>\n')
                xml_file.write('        <name>' + dict[oneline[0]] + '</name>\n')
                '''  按需添加这里和上面
                xml_file.write('        <pose>Unspecified</pose>\n')
                xml_file.write('        <truncated>0</truncated>\n')
                xml_file.write('        <difficult>0</difficult>\n')
                '''
                xml_file.write('        <bndbox>\n')
                xml_file.write('            <xmin>' + str(
                    int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth)) + '</xmin>\n')
                xml_file.write('            <ymin>' + str(
                    int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight)) + '</ymin>\n')
                xml_file.write('            <xmax>' + str(
                    int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth)) + '</xmax>\n')
                xml_file.write('            <ymax>' + str(
                    int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight)) + '</ymax>\n')
                xml_file.write('        </bndbox>\n')
                xml_file.write('    </object>\n')
                xml_file.close()

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


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

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

三、总结

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