目标检测自动标注python脚本

在需要大量标注数据时，可以先用少量样本训练一个模型，在模型大致可以正常输出类别和位置坐标的情况下进行自动标注，后续再进行少量的人工调整，可以节省很大的精力（自动标注的格式仿照labelimg标注的VOC格式）：

import shutil
import os
from xml.etree import cElementTree as ET
from PIL import Image
import cv2


def write_xml(img_path, boxes):
    '''
    img_path :  图片的绝对路径
    boxes   :   图像的预测结果 [[cls,x1,y1,x2,y2],
                             [cls,x1,y1,x2,y2],
                             ................]
    '''

    img = cv2.imread(img_path)
    h, w, d = img.shape
    img_name = img_path.split("/")[-1:][0]
    file_dir = img_path.split("/")[-2]                       #图片存放文件夹的名称
    xml_dir = os.path.join(os.path.dirname(img_path),"xml_labels")

    if os.path.exists(xml_dir)==False:
        os.mkdir(xml_dir)
    xml_path = os.path.join(xml_dir, img_name.split(".")[0]+".xml" )       #需要写入的xml的绝对路径

    doc = ET.Element('annotation')                          # 树根
    folder = ET.SubElement(doc, "folder")                   # 文件夹
    folder.text = file_dir                                  # 文件夹--------填值
    filename = ET.SubElement(doc, "filename")               # 文件名
    filename.text = img_path.split("/")[-1:][0]             # 文件名--------填值

    path = ET.SubElement(doc, 'path')                       # 绝对路径
    path.text = img_path                                    # 绝对路径------填值
    source = ET.SubElement(doc, "source")                   # 来源
    database = ET.SubElement(source, "database")            # 数据库
    database.text = str('Unknown')                          # 数据库--------填值      --默认
    size = ET.SubElement(doc, 'size')                       # 图片大小
    width = ET.SubElement(size, 'width')                    # 图片的宽度
    width.text = str(w)                                     # 图片宽度-----填值
    height = ET.SubElement(size, 'height')                  # 图片高度
    height.text = str(h)                                    # 图片高度-----填值
    depth = ET.SubElement(size, 'depth')                    # 通道数
    depth.text = str(d)                                     # 通道数-------填值
    segmented = ET.SubElement(doc, "segmented")             # 分割
    segmented.text = str(0)                                 # 分割--------填值
    for i,box in enumerate(boxes):
        cls,x1,y1,x2,y2 = box
        object = ET.SubElement(doc, "object")               # 物体
        name = ET.SubElement(object, 'name')                # 物体名称
        name.text = cls                                     # 物体名称-----填值
        pose = ET.SubElement(object, "pose")                # 姿势
        pose.text = str("Unspecified")                      # 姿势---------填值      --默认
        truncated = ET.SubElement(object, "truncated")      # 截断
        truncated.text = str('0')                           # 截断---------填值      --默认
        difficult = ET.SubElement(object, "difficult")      # 困难样本
        difficult.text = str('0')                           # 困难样本------填值      --默认
        bndbox = ET.SubElement(object, 'bndbox')            # 矩形框
        xmin = ET.SubElement(bndbox, 'xmin')                # x1
        xmin.text = str(x1)                                 # x1-----------填值
        ymin = ET.SubElement(bndbox, 'ymin')                # y1
        ymin.text = str(y1)                                 # y1-----------填值
        xmax = ET.SubElement(bndbox, 'xmax')                # x2
        xmax.text = str(x2)                                 # x2-----------填值
        ymax = ET.SubElement(bndbox, 'ymax')                # y2
        ymax.text = str(y2)                                 # y2-----------填值

        tree = ET.ElementTree(doc)
        tree.write(xml_path, encoding='utf-8')
    print("[{}] -----------------------------finished ！！ !".format(xml_path))

if __name__ == '__main__':

    img_path = "1.jpg"          
    boxes = [['cls1', 23, 46, 90, 120], ['cls2', 60, 100, 300, 200]]
    write_xml(img_path,boxes)                   #传入预测结果和图片的绝对路径