python使用yolov3/yolov3-tiny訓練好的權重文件.weights進行行人檢測,批量測試自定義文件夾下的圖片並輸出至指定文件夾
目錄
python使用yolov3/yolov3-tiny訓練好的權重文件.weights進行行人檢測,批量測試自定義文件夾下的圖片並輸出至指定文件夾
一、寫在開頭
最近在做畢業設計,研究的是目標識別與追蹤,前段時間打算只用opencv識別個簡單的目標就算了,但參考着論文硬着頭皮擼了一星期的圖片預處理,到了後面的識別部分,實在做不下去了,太南了!也不知道怎麼具體去改其中的一些參數(數學不太好,害...),所以經過多方的baidu,用了yolov3跑訓練模型,昨天剛跑完個只檢測person的demo,買的東西還沒有來到,也不知道實際效果怎麼樣,只是網上找的一些照片隨便測試了一下,但效果其實一般般,看過幾天實際效果怎麼樣吧,不太好的話就自己再做一部分數據加入到訓練集中去,效果應該會好很多。所以,在此先記錄下,如何使用已經訓練好的權重文件來批量測試自定義文件夾下的照片並保存至指定文件夾。之前也找了這方面的教程,發現多數是去修改的.c的配置文件,我又懶,記性又不好,修改了以後一段時間我鐵定忘記再改回來。所以,也就有了這篇blog。
二、已有的環境條件
1. pycharm--python
我在win10的pycharm中進行的測試,Ubuntu或其他環境也一樣的。
2. opencv3.4
我win10和Ubuntu上都是用的opencv3.4,保持版本一致,可能會避免一些沒必要的麻煩。
3. 用yolov3訓練好了自己的權重文件.weights
由於我是深度學習方面的新手,所以暫時沒有製作自己的數據集,直接用的VOC的數據集。在此程序中,需要用到訓練好的三個文件,分別是voc.names,yolo3-tiny_person.cfg,yolov3-tiny_person800.weights。如下圖所示,三個文件在yolo_voc文件夾下,
三、文件目錄結構
如下所示:
py_person_yolov3
|
│
├─yolo_voc
│ voc.names
│ yolov3-tiny_person.cfg
│ yolov3-tiny_person800.weights
│
├─test-img1 # 要測試的圖片文件夾,測試的圖片都放在該文件夾下
│ 1.jpg
│ 2.jpg
│ 3.jpg
│ 4.jpg
│ 5.jpg
│ 6.jpg
│ 7.jpg
│ 8.jpg
│ 9.jpg
├─detect_person_fromDir.py
四、批量測試圖片測試程序
detect_person_fromDir.py程序如下:
# -*- coding: utf-8 -*-
# 載入所需庫
# author:cpz
# date:2020-04-10
"""
該程序在程序detect_person.py的基礎上,改進了只能對一張照片進行目標檢測,
可以實現對指定文件夾下的所有圖片進行目標檢測。
加入命令行參數--inDir——即待檢測的所有圖片的文件夾;
--outDir——即檢測完之後的圖片的保存路徑。
"""
import cv2
import numpy as np
import os
import time
import argparse
def yolo_detect(inDir,
outDir,
label_path='./yolo_voc/voc.names',
config_path='./yolo_voc/yolov3-tiny_person.cfg',
weights_path='./yolo_voc/yolov3-tiny_person800.weights',
confidence_thre=0.3,
nms_thre=0.3,
jpg_quality=80):
'''
pathIn:原始圖片的路徑
pathOut:結果圖片的路徑
label_path:類別標籤文件的路徑
config_path:模型配置文件的路徑
weights_path:模型權重文件的路徑
confidence_thre:0-1,置信度(概率/打分)閾值,即保留概率大於這個值的邊界框,默認爲0.5
nms_thre:非極大值抑制的閾值,默認爲0.3
jpg_quality:設定輸出圖片的質量,範圍爲0到100,默認爲80,越大質量越好
'''
# 輸出結果圖片--->用絕對路徑,用了一次相對路徑,發現並沒有結果
# os.getcwd()——用於返回當前工作目錄
currentDir = os.getcwd()
print("當前工作路徑爲:%s" % (currentDir))
# 如果輸出文件夾在命令行中未輸入
# 則先創建一個文件夾:"outDir-from+輸入照片的文件夾名字"
if outDir is None:
cDir = "outDir-from_" + inDir
# 判斷是否已經存在了文件夾----> 存在的話就跳過,什麼也不做
if os.path.exists(cDir):
print("由於未輸入輸出圖片保存的文件夾,且已經存在了這個文件夾 %s" %(cDir))
pass
else:
os.mkdir(cDir)
print("由於未輸入輸出圖片保存的文件夾,則創建了文件夾 %s" %(cDir))
# 加載類別標籤文件
LABELS = open(label_path).read().strip().split("\n")
nclass = len(LABELS)
# 爲每個類別的邊界框隨機匹配相應顏色
np.random.seed(42)
COLORS = np.random.randint(0, 255, size=(nclass, 3), dtype='uint8')
"cpz在2020-4-10新改的,爲了對一個目錄中所有的圖片進行檢測"
# 試試只加載一次模型行不行
# 加載模型配置和權重文件
print('從硬盤加載YOLO......')
net = cv2.dnn.readNetFromDarknet(config_path, weights_path)
# 獲取YOLO輸出層的名字
ln = net.getLayerNames()
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
# os.listdir()——用於返回指定的文件夾包含的文件或文件夾的名字的列表
img_list = os.listdir(inDir)
for i in range(0, len(img_list)):
# os.path.join()——連接兩個或更多的路徑名組件
img_path = os.path.join(inDir, img_list[i])
# 載入圖片並獲取其維度
base_path = os.path.basename(img_path) # base_path爲照片的名字(有擴展名)
img = cv2.imread(img_path)
(H, W) = img.shape[:2]
# 將圖片構建成一個blob,設置圖片尺寸,然後執行一次
# YOLO前饋網絡計算,最終獲取邊界框和相應概率
blob = cv2.dnn.blobFromImage(img, 1 / 255.0, (416, 416), swapRB=True, crop=False)
net.setInput(blob)
start = time.time()
layerOutputs = net.forward(ln)
end = time.time()
# 顯示預測所花費時間
print('YOLO模型花費 {:.2f} 秒來預測一張圖片'.format(end - start))
# 初始化邊界框,置信度(概率)以及類別
boxes = []
confidences = []
classIDs = []
# 迭代每個輸出層,總共三個
for output in layerOutputs:
# 迭代每個檢測
for detection in output:
# 提取類別ID和置信度
scores = detection[5:]
classID = np.argmax(scores)
confidence = scores[classID]
# 只保留置信度大於某值的邊界框
if confidence > confidence_thre:
# 將邊界框的座標還原至與原圖片相匹配,記住YOLO返回的是
# 邊界框的中心座標以及邊界框的寬度和高度
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, confidence_thre, nms_thre)
# 確保至少一個邊界框
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])
# 繪製邊界框以及在左上角添加類別標籤和置信度
color = [int(c) for c in COLORS[classIDs[i]]]
cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
text = '{}: {:.3f}'.format(LABELS[classIDs[i]], confidences[i])
(text_w, text_h), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
cv2.rectangle(img, (x, y-text_h-baseline), (x + text_w, y), color, -1)
cv2.putText(img, text, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)
# 要輸出的圖片的名稱(含有擴展名)
outImg_name = 'with_box_' + base_path
outImg_path = os.path.join(currentDir, cDir, outImg_name)
cv2.imwrite(outImg_path, img, [int(cv2.IMWRITE_JPEG_QUALITY), jpg_quality])
# 輸入了要輸出的照片的文件夾
if __name__ == '__main__':
ap = argparse.ArgumentParser(description='Input the image to detect person, and the path to save the detected image.')
ap.add_argument("-i", "--inDir", required=True, help="path dir of input image")
ap.add_argument("-o", "--outDir", required=False, help="path to save detected image, default: ./outDir-from+inDir/")
args = vars(ap.parse_args())
# 從命令行輸入中獲得輸入圖片的路徑
pathIn = args["inDir"]
# 從命令行輸入中獲得檢測完之後的結果圖片的保存路徑
pathOut = args["outDir"]
# 進行檢測咯!
yolo_detect(inDir=pathIn, outDir=pathOut)
五、進行測試
1. 打開pycharm的Terminal,並打開當前程序所在的文件夾,如下所示:
2. 運行程序,並輸入所必需的參數--inDir:
python detect_person_fromDir.py --inDir test-img1
3. 運行結果:
4. 打開輸出文件夾 outDir-from_test-img1,下面這兩張是我挑的效果比較好的兩張(先欺騙下自己吧):
六、寫在最後
經過一段時間的學習,對yolov3的原理還每細看,這是我下面準備要做的。還有,感謝寫blog分享經驗的各位,用自己的一段時間寫的blog真的可以幫到很多人,這也是我blog名字的由來。最後的最後,該程序是我在Ubuntu上搜到的一個程序,並做了進一步的改進(yolo_detect()函數是主要出自該博主),等我下次打開Ubuntu系統的時候我記得把原鏈接貼上,非常感謝這位博主。有找到的同學可以評論下,謝謝啦,就先寫到這裏了。