如何在PyTorch中从头开始实现YOLO（v3）对象检测器：第5部分
图片来源：Karol Majek。在这里查看他的YOLO v3实时检测视频
这是从头开始实现YOLO v3探测器的教程的第5部分。在最后一部分中，我们实现了一个函数，将网络输出转换为检测预测。手头有一个工作探测器，剩下的就是创建输入和输出管道。
本教程的代码旨在在Python 3.5和PyTorch 0.4上运行。它可以在这个Github回购中找到它的全部内容。
本教程分为5个部分：
-
第5部分（本文）：设计输入和输出管道
先决条件
- 本教程的第1-4部分。
- PyTorch的基本工作知识，包括如何使用nn.Module，nn.Sequential和torch.nn.parameter类创建自定义体系结构。
- OpenCV的基础知识
编辑：如果您在30/03/2018之前访问过这篇文章，我们将一个任意大小的图像调整为Darknet输入大小的方式就是简单地重新调整尺寸。然而，在最初的实现中，调整图像的大小，保持纵横比不变，并填充左边部分。例如，如果我们要将1900 x 1280图像的大小调整为416 x 415，则调整大小的图像将如下所示。
准备输入的这种差异导致早期实现的性能略低于原始性能。但是，该帖子已经更新，以包含原始实施中遵循的调整大小方法
在这一部分中，我们将构建探测器的输入和输出管道。这涉及从磁盘读取图像，进行预测，使用预测在图像上绘制边界框，然后将它们保存到磁盘。我们还将介绍如何让探测器在摄像机输入或视频上实时工作。我们将介绍一些命令行标志，以允许对网络的各种超级参数进行一些实验。让我们开始吧。
注意：您需要为此部件安装OpenCV 3。
detector.py在巡视检测器文件中创建文件。在它的顶部添加neccasary导入。
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数据集中对象的名称。data在检测器目录中创建一个文件夹。同样，如果你在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
初始化网络并加载权重。
#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。因此，我们写的函数prep_image在util.py给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 x 8，函数输出）。
对于每个批次，我们将检测所花费的时间测量为获取输入和生成write_results函数输出之间所花费的时间。在返回的输出中write_prediction，其中一个属性是批量映像的索引。我们以这样一种方式转换该特定属性，即它现在表示图像的索引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 #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同步。否则，一旦GPU作业排队，CUDA内核就会在GPU作业完成之前将控制权返回给CPU（异步调用）。如果end = time.time()在GPU作业实际结束之前打印，则可能会导致误导时间。
现在，我们在Tensor输出中检测到所有图像。让我们在图像上绘制边界框。
在图像上绘制边界框
我们使用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])
如果图像中有太多的边界框，将它们全部用一种颜色绘制可能不是一个好主意。将此文件下载到检测器文件夹。这是一个腌制文件，包含许多颜色可供随机选择。
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。它还在边界框的左上角创建一个填充矩形，并在填充矩形中写入检测到的对象的类。-1函数的参数cv2.rectangle用于创建填充矩形。
我们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()
打印时间摘要
在我们的检测器结束时，我们将打印一个摘要，其中包含代码的哪一部分需要多长时间才能执行。当我们必须比较不同的超参数如何影响探测器的速度时，这很有用。超参数如批量大小，对象性的信心和NMS阈值时，（用传递bs，confidence，nms_thresh分别标记）可以在执行脚本中设置detection.py在命令行上。
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
产生输出
以下代码在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目录中。
在视频/网络摄像头上运行检测器
为了在视频或网络摄像头上运行探测器，代码几乎保持不变，除了我们不必迭代批量，而是迭代视频帧。
可以video.py在github存储库的文件中找到在视频上运行检测器的代码。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
我们以类似于迭代图像的方式迭代帧。
许多地方已经简化了许多代码，因为我们不再需要处理批处理，而是一次只处理一个图像。这是因为一次只能有一个帧。这包括使用元组代替张量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
结论
在这一系列的教程中，我们从头开始实现了一个对象检测器，并为达到这个目标而欢呼。我仍然认为能够生成有效的代码是深度学习从业者可以拥有的最被低估的技能之一。无论你有什么革命性的想法，除非你能测试它，否则没用。为此，您需要具备强大的编码技能。
我还了解到深入学习任何主题的最佳方法是实现代码。它会迫使您浏览一下您在阅读论文时可能错过的主题的微妙但基本的微妙之处。我希望这个教程系列可以作为一种练习，可以培养你作为深度学习练习者的技能。
进一步阅读
Ayoosh Kathuria目前是印度国防研究与发展组织的实习生，他正致力于改善粒状视频中的物体检测。当他不工作时，他正在睡觉或者在他的吉他上玩粉红色弗洛伊德。您可以在LinkedIn上与他联系，或者查看他在GitHub上做的更多内容