从头实现YOLOv3：第5部分

原文为英文，进行了翻译和部分修改，原文地址
代码地址：github仓库、ACgit仓库
相关内容：
YOLOv3论文翻译
YOLOv3原理及流程简述
从头实现YOLOv3：第1部分
从头实现YOLOv3：第2部分
从头实现YOLOv3：第3部分
从头实现YOLOv3：第4部分

第5部分：设计输入和输出管道

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

​ 在这一部分中，我们将构建检测器的输入和输出管道。这涉及从磁盘读取图像，进行预测，使用预测在图像上绘制边界框，然后将它们保存到磁盘。我们还将介绍如何让检测器在摄像头或视频上实时工作。我们将介绍一些命令行标志位，以允许对网络的各种超参数进行一些实验。那么让我们开始吧。

注意：需要安装 OpenCV 3

​ 创建检测器文件detect.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_thresh = float(args.nms_thresh)
start = 0
CUDA = torch.cuda.is_available()

​ 其中，重要的参数是 images（用于指定输入图像或图像目录）、det（保存检测结果的目录）、reso（输入图像的分辨率，可用于速度-精度权衡）、cfg（替代配置文件) 和权重文件。

加载网络

​ 从此处下载文件 coco.names，该文件包含 COCO 数据集中的对象名称。在您的检测器目录中创建一个文件夹 data。

​ 然后，在程序中加载文件。

num_classes = 80  # coco 数据集
classes = load_classes("data/coco.names")

​ load_classes 是在 util.py 中定义的函数，它返回一个字典，该字典将每个类的索引映射到它的name字符串。

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()  # eval() 自动把BN和DropOut固定住，不会取平均，而是用训练好的值

读取输入图像

​ 从磁盘读取图像，或从目录中读取图像。图像的路径存储在名为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)
    # 图像中心是缩小的原图，周边用 128 填充
    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
    """
    # (1 x c x h x w)
    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
# 将加载的图像转换为(批次 x 通道 x 高度 x 宽度)
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()

创建Batch

# 将图像分批次存在 im_batches
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)]

循环检测

​ 迭代所有批次，生成预测，并连接起来（形状D x 8，write_results 函数的输出）。

​ 对于每个批次，测量检测所花费的时间，即从获取输入图像到生成 write_results 函数的输出之间所花费的时间。在 write_prediction 返回的输出中，其中一个属性是批量图像的索引。我们转换该索引属性，使其表示 imlist 中图像的索引，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), CUDA)
    # 目标得分阈值化和非最大值抑制
    prediction = write_result(prediction, confidence, num_classes, nms_conf=nms_thresh)
    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  # 只打印一次
    # transform the attribute from index in batch to index in imlist
    # 将 prediction 中的索引属性转换成 imlist 中的索引
    prediction[:, 0] += i * batch_size

    if not write:  # If we haven'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 作业完成之前，并且 GPU 作业排队时（异步调用），CUDA 内核就会将控制权返回给 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])

​ 如果图像中有太多边界框，用一种颜色绘制它们不太好。将此文件下载到您的检测器文件夹。这是一个pickled文件，其中包含多种颜色可供随机选择。

class_load = time.time()
# RGB颜色列表
colors = pkl.load(open("pallete", "rb"))

​ 现在编写一个函数来绘制检测框。

def write(x, results, color):
    # 坐标
    c1 = tuple([int(x[1]), int(x[2])])
    c2 = tuple([int(x[3]), int(x[4])])
    # 图像
    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, [255, 255, 255], 1)
    return img

​ 上面的函数绘制了一个矩形，颜色是从colors中随机选择的。它还在边界框的左上角创建一个填充矩形，并在填充矩形中写入检测到的目标类别。 cv2.rectangle 函数的 -1 参数用于创建填充矩形。

​ 现在在图像上绘制边界框。

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

测试目标检测器

​ 将测试的图像 dog-cycle-car.png 存在 imgs 目录下，然后在终端运行

python detect.py --images imgs --det det

​ 结果如下

Loading network ......
Network successfully loaded
D:\learning\YOLO\YOLOv3_byme\imgs\cats-dogs.jpg predicted in  0.594 seconds
Objects Detected:    cat cat cat dog
----------------------------------------------------------
D:\learning\YOLO\YOLOv3_byme\imgs\dog-cycle-car.png predicted in  0.561 seconds
Objects Detected:    bicycle truck dog
----------------------------------------------------------
SUMMARY
----------------------------------------------------------
Task                     : Time Taken (in seconds)

Reading addresses        : 0.000
Loading batch            : 0.033
Detection (2 images)     : 1.156
Output Processing        : 0.000
Drawing Boxes            : 0.037
Average time_per_img     : 0.614
----------------------------------------------------------

​ 名为 det_dog-cycle-car.png 的图像保存在 det 目录中。

结论

​ 在本系列教程中，我们从头开始实现了一个目标检测器，很高兴最终实现成功。我仍然认为能够编写出高效的代码是深度学习从业者可以拥有的最被低估的技能之一。无论您的想法多么具有革命性，除非您可以对其进行验证，否则它毫无用处。为此，您需要具备强大的编程技能。

​ 我还了解到，了解深度学习中任何主题的最佳方式是实现代码。它迫使您浏览一个主题的微小而基本的微妙之处，而您在阅读论文时可能会错过这些细节。我希望本系列教程可以作为磨练您作为深度学习从业者技能的练习。