Pytorch從0開始實現YOLO V3指南 part5——設計輸入和輸出的流程

本節翻譯自：https://blog.paperspace.com/how-to-implement-a-yolo-v3-object-detector-from-scratch-in-pytorch-part-5/

在前一節最後，咱們實現了一個將網絡輸出轉換爲檢測預測的函數。如今咱們已經有了一個檢測器了，剩下的就是建立輸入和輸出的流程。

 

必要條件：

1.此係列教程的Part1到Part4。

2.Pytorch的基本知識，包括如何使用nn.Module，nn.Sequential，torch.nn.parameter類構建常規的結構

3.OpenCV的基礎知識

 

EDIT: 若是你在2018年3月30日以前訪問過這篇文章，咱們將任意大小的圖片調整爲Darknet的輸入大小的方法就是resize。然而在原始的實現中，調整圖像的大小時，須要保持長寬比不變，並填充遺漏的部分。例如，若是咱們將1900 x 1280的圖像調整爲416 x 415，那麼調整後的圖像應該是這樣的。

對於輸入處理的差別致使早期實現的性能略低於原始實現。如今這篇文章已經進行了更新，遵循了原始實現中調整大小的方法。

在這一部分中，咱們將構建檢測器的輸入和輸出管道。這包括從磁盤讀取圖像，進行預測，使用預測結果在圖像上繪製邊界框，而後將它們保存到磁盤。咱們還將介紹如何讓檢測器實時工做在一個攝像機或視頻中。咱們將介紹一些命令行標誌，以容許對網絡的各類超參數進行一些實驗。那麼讓咱們開始吧！

 注意:這部分須要安裝opencv3。

 

 建立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。

 

加載網絡：

從這裏下載coco.names文件，該文件包含COCO數據集中對象的名稱。在檢測器目錄中建立文件夾數據。一樣若是你在linux上工做，能夠輸入。

mkdir data
cd data
wget https://raw.githubusercontent.com/ayooshkathuria/YOLO_v3_tutorial_from_scratch/master/data/coco.name

而後，咱們在程序中加載該文件。

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

初始化網絡並加載權重。

#Set up the neural network
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

#If there's a GPU availible, put the model on GPU
if CUDA:
    model.cuda()

#Set the model in evaluation mode
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的圖像輸入格式爲(批量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 images
im_batches = list(map(prep_image, loaded_ims, [inp_dim for x in range(len(imlist))]))

#List containing dimensions of original images
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)] 

 

 

檢測循環體：

咱們對批次進行迭代，生成預測並鏈接全部圖像的預測張量(write_results函數的輸出，維度爲D*8)。

對於每一個批，咱們將檢測所花費的時間定義爲從接收輸入到write_results函數產生輸出之間所花費的時間。在write_prediction返回的輸出中，其中一個屬性是批中圖像的索引。咱們將其轉換成在imlist中圖像的索引，imlist列表包含全部圖像的地址。

以後，咱們打印每次檢測所花費的時間以及在每張圖像中檢測到的對象。若是批的write_results函數的輸出是int(0)就意味着沒有檢測，那麼咱們使用continue跳過rest循環。

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    #transform the atribute from index in batch to index in imlist 

    if not write:                      #If we have't 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做業排隊以後，甚至在GPU做業完成以前(異步調用)就將控制權返回給CPU。若是end = time() 在GPU做業實際結束以前打印出來，可能會致使錯誤的時間。

如今，咱們已經在tensor輸出中檢測到了全部的圖像。讓咱們在圖像上繪製邊界框吧！

 

在圖像上繪製邊界框：

咱們使用try-catch塊來檢查是否進行了一次檢測。若是沒有則退出程序。

try:
    output
except NameError:
    print ("No detections were made")
    exit()

 

在繪製邊界框以前，輸出tensor中包含的預測符合網絡的輸入大小而不是圖像的原始大小。所以，在繪製邊界框以前，讓咱們將每一個邊界框的角屬性轉換爲圖像的原始維度。

在繪製邊界框以前，輸出tensor中包含的預測是對填充圖像的預測，而不是對原始圖像的預測。僅僅將它們從新縮放到輸入圖像的維數在這裏是行不通的。首先，咱們須要將邊界框的座標轉換到相對於包含原始圖像的填充圖像上的邊界。

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])

 

 若是圖像中有不少邊界框，用一種顏色繪製它們可能不太好。將此文件下載到檢測器文件夾，這是一個pickle文件，其中包含許多顏色可供隨機選擇。

class_load = time.time()
colors = pkl.load(open("pallete", "rb"))

 

 如今咱們來寫繪製邊界框的函數。（x中的信息是圖像索引、4個角座標、目標置信度得分、最大置信類得分、該類的索引）

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中隨機選擇一個顏色來繪製矩形。它還在包圍框的左上角建立一個填充矩形，並在填充矩形中寫入檢測到的對象的類。-1是cv2.rectangle函數用於建立填充矩形的參數。

 咱們的write函數是局部定義的以便它能夠訪問顏色列表。咱們也能夠用顏色做爲參數，可是那樣咱們就只能用一種顏色。

完成這個函數定義後，如今讓咱們在圖像上繪製邊界框。

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()

 

 

打印時間日誌：

在檢測器的最後，咱們將打印一個日誌，其中包含執行代碼的哪一部分花費了多長時間。這對咱們比較不一樣的超參數如何影響檢測器的速度時很重要。能夠在命令行上執行腳本檢測.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()

 

 

測試目標檢測器：

在終端輸入：

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

 

產生輸出：

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文件中能夠找到在視頻上運行檢測器的代碼。除了一些更改以外，代碼與detector .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

 

咱們在幀上迭代的方式與在圖像上迭代的方式類似。

許多地方都簡化了不少代碼，由於每次只須要處理一個圖像，再也不須要處理批。咱們使用一個元組來代替im_dim_list的張量，在write函數中進行了微小的更改。

每次迭代時咱們使用一個變量frames。而後咱們用這個數字除以從第一個幀開始的時間，打印視頻的FPS。

如今咱們不是使用cv2將檢測圖像寫入磁盤，而是用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     

 

 

結論：

在本系列教程中，咱們從零開始實現了一個目標檢測器，併爲達到這個目標而歡呼。我認爲可以寫出高效的代碼是深度學習實踐者被低估的技能之一。不管你的想法多麼具備革命性，除非你能對它進行測試，不然它毫無用處，爲此你須要有很強的編程技能。

我還意識到，在深度學習中學習任何topic的最佳方法都是實現代碼。當你在閱讀一篇文章的時候一些細微之處你可能會錯過，編程會迫使你注意topic的每一個細微之處。我但願本系列教程可以做爲一個練習，鍛鍊你做爲一個深度學習實踐者的技能。

 

Further Reading

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