從零開始實現YOLO v3（part5）

9 人贊了文章

僅供學術交流，未經同意，請勿轉載）

（本文翻譯自：Tutorial on implementing YOLO v3 from scratch in PyTorch）

（這篇文章的原作者，原作者，原作者（重要的話說3遍）真的寫得很好很用心，去github上給他打個星星?吧）

這是從零開始實現YOLO v3檢測器的教程的第5部分。在上一部分中，我們實現了將網路輸出轉換為檢測的預測結果的函數。我們現在已經有了檢測器，剩下的就是創建輸入和輸出流程。

本教程的代碼旨在運行在Python 3.5和PyTorch 0.4上。它可以在這個Github中找到。

本教程分為5個部分：

第1部分：了解YOLO如何工作

第2部分：創建網路結構的層

第3部分：實現網路的前向傳播

第4部分：目標分數閾值和非最大值抑制

第5部分（本文）：設計輸入和輸出流程

預備知識

• 本教程的第1-4部分。
• PyTorch的基本知識，包括如何使用nn.Module，nn.Sequential和torch.nn.parameter類創建自定義體系結構。
• OpenCV的基本知識

編輯：在30/03/2018前，我們將任意大小的圖像調整為Darknet的輸入尺寸的方法是簡單地重新縮放尺寸。但是，在官方的實現中，圖像以保持寬高比不變的方式縮放，並填充空白的部分。例如，如果我們要將1900 x 1280圖像的調整為416 x 416，則調整大小後的圖像將如下所示。

輸入圖像的預處理的不同導致較早的實現的性能略低於原論文的性能。但是，現在已經更新了，按照原論文的實現中採用的調整大小的方法。

在這部分我們將建立探測器的輸入和輸出流程。這涉及從磁碟讀取圖像，進行預測，根據預測在圖像上繪製邊界框，然後將其保存到磁碟。我們還將介紹如何使檢測器在攝像頭或視頻上實時工作。我們將引入一些命令行標誌來允許對網路的各種超參數進行一些實驗。現在我們開始吧。

注意：您需要安裝OpenCV 3。

創建detector.py，在它的頂部導入一些必要的庫。

from __future__ import divisionimport timeimport torch import torch.nn as nnfrom torch.autograd import Variableimport numpy as npimport cv2 from util import *import argparseimport os import os.path as ospfrom darknet import Darknetimport pickle as pklimport pandas as pdimport 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.imagesbatch_size = int(args.bs)confidence = float(args.confidence)nms_thesh = float(args.nms_thresh)start = 0CUDA = torch.cuda.is_available()

其中重要的標誌是images（用於指定輸入圖像或圖像目錄），det（保存檢測的目標），reso（輸入圖像的解析度，調整這個值可以調節速度與精度之間的折衷），cfg（備用配置文件）和weightfile。

載入網路

從這裡下載文件coco.names，該文件包含COCO數據集中目標的名稱。在檢測器目錄中創建一個文件夾data。如果你使用的是Linux，你可以輸入。

mkdir datacd datawget https://raw.githubusercontent.com/ayooshkathuria/YOLO_v3_tutorial_from_scratch/master/data/coco.names

然後，我們在程序中載入類別文件。

num_classes = 80 #For COCOclasses = load_classes("data/coco.names")

load_classes是util.py中定義的函數，它返回一個將每個類的索引映射到名稱字元串的字典。

def load_classes(namesfile): fp = open(namesfile, "r") names = fp.read().split("
")[:-1] return names

初始化網路並載入權重。

#Set up the neural networkprint("Loading network.....")model = Darknet(args.cfgfile)model.load_weights(args.weightsfile)print("Network successfully loaded")model.net_info["height"] = args.resoinp_dim = int(model.net_info["height"])assert inp_dim % 32 == 0 assert inp_dim > 32#If theres a GPU availible, put the model on GPUif CUDA: model.cuda()#Set the model in evaluation modemodel.eval()

讀取輸入圖片

從磁碟讀取一張圖片或從目錄讀取多張圖片。一張圖片／多張圖片的路徑存儲在名為imlist的列表中。

read_dir = time.time()#Detection phasetry: 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標誌指定的檢測目錄不存在，請創建它。

f 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的圖像輸入格式是（批x通道x高x寬），通道順序為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 Variables for imagesim_batches = list(map(prep_image, loaded_ims, [inp_dim for x in range(len(imlist))]))#List containing dimensions of original imagesim_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()

創建批（batch）

leftover = 0if (len(im_dim_list) % batch_size): leftover = 1if 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)]

檢測循環

我們按批迭代，生成預測結果，並把執行檢測的所有圖像的預測結果的張量（它的形狀是D x 8，，來自write_results函數的輸出）連接起來。

對於每個批，我們將測量檢測所花費的時間，即獲取輸入和生成write_results函數輸出之間的時間。在由write_prediction返回的輸出中，其中一個屬性是批中圖像的索引。我們對該特定屬性（索引）進行轉換，使其成為imlist（該列表包含所有圖像的地址）中圖像的索引。

之後，我們會列印每個檢測的時間以及每個圖像中檢測到的目標。

如果批的write_results函數的輸出是int（0），意味著沒有檢測，我們使用continue繼續跳過剩下的循環。

write = 0start_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 #transform the atribute from index in batch to index in imlist if not write: #If we havet initialised output 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同步。否則，CUDA內核會在GPU作業排隊後立即將控制返回給CPU，這時GPU作業尚未完成（非同步調用）。如果在GPU作業實際結束之前end = time.time（）被列印出來，這可能會導致錯誤的時間。

現在，我們的Output張量擁有了所有圖像的輸出。讓我們在圖像上繪製邊界框。

在圖像上繪製邊界框

我們使用try-catch塊來檢查是否已經有檢測結果。如果沒有，則退出程序。

try: outputexcept 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))/2output[:,[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])

如果圖像中的邊界框太多，將它們全部繪製成同一種顏色可能不大好。將此文件下載到您的檢測器文件夾。這是一個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標誌傳遞）。

測試目標檢測器

例如，在終端上運行下面的命令，

python detect.py --images dog-cycle-car.png --det det

產生的輸出如下：

以下代碼在CPU上運行。預計GPU上的檢測時間要快得多。它在Tesla K80上約0.1秒/圖像。

Loading network.....Network successfully loadeddog-cycle-car.png predicted in 2.456 secondsObjects Detected: bicycle truck dog----------------------------------------------------------SUMMARY----------------------------------------------------------Task : Time Taken (in seconds)Reading addresses : 0.002Loading batch : 0.120Detection (1 images) : 2.457Output Processing : 0.002Drawing Boxes : 0.076Average 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 webcamassert cap.isOpened(), Cannot capture sourceframes = 0

用對圖像進行迭代的類似方式迭代幀。

很多地方的許多代碼都被簡化了，因為我們不再需要處理批，而是一次只需處理一個圖像，這是因為一次只能有一幀。這包括使用元組來代替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

結論

在這一系列教程中，我們從零開始實現了一個目標檢測器，並為達到這個目標而歡呼雀躍。我仍然認為能夠生成高效的代碼是深度學習實踐者該擁有的最被低估的技能之一。無論你的想法是多麼的新穎，除非你能測試該想法，否則它也是無用的。因此，你需要有強大的編碼技能。

我還明白到學習任何關於深度學習的主題的最佳途徑就是實現代碼。它強制你瀏覽一個主題的基本細微之處，而這些在閱讀論文時可能會被你忽略。我希望這個系列教程能夠鍛煉你作為一名深度學習實踐者的技能。

擴展閱讀

1. PyTorch tutorial
2. OpenCV Basics
3. Python ArgParse

推薦閱讀：