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

 


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM