yolov3预训练模型自动标注

序：想要真正准确的的自动标注，的确不太现实，都能准确的自动标注了，还训练模型干嘛！

所以本文所写方法是小量数据集预训练模型后，自动打标最后微调。

 

 

 （上图是我的文件夹格式，将自己预训练后的模型放到指定位置）

代码包含调用yolo模型。废话不多说！

# coding=utf-8
'''
author : Helen
date : 2020-11-12 16:15
'''

import cv2
import numpy as np
import os
import xml.dom.minidom
# import pypinyin
import time


def ModelYoloV3(frame, confidence=0.5, threshold=0.4):
    # 加载文件路径
    yolo_dir = './previousTrainModel'  # YOLO文件路径
    weightsPath = os.path.join(yolo_dir, 'myVLPChar_yolov3_21000.weights')  # 权重文件
    configPath = os.path.join(yolo_dir, 'myVLPChar_yolov3.cfg')  # 配置文件
    labelsPath = os.path.join(yolo_dir, 'myVLPCharData.names')  # label名称

    print("[INFO] loading YOLO from disk...")  # # 可以打印下信息

    # 加载网络、配置权重
    net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)  # #  利用下载的文件
    # net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
    # net.setPreferableTarget(cv.dnn.DNN_TARGET_CPU)

    CONFIDENCE = confidence  # 过滤弱检测的最小概率,默认0.5
    THRESHOLD = threshold  # 非最大值抑制阈值默认0.4

    # 加载图片、转为blob格式、送入网络输入层
    img = frame.copy()
    blobImg = cv2.dnn.blobFromImage(img, 1.0 / 255.0, (416, 416), None, True,
                                   False)  # # net需要的输入是blob格式的，用blobFromImage这个函数来转格式
    net.setInput(blobImg)  # # 调用setInput函数将图片送入输入层

    # 获取网络输出层信息（所有输出层的名字），设定并前向传播
    outInfo = net.getUnconnectedOutLayersNames()  # # 前面的yolov3架构也讲了，yolo在每个scale都有输出，outInfo是每个scale的名字信息，供net.forward使用
    start = time.time()
    layerOutputs = net.forward(outInfo)  # 得到各个输出层的、各个检测框等信息，是二维结构。
    end = time.time()
    print("[INFO] YOLO took {:.6f} seconds".format(end - start))  # # 可以打印下信息

    # 拿到图片尺寸
    (H, W) = img.shape[:2]
    # 过滤layerOutputs
    # layerOutputs的第1维的元素内容: [center_x, center_y, width, height, objectness, N-class score data]
    # 过滤后的结果放入：
    boxes = []  # 所有边界框（各层结果放一起）
    confidences = []  # 所有置信度，概率
    classIDs = []  # 所有分类ID
    rectsAndClasses = []  # 所有概率超过阈值的框位置列表

    # # 1）过滤掉置信度低的框框
    for out in layerOutputs:  # 各个输出层
        for detection in out:  # 各个框框
            # 拿到置信度
            scores = detection[5:]  # 各个类别的置信度
            classID = np.argmax(scores)  # 最高置信度的id即为分类id
            confidence = scores[classID]  # 拿到置信度

            # 根据置信度筛查
            if confidence > CONFIDENCE:
                box = detection[0:4] * np.array([W, H, W, H])  # 将边界框放会图片尺寸
                (centerX, centerY, width, height) = box.astype("int")
                x = int(centerX - (width / 2))
                y = int(centerY - (height / 2))
                boxes.append([x, y, int(width), int(height)])
                confidences.append(float(confidence))
                classIDs.append(classID)

    # # 2）应用非最大值抑制(non-maxima suppression，nms)进一步筛掉
    idxs = cv2.dnn.NMSBoxes(boxes, confidences, CONFIDENCE, THRESHOLD)  # boxes中，保留的box的索引index存入idxs
    # 得到labels列表
    with open(labelsPath, 'rt') as f:
        labels = f.read().rstrip('\n').split('\n')
    # 应用检测结果
    np.random.seed(42)
    COLORS = np.random.randint(0, 255, size=(len(labels), 3),
                               dtype="uint8")  # 框框显示颜色，每一类有不同的颜色，每种颜色都是由RGB三个值组成的，所以size为(len(labels), 3)
    if len(idxs) > 0:
        for i in idxs.flatten():  # indxs是二维的，第0维是输出层，所以这里把它展平成1维
            (x, y) = (boxes[i][0], boxes[i][1])
            (w, h) = (boxes[i][2], boxes[i][3])

            color = [int(c) for c in COLORS[classIDs[i]]]
            cv2.rectangle(img, (x, y), (x + w, y + h), color, 3)  # 线条粗细为2px
            text = "{}: {:.4f}".format(labels[int(classIDs[i])], confidences[i])
            cv2.putText(img, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color,
                       2)  # cv.FONT_HERSHEY_SIMPLEX字体风格、0.5字体大小、粗细2px
            print(labels[int(classIDs[i])], ":", confidences[i])
            rectsAndClasses.append([x, y, w, h, labels[int(classIDs[i])], confidences[i], color])
    return img, rectsAndClasses  # 返回画过框的图片和NMS后的框列表


def genXML(imgName, shape, labels, classes, xmlPath):  # 

    new_txtname = imgName.split('.')[0]  #

    # 创建空的Dom文档对象
    doc = xml.dom.minidom.Document()
    # 创建根结点，根节点名为 annotation
    annotation = doc.createElement('annotation')  # 根节点
    # 将根节点添加到Dom文档对象中
    doc.appendChild(annotation)

    # folder节点
    folder = doc.createElement('folder')  # 创建一个名叫folder的节点
    # 内容写入
    folder_text = doc.createTextNode('JPEGImages')  # folder节点里面要写的内容
    folder.appendChild(folder_text)  # 添加到folder节点下，如果是内容，节点内容createTextNode类型，就作为内容写入；如果是createElement类型，就作为子节点添加进去
    annotation.appendChild(folder)  # 之后将添加好内容的folder节点，作为子节点添加到annotation节点中

    # filename节点
    filename = doc.createElement('filename')
    filename_text = doc.createTextNode(str(new_txtname) + '.jpg')
    filename.appendChild(filename_text)
    #
    annotation.appendChild(filename)

    # path节点
    path = doc.createElement('path')
    path_text = doc.createTextNode('E:\\darknet-master\\build\\darknet\\myVLPCharData\\JPEGImages\\%s.jpg' % new_txtname)  # 框架路径，根据自己修改
    path.appendChild(path_text)
    #
    annotation.appendChild(path)

    # sourch节点
    source = doc.createElement('source')
    #
    database = doc.createElement('database')
    database_text = doc.createTextNode('Unknown')
    database.appendChild(database_text)
    #
    source.appendChild(database)
    #
    annotation.appendChild(source)

    # size节点
    size = doc.createElement('size')

    width = doc.createElement('width')
    width_text = doc.createTextNode(str(shape[1]))
    width.appendChild(width_text)
    size.appendChild(width)

    height = doc.createElement('height')
    height_text = doc.createTextNode(str(shape[0]))
    height.appendChild(height_text)
    size.appendChild(height)

    depth = doc.createElement('depth')
    depth_text = doc.createTextNode(str(shape[-1]))
    depth.appendChild(depth_text)
    size.appendChild(depth)
    #
    annotation.appendChild(size)

    # segmented节点
    segmented = doc.createElement('segmented')
    segmented_text = doc.createTextNode('0')
    segmented.appendChild(segmented_text)
    #
    annotation.appendChild(segmented)

    # object节点
    for [y1, y2, x1, x2], pChar in zip(labels, classes):  # 分类坐标和分类名称
        object = doc.createElement('object')

        name = doc.createElement('name')
        name_text = doc.createTextNode(pChar)  # 这个地方是标签的name，也就是分类名称
        name.appendChild(name_text)
        object.appendChild(name)

        pose = doc.createElement('pose')
        pose_text = doc.createTextNode("Unspecified")
        pose.appendChild(pose_text)
        object.appendChild(pose)

        truncated = doc.createElement('truncated')
        truncated_text = doc.createTextNode("0")
        truncated.appendChild(truncated_text)
        object.appendChild(truncated)

        difficult = doc.createElement('difficult')
        difficult_text = doc.createTextNode("0")
        difficult.appendChild(difficult_text)
        object.appendChild(difficult)

        bndbox = doc.createElement('bndbox')
        #
        xmin = doc.createElement('xmin')
        xmin_text = doc.createTextNode(str(x1))
        xmin.appendChild(xmin_text)
        bndbox.appendChild(xmin)
        #
        ymin = doc.createElement('ymin')
        ymin_text = doc.createTextNode(str(y1))
        ymin.appendChild(ymin_text)
        bndbox.appendChild(ymin)
        #
        xmax = doc.createElement('xmax')
        xmax_text = doc.createTextNode(str(x2))
        xmax.appendChild(xmax_text)
        bndbox.appendChild(xmax)
        #
        ymax = doc.createElement('ymax')
        ymax_text = doc.createTextNode(str(y2))
        ymax.appendChild(ymax_text)
        bndbox.appendChild(ymax)
        #
        object.appendChild(bndbox)
        #
        annotation.appendChild(object)

    # 写入xml文本文件中
    if not os.path.exists(xmlPath):
        os.mkdir(xmlPath)
    fp = open(xmlPath + '/%s.xml' % new_txtname, 'w+')
    doc.writexml(fp, indent='\n', addindent='\t', newl='', encoding='utf-8')
    fp.close()


def run(imgPath):
    filesName = os.listdir(imgPath)  # 读取所有图片名
    for f in filesName:
        p = os.path.join(imgPath, f)  # 单独图片的路径
        frame = cv2.imread(p)
        classesImg, rectsAndClasses = ModelYoloV3(frame)
        cv2.imshow("Labels", classesImg)

        labels = []
        classes = []
        for x, y, w, h, cn, confidence, color in rectsAndClasses:
            X1 = x
            Y1 = y
            X2 = x + w
            Y2 = y + h
            labels.append([Y1, Y2, X1, X2])  # 分类坐标
            classes.append(cn)  # 分类名称
            genXML(f, frame.shape, labels, classes, 'E:\\darknet-master\\build\\xml')  # args: 图片名，图片shape， 标签坐标列表， 类名列表， xml文件保存路径
        cv2.waitKey(10)

class AutoLabels:

    def __init__(self):
        pass


if __name__ == '__main__':

    path = 'E:\\darknet-master\\build\\image'
    run(path)

参考博客：立即前往 https://www.cnblogs.com/study-/p/13959391.html