YOLO v3实现 Part5

这是关于从头实现YOLO v3检测器的教程的第5部分。在上一部分中，我们实现了一个将网络输出转换为检测预测的函数。有了可用的检测器，剩下的就是创建输入和输出管道。

本教程的代码设计为在Python 3.5和PyTorch 0.4上运行。完整代码可以在这里找到 Github repo.

本教程分为5个部分：

1. Part 1 : Understanding How YOLO works
2. Part 2 : Creating the layers of the network architecture
3. Part 3 : Implementing the the forward pass of the network
4. Part 4 : Confidence Thresholding and Non-maximum Suppression
5. Part 5 (This one): Designing the input and the output pipelines

预备知识

1. Part 1-4 of the tutorial.
2. Basic working knowledge of PyTorch, including how to create custom architectures with nn.Module, nn.Sequential and torch.nn.parameter classes.
3. 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     

结论

在本系列教程中，我们从头开始实现了一个对象检测器，并为此欢呼。我仍然认为，能够写出高效的代码是深度学习实践者可能拥有的最被低估的技能之一。不管你的想法多么具有革命性，除非你能测试它，否则它是没有用的。为此，您需要具有强大的编码技能。

我还了解到深入学习任何主题的最佳方法是实现代码。 它会迫使您浏览一下您在阅读论文时可能错过的主题的微妙之处。我希望这个教程系列可以作为一种练习，可以培养你作为深度学习练习者的技能。

扩展阅读

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