yolo預訓練模型自動標註
阿新 • • 發佈：2020-11-12
序：想要真正準確的的自動標註，的確不太現實，都能準確的自動標註了，還訓練模型幹嘛！
所以本文所寫方法是小量資料集預訓練模型後，自動打標最後微調。
（上圖是我的資料夾格式，將自己預訓練後的模型放到指定位置）
程式碼包含呼叫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)