這是關於從頭實現YOLO v3檢測器的教程的第5部分。在上一部分中，我們實現了一個將網絡輸出轉換爲檢測預測的函數。有了可用的檢測器，剩下的就是創建輸入和輸出管道。
本教程的代碼設計爲在Python 3.5和PyTorch 0.4上運行。完整代碼可以在這裏找到 Github repo.
本教程分爲5個部分：
- Part 1 : Understanding How YOLO works
- Part 2 : Creating the layers of the network architecture
- Part 3 : Implementing the the forward pass of the network
- Part 4 : Confidence Thresholding and Non-maximum Suppression
- Part 5 (This one): Designing the input and the output pipelines
預備知識
- Part 1-4 of the tutorial.
- Basic working knowledge of PyTorch, including how to create custom architectures with nn.Module, nn.Sequential and torch.nn.parameter classes.
- Basic knowledge of OpenCV
在這一部分中，我們將構建檢測器的輸入和輸出管道。這包括從磁盤上讀取圖像，進行預測，使用預測在圖像上繪製邊框，然後將它們保存到磁盤上。我們還將介紹如何讓探測器在攝像機輸入或視頻上實時工作。我們將介紹一些命令行標誌，以便對網絡中的各種超參數進行一些試驗。那麼讓我們開始吧。
Note: You will need to install OpenCV 3 for this part.
在tour detector文件中創建一個文件detector.py。在頂部添加必要的導入。
from __future__ import division
import time
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
import cv2
from util import *
import argparse
import os
import os.path as osp
from darknet import Darknet
import pickle as pkl
import pandas as pd
import random
創建命令行參數
因爲detector.py是我們將執行以運行檢測器的文件，所以最好有命令行參數可以傳遞給它。我已經使用python的ArgParse模塊來實現這一點。
def arg_parse():
"""
Parse arguements to the detect module
"""
parser = argparse.ArgumentParser(description='YOLO v3 Detection Module')
parser.add_argument("--images", dest = 'images', help =
"Image / Directory containing images to perform detection upon",
default = "imgs", type = str)
parser.add_argument("--det", dest = 'det', help =
"Image / Directory to store detections to",
default = "det", type = str)
parser.add_argument("--bs", dest = "bs", help = "Batch size", default = 1)
parser.add_argument("--confidence", dest = "confidence", help = "Object Confidence to filter predictions", default = 0.5)
parser.add_argument("--nms_thresh", dest = "nms_thresh", help = "NMS Threshhold", default = 0.4)
parser.add_argument("--cfg", dest = 'cfgfile', help =
"Config file",
default = "cfg/yolov3.cfg", type = str)
parser.add_argument("--weights", dest = 'weightsfile', help =
"weightsfile",
default = "yolov3.weights", type = str)
parser.add_argument("--reso", dest = 'reso', help =
"Input resolution of the network. Increase to increase accuracy. Decrease to increase speed",
default = "416", type = str)
return parser.parse_args()
args = arg_parse()
images = args.images
batch_size = int(args.bs)
confidence = float(args.confidence)
nms_thesh = float(args.nms_thresh)
start = 0
CUDA = torch.cuda.is_available()
其中，重要的標誌包括images(用於指定輸入圖像或圖像目錄)、det(保存檢測的目錄)、reso(輸入圖像的分辨率，可用於速度-精度權衡)、cfg(可選配置文件)和weightfile權重文件。
下載網絡
從這裏 here下載coco.names文件，一個包含COCO數據集中對象名稱的文件。 在檢測器目錄中創建文件夾數據。 同樣，如果你在Linux上，你可以操作如下。
mkdir data
cd data
wget https://raw.githubusercontent.com/ayooshkathuria/YOLO_v3_tutorial_from_scratch/master/data/coco.names
然後，我們在程序中下載類文件
num_classes = 80 #For COCO
classes = load_classes("data/coco.names")
load_classes是util.py中定義的函數，它返回一個字典，該字典將每個類的索引映射到它的名稱字符串。
def load_classes(namesfile):
fp = open(namesfile, "r")
names = fp.read().split("\n")[:-1]
return names
初始化網絡和下載權重。
#建立神經網絡
print("Loading network.....")
model = Darknet(args.cfgfile)
model.load_weights(args.weightsfile)
print("Network successfully loaded")
model.net_info["height"] = args.reso
inp_dim = int(model.net_info["height"])
assert inp_dim % 32 == 0
assert inp_dim > 32
#如果有一個GPU可用，把模型放在GPU上
if CUDA:
model.cuda()
#在評估模式下設置模型
model.eval()
讀取輸入圖片
從磁盤或目錄讀取圖片。圖像的路徑存儲在一個名爲imlist的列表中。
read_dir = time.time()
#Detection phase
try:
imlist = [osp.join(osp.realpath('.'), images, img) for img in os.listdir(images)]
except NotADirectoryError:
imlist = []
imlist.append(osp.join(osp.realpath('.'), images))
except FileNotFoundError:
print ("No file or directory with the name {}".format(images))
exit()
read_dir是用於測量時間的檢查點。(我們會遇到其中的幾個)。
如果保存檢測的目錄不存在(由det標誌定義)，那就創建它。
if not os.path.exists(args.det):
os.makedirs(args.det)
我們使用OpenCV來下載圖片。
load_batch = time.time()
loaded_ims = [cv2.imread(x) for x in imlist]
load_batch 又是一個檢查點。
OpenCV以numpy數組的形式加載圖像，BGR作爲顏色通道的順序。PyTorch的圖像輸入格式爲(batch x Channels x Height x Width)，通道順序爲RGB。因此，我們在util.py中編寫函數prep_image，將numpy數組轉換爲PyTorch的輸入格式。
在編寫這個函數之前，我們必須編寫一個函數letterbox_image，它可以調整圖像的大小，保持長寬比一致，並用顏色(128,128,128)填充剩下的區域。
def letterbox_image(img, inp_dim):
'''resize image with unchanged aspect ratio using padding'''
img_w, img_h = img.shape[1], img.shape[0]
w, h = inp_dim
new_w = int(img_w * min(w/img_w, h/img_h))
new_h = int(img_h * min(w/img_w, h/img_h))
resized_image = cv2.resize(img, (new_w,new_h), interpolation = cv2.INTER_CUBIC)
canvas = np.full((inp_dim[1], inp_dim[0], 3), 128)
canvas[(h-new_h)//2:(h-new_h)//2 + new_h,(w-new_w)//2:(w-new_w)//2 + new_w, :] = resized_image
return canvas
現在，我們編寫一個函數，該函數獲取OpenCV圖像並將其轉換爲網絡的輸入。
def prep_image(img, inp_dim):
"""
Prepare image for inputting to the neural network.
Returns a Variable
"""
img = cv2.resize(img, (inp_dim, inp_dim))
img = img[:,:,::-1].transpose((2,0,1)).copy()
img = torch.from_numpy(img).float().div(255.0).unsqueeze(0)
return img
除了轉換後的圖像之外，我們還維護原始圖像的列表，以及包含原始圖像維度的im_dim_list列表。
#圖像的PyTorch變量
im_batches = list(map(prep_image, loaded_ims, [inp_dim for x in range(len(imlist))]))
#包含原始圖像尺寸的列表
im_dim_list = [(x.shape[1], x.shape[0]) for x in loaded_ims]
im_dim_list = torch.FloatTensor(im_dim_list).repeat(1,2)
if CUDA:
im_dim_list = im_dim_list.cuda()
創建一個批次
leftover = 0
if (len(im_dim_list) % batch_size):
leftover = 1
if batch_size != 1:
num_batches = len(imlist) // batch_size + leftover
im_batches = [torch.cat((im_batches[i*batch_size : min((i + 1)*batch_size,
len(im_batches))])) for i in range(num_batches)]
檢測循環
我們對批次進行迭代，生成預測，並連接預測張量(形狀，Dx8, write_results函數的輸出)所有我們需要檢測的圖像。
對於每個批次，我們將度量檢測所花費的時間，即從獲取輸入到生成write_results函數輸出之間所花費的時間。在write_forecast返回的輸出中，其中一個屬性是批次圖像的索引。我們轉換該特定屬性的方式是，現在它表示imlist中圖像的索引，即列表包含所有圖像地址。
然後，我們打印每次檢測所花費的時間和每張圖像中檢測到的對象。
如果批處理的write_results函數的輸出是int(0)，這意味着沒有檢測，則使用continue跳過其餘的循環。
write = 0
start_det_loop = time.time()
for i, batch in enumerate(im_batches):
#load the image
start = time.time()
if CUDA:
batch = batch.cuda()
prediction = model(Variable(batch, volatile = True), CUDA)
prediction = write_results(prediction, confidence, num_classes, nms_conf = nms_thesh)
end = time.time()
if type(prediction) == int:
for im_num, image in enumerate(imlist[i*batch_size: min((i + 1)*batch_size, len(imlist))]):
im_id = i*batch_size + im_num
print("{0:20s} predicted in {1:6.3f} seconds".format(image.split("/")[-1], (end - start)/batch_size))
print("{0:20s} {1:s}".format("Objects Detected:", ""))
print("----------------------------------------------------------")
continue
prediction[:,0] += i*batch_size #將atribute從批量索引轉換爲imlist中的索引
if not write: #如果我們沒有初始化輸出
output = prediction
write = 1
else:
output = torch.cat((output,prediction))
for im_num, image in enumerate(imlist[i*batch_size: min((i + 1)*batch_size, len(imlist))]):
im_id = i*batch_size + im_num
objs = [classes[int(x[-1])] for x in output if int(x[0]) == im_id]
print("{0:20s} predicted in {1:6.3f} seconds".format(image.split("/")[-1], (end - start)/batch_size))
print("{0:20s} {1:s}".format("Objects Detected:", " ".join(objs)))
print("----------------------------------------------------------")
if CUDA:
torch.cuda.synchronize()
torch.cuda.synchronize 行確保CUDA內核與CPU同步。否則，一旦GPU作業排隊，CUDA內核就會在GPU作業完成之前將控制權返回給CPU（異步調用）。 如果在GPU作業實際結束之前打印end = time.time（），這可能會導致誤導時間。
現在，我們有了張量輸出中所有圖像的檢測。讓我們在圖像上繪製邊界框。
在圖像上繪製邊界框
我們使用try-catch塊來檢查是否已經進行過單次檢測。 如果不是這樣，請退出程序。
try:
output
except NameError:
print ("No detections were made")
exit()
在繪製邊界框之前，輸出張量中包含的預測符合網絡的輸入大小，而不是圖像的原始大小。因此，在繪製邊界框之前，讓我們將每個邊界框的角屬性轉換爲圖像的原始維度。
在繪製邊界框之前，輸出張量中包含的預測是對填充圖像的預測，而不是對原始圖像的預測。只是，將它們重新縮放到輸入圖像的尺寸在這裏不起作用。我們首先需要根據包含原始圖像的填充圖像上的區域邊界來變換要測量的框的座標。
im_dim_list = torch.index_select(im_dim_list, 0, output[:,0].long())
scaling_factor = torch.min(inp_dim/im_dim_list,1)[0].view(-1,1)
output[:,[1,3]] -= (inp_dim - scaling_factor*im_dim_list[:,0].view(-1,1))/2
output[:,[2,4]] -= (inp_dim - scaling_factor*im_dim_list[:,1].view(-1,1))/2
現在，我們的座標與填充區域上圖像的尺寸一致。然而，在letterbox_image函數中，我們通過縮放因子調整了圖像的兩個維度(記住，這兩個維度都是用一個公共因子來劃分的，以保持長寬比)。現在，我們撤銷這個縮放，以獲取原始圖像上的邊框座標。
output[:,1:5] /= scaling_factor
現在讓我們將可能在圖像外部有邊界的任何邊框剪輯到圖像的邊緣。
for i in range(output.shape[0]):
output[i, [1,3]] = torch.clamp(output[i, [1,3]], 0.0, im_dim_list[i,0])
output[i, [2,4]] = torch.clamp(output[i, [2,4]], 0.0, im_dim_list[i,1])
如果圖像中有太多的邊框，將它們全部用一種顏色繪製可能不是一個好主意。將此文件 file 下載到檢測器文件夾。這是一個pickle文件，其中包含許多顏色可供隨機選擇。
class_load = time.time()
colors = pkl.load(open("pallete", "rb"))
現在讓我們寫一個函數來畫這些邊框。
draw = time.time()
def write(x, results, color):
c1 = tuple(x[1:3].int())
c2 = tuple(x[3:5].int())
img = results[int(x[0])]
cls = int(x[-1])
label = "{0}".format(classes[cls])
cv2.rectangle(img, c1, c2,color, 1)
t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1 , 1)[0]
c2 = c1[0] + t_size[0] + 3, c1[1] + t_size[1] + 4
cv2.rectangle(img, c1, c2,color, -1)
cv2.putText(img, label, (c1[0], c1[1] + t_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 1, [225,255,255], 1);
return img
上面的函數從colors中隨機選擇顏色繪製一個矩形。它還在邊界框的左上角創建一個填充矩形，並在填充矩形中寫入檢測到的對象類。cv2.rectangle函數的-1參數用於創建填充矩形。
我們在本地定義write函數，以便它能夠訪問colors列表。我們也可以把colors作爲參數，但是那樣的話，每個圖像只能使用一種顏色，這違背了我們到目標。
一旦我們定義了這個函數，現在讓我們在圖像上繪製邊界框。
list(map(lambda x: write(x, loaded_ims), output))
上面的代碼片段修改了loaded_ims中的圖像。
每個圖像都是通過在圖像名稱前面加上前綴“det_”保存的。我們創建一個地址列表，將檢測圖像保存到該地址。
det_names = pd.Series(imlist).apply(lambda x: "{}/det_{}".format(args.det,x.split("/")[-1]))
最後，將檢測到的圖片寫入det_names地址中。
list(map(cv2.imwrite, det_names, loaded_ims))
end = time.time()
打印總計時間
在檢測器的末尾，我們將打印一個摘要，其中包含代碼的哪一部分需要執行多長時間。當我們需要比較不同的超參數如何影響探測器的速度時，這是有用的。在命令行上執行腳本 detection.py 時，可以設置批處理大小、對象置信度和NMS閾值等超參數(分別使用bs、confidence, nms_thresh標誌傳遞)。
print("SUMMARY")
print("----------------------------------------------------------")
print("{:25s}: {}".format("Task", "Time Taken (in seconds)"))
print()
print("{:25s}: {:2.3f}".format("Reading addresses", load_batch - read_dir))
print("{:25s}: {:2.3f}".format("Loading batch", start_det_loop - load_batch))
print("{:25s}: {:2.3f}".format("Detection (" + str(len(imlist)) + " images)", output_recast - start_det_loop))
print("{:25s}: {:2.3f}".format("Output Processing", class_load - output_recast))
print("{:25s}: {:2.3f}".format("Drawing Boxes", end - draw))
print("{:25s}: {:2.3f}".format("Average time_per_img", (end - load_batch)/len(imlist)))
print("----------------------------------------------------------")
torch.cuda.empty_cache()
Testing The Object Detector
比如，在終端運行
python detect.py --images dog-cycle-car.png --det det
產生如下輸出
以下代碼在CPU上運行。 預計GPU上檢測時間要快得多。 在特斯拉K80上大約0.1秒/幅。
Loading network.....
Network successfully loaded
dog-cycle-car.png predicted in 2.456 seconds
Objects Detected: bicycle truck dog
----------------------------------------------------------
SUMMARY
----------------------------------------------------------
Task : Time Taken (in seconds)
Reading addresses : 0.002
Loading batch : 0.120
Detection (1 images) : 2.457
Output Processing : 0.002
Drawing Boxes : 0.076
Average time_per_img : 2.657
----------------------------------------------------------
名爲det_dog-cycle-car.png的圖像保存在det目錄中。
在視頻和網頁上運行檢測器
爲了在視頻或網絡攝像頭上運行檢測器，代碼幾乎保持不變，除了我們不必迭代批次，而是迭代視頻幀。
在視頻上運行檢測器的代碼可以在github存儲庫中的video.py文件中找到。 除了一些更改之外，代碼與detect.py的代碼非常相似。
首先，我們在OpenCV中打開視頻/攝像頭。
videofile = "video.avi" #or path to the video file.
cap = cv2.VideoCapture(videofile)
#cap = cv2.VideoCapture(0) for webcam
assert cap.isOpened(), 'Cannot capture source'
frames = 0
我們迭代幀的方式類似於我們迭代圖片的方式。
許多地方簡化了許多代碼，因爲我們不再需要處理批處理，而是一次只處理一張圖片。這是因爲一次只能有一個幀。這包括使用tuple代替im_dim_list中的張量，以及在write函數中進行細微的更改。
在每次迭代中，我們都會跟蹤變量frames中捕獲的幀的數量。然後，我們將該數字除以自第一幀以來經過的時間以打印視頻的FPS。
打算使用cv2.imwrite將檢測圖片寫入磁盤，我們使用cv2.imshow來顯示帶有邊框的幀。如果用戶按下Q按鈕，就會導致代碼中斷循環，視頻結束。
frames = 0
start = time.time()
while cap.isOpened():
ret, frame = cap.read()
if ret:
img = prep_image(frame, inp_dim)
# cv2.imshow("a", frame)
im_dim = frame.shape[1], frame.shape[0]
im_dim = torch.FloatTensor(im_dim).repeat(1,2)
if CUDA:
im_dim = im_dim.cuda()
img = img.cuda()
output = model(Variable(img, volatile = True), CUDA)
output = write_results(output, confidence, num_classes, nms_conf = nms_thesh)
if type(output) == int:
frames += 1
print("FPS of the video is {:5.4f}".format( frames / (time.time() - start)))
cv2.imshow("frame", frame)
key = cv2.waitKey(1)
if key & 0xFF == ord('q'):
break
continue
output[:,1:5] = torch.clamp(output[:,1:5], 0.0, float(inp_dim))
im_dim = im_dim.repeat(output.size(0), 1)/inp_dim
output[:,1:5] *= im_dim
classes = load_classes('data/coco.names')
colors = pkl.load(open("pallete", "rb"))
list(map(lambda x: write(x, frame), output))
cv2.imshow("frame", frame)
key = cv2.waitKey(1)
if key & 0xFF == ord('q'):
break
frames += 1
print(time.time() - start)
print("FPS of the video is {:5.2f}".format( frames / (time.time() - start)))
else:
break
結論
在本系列教程中，我們從頭開始實現了一個對象檢測器，併爲此歡呼。我仍然認爲，能夠寫出高效的代碼是深度學習實踐者可能擁有的最被低估的技能之一。不管你的想法多麼具有革命性，除非你能測試它，否則它是沒有用的。爲此，您需要具有強大的編碼技能。
我還瞭解到深入學習任何主題的最佳方法是實現代碼。 它會迫使您瀏覽一下您在閱讀論文時可能錯過的主題的微妙之處。我希望這個教程系列可以作爲一種練習，可以培養你作爲深度學習練習者的技能。
擴展閱讀
- PyTorch tutorial
- OpenCV Basics
- Python ArgParse