zoukankan      html  css  js  c++  java
  • 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('
    ').split('
    ')
        # 应用检测结果
        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='
    ', addindent='	', 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

  • 相关阅读:
    学习笔记——Maven实战(九)打包的技巧
    学习笔记——Maven实战(八)常用Maven插件介绍(下)
    学习笔记——Maven实战(七)常用Maven插件介绍(上)
    学习笔记——Maven实战(六)Gradle,构建工具的未来?
    学习笔记——Maven实战(五)自动化Web应用集成测试
    在Google的GKE上创建支持Internal Load Balancer的Service
    Bash命令查找本机公网IP
    Google Cloud IAM中添加自定义域名
    Debian上启用Apache2服务
    Google Cloud VM上在线扩硬盘
  • 原文地址:https://www.cnblogs.com/buxian/p/13964606.html
Copyright © 2011-2022 走看看