link

前言

在訓練目標檢測時，標注資料是一項簡單而又浪費時間的事情，如果能夠自動標注資料將可以高效的擴充資料集，從而提高訓練模型的效果，
目前能想到的一種自動標注方法是先訓練一個檢測效果較好的模型，然后使用這個模型對需要標注的圖片進行檢測，將檢測到的BBox寫入到標注檔案中，這里來簡單說明一下相關代碼，
python腳本檔案下載地址：https://download.csdn.net/download/qq_43019451/12836771

一、使用opencv呼叫模型

此處以Darknet模型為例，使用OpenCV呼叫模型來檢測圖片中的目標，

#參考opencv-python庫
import cv2

#此處設定相關的檔案路徑，我使用的時人臉檢測的模型，所示是face.weights
weightsPath = "./face/face.weights"
configPath = "./face/face.cfg"
labelsPath = "./face/face.names"

#讀取names檔案中的類別名
LABELS = open(labelsPath).read().strip().split("\n")

#使用opencv加載Darknet模型
net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

#下面是通過檢測獲取坐標的函式
def coordinate_get(img):
    coordinates_list=[] # 創建坐標串列
    boxes = []
    confidences = []
    classIDs = []
    (H, W) = img.shape[:2]
    # 得到 YOLO需要的輸出層
    ln = net.getLayerNames()
    ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
    # 從輸入影像構造一個blob，然后通過加載的模型，給我們提供邊界框和相關概率
    blob = cv2.dnn.blobFromImage(img, 1 / 255.0, (416, 416), swapRB=True, crop=False)
    net.setInput(blob)
    layerOutputs = net.forward(ln)

    # 在每層輸出上回圈
    for output in layerOutputs:
        # 對每個檢測進行回圈
        for detection in output:
            scores = detection[5:]
            classID = np.argmax(scores)
            confidence = scores[classID]
            # 過濾掉那些置信度較小的檢測結果
            if confidence > 0.01:
                # 框后接框的寬度和高度
                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) 

    idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.2, 0.3)
    if len(idxs) > 0:
        for i in idxs.flatten():
            (x, y) = (boxes[i][0], boxes[i][1])
            (w, h) = (boxes[i][2], boxes[i][3])

            xmin = int(x)
            ymin = int(y)
            xmax = int(x + w)
            ymax = int(y + h)
            coordinates_list.append([xmin,ymin,xmax,ymax,classIDs[i]])

    return coordinates_list

二、在xml檔案中創建樹并寫入坐標

需要參考ElementTree來創建xml檔案，并在xml檔案中加入tree結構，相關函式如下：

import os
from os import getcwd
from xml.etree import ElementTree as ET

# 定義一個創建一級分支object的函式
def create_object(root,xi,yi,xa,ya,obj_name):   # 引數依次，樹根，xmin，ymin，xmax，ymax
    #創建一級分支object
    _object=ET.SubElement(root,'object')
    #創建二級分支
    name=ET.SubElement(_object,'name')
    print(obj_name)
    name.text= str(obj_name)
    pose=ET.SubElement(_object,'pose')
    pose.text='Unspecified'
    truncated=ET.SubElement(_object,'truncated')
    truncated.text='0'
    difficult=ET.SubElement(_object,'difficult')
    difficult.text='0'
    #創建bndbox
    bndbox=ET.SubElement(_object,'bndbox')
    xmin=ET.SubElement(bndbox,'xmin')
    xmin.text='%s'%xi
    ymin = ET.SubElement(bndbox, 'ymin')
    ymin.text = '%s'%yi
    xmax = ET.SubElement(bndbox, 'xmax')
    xmax.text = '%s'%xa
    ymax = ET.SubElement(bndbox, 'ymax')
    ymax.text = '%s'%ya

# 創建xml檔案的函式
def create_tree(image_name, h, w):
    global annotation
    # 創建樹根annotation
    annotation = ET.Element('annotation')
    #創建一級分支folder
    folder = ET.SubElement(annotation,'folder')
    #添加folder標簽內容
    folder.text=(imgdir)

    #創建一級分支filename
    filename=ET.SubElement(annotation,'filename')
    filename.text=image_name

    #創建一級分支path
    path=ET.SubElement(annotation,'path')

    path.text= getcwd() + '\{}\{}'.format(imgdir,image_name)  # 用于回傳當前作業目錄
    
    #創建一級分支source
    source=ET.SubElement(annotation,'source')
    #創建source下的二級分支database
    database=ET.SubElement(source,'database')
    database.text='Unknown'

    #創建一級分支size
    size=ET.SubElement(annotation,'size')
    #創建size下的二級分支影像的寬、高及depth
    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 = '3'

    #創建一級分支segmented
    segmented = ET.SubElement(annotation,'segmented')
    segmented.text = '0'

定義一個main函式

最后定義一個main函式：

    for image_name in IMAGES_LIST:
        #判斷后綴只處理jpg檔案
        if image_name.endswith('jpg'):
            image = cv2.imread(os.path.join(imgdir, image_name))
            coordinates_list = coordinate_get(image)
            (h, w) = image.shape[:2]
            create_tree(image_name, h, w)

            for coordinate in coordinates_list:
                label_id = coordinate[4]
                create_object(annotation, coordinate[0], coordinate[1], coordinate[2], coordinate[3], LABELS[label_id])
                # if coordinates_list==[]:
                #     break

            # 將樹模型寫入xml檔案
            tree = ET.ElementTree(annotation)       
            tree.write('.\{}\{}.xml'.format(imgdir, image_name.strip('.jpg')))

總結

提示：在腳本運行前需要在python環境中安裝opencv等依賴庫，使用時修改模型組態檔的路徑，注意main函式中有判斷后綴的代碼，如果圖片不是jpg不會進行處理，