从零开始Pytorch-YOLOv3【笔记】（五）设计输入和输出流程

前言

上一篇：从零开始Pytorch-YOLOv3【笔记】（四）置信度阈值和非极大值抑制

这一部分，原作者又进行了一部分更新，因此机器之心的翻译有所出入，这里给出原文链接：How to implement a YOLO (v3) object detector from scratch in PyTorch: Part 5

更新的部分：如果你在30/03/2018之前访问过这篇文章，我们将一个大小适中的图片调整为 Darknet 的输入大小的方法就是简单地重新调整尺寸。然而，在最初的实现中，图像会调整大小，保持长宽比完整，并填充左边的部分。例如，如果我们将一个1900 x 1280的图像调整为416 x 415，调整后的图像看起来会是这样的。
在准备输入方面的这种差异导致早期实现的性能略低于原始实现。然而，这篇文章已经被更新以包含在最初的实现中遵循的调整大小的方法。

上一篇我们通过置信度阈值和非极大值抑制过滤得到了张量形式的预测结果，在这一部分，我们将为我们的检测器构建输入和输出流程。

这涉及到从磁盘读取图像，做出预测，使用预测结果在图像上绘制边界框，然后将它们保存到磁盘上。我们也会介绍如何让检测器在相机馈送或视频上实时工作。我们将引入一些命令行标签，以便能使用该网络的各种超参数进行一些实验。接下来就开始吧。

注：这部分需要安装 OpenCV 3。

在我们的检测器文件中创建一个 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 文件：https://raw.githubusercontent.com/ayooshkathuria/YOLO_v3_tutorial_from_scratch/master/data/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 的图像输入格式是（batch x 通道 x 高度 x 宽度），其通道顺序为 RGB。因此，我们在 util.py 中写了一个函数 prep_image 来将 numpy 数组转换成 PyTorch 的输入格式。

在编写这个函数之前，我们必须编写一个函数letter_box来调整图像的大小，保持长宽比的一致性，并用color(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 = (letterbox_image(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（比原始图像维度大，height这一维是原维度的二倍）

#PyTorch Variables for images
im_batches = list(map(prep_image, loaded_ims, [inp_dim for x in range(len(imlist))]))  # Python map() 函数：根据函数对指定序列做映射
#List containing dimensions of original images
im_dim_list = [(x.shape[1], x.shape[0]) for x in loaded_ims]  # img.shape = [height, weight]
im_dim_list = torch.FloatTensor(im_dim_list).repeat(1,2)  # repeat是复制Tensor的对应通道，这里height的维度是之前的2倍

# 创建batches

if CUDA:
    im_dim_list = im_dim_list.cuda()

创建baches

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

检测循环

迭代我们的batches的所有图片，生成shape=D×8的预测结果（write_results函数的输出）。

对于每个 batch，我们都会测量检测所用的时间，即测量获取输入到 write_results 函数得到输出之间所用的时间。在 write_prediction 返回的输出中，其中一个属性是 batch 中图像的索引。我们对这个特定属性执行转换，使其现在能代表 imlist 中图像的索引，该列表包含了所有图像的地址。

在那之后，我们 print 每个检测结果所用的时间以及每张图像中检测到的目标。

如果 write_results 函数在 batch 上的输出是一个 int=0，也就是说没有检测结果，那么我们就使用continue跳过后续对该检测结果的处理。

write 作为给检测输出所属于哪一batch，统一batch的检测结果合并需要用到的锚点，用来初始化每个batch的第一个检测结果。

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()
    with torch.no_grad():
        prediction = model(Variable(batch), 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()       
try:
    output
except NameError:
    print ("No detections were made")
    exit()

torch.cuda.synchronize 这一行是为了确保 CUDA 核与 CPU 同步。否则，一旦 GPU 工作还远未完成并且正在排队，那么 CUDA 核将通过异步调用的方式将控制返回给 CPU。此时 end = time.time() 可能在 GPU 工作实际完成前就 print ，产生计时错误。

现在，所有图像的检测结果都在张量输出中了。让我们在图像上绘制边界框。

在图像上绘制边界框

我们使用一个 try-catch 模块来检查是否存在单个检测结果。如果不存在，就退出程序。

在我们绘制边界框之前，我们的输出张量中包含的预测结果对应的是该网络的输入大小（经过letter_box处理后的填充图像），而不是图像的原始大小。因此，在我们绘制边界框之前，让我们将每个边界框的尺寸属性转换为图像的原始尺寸。

im_dim_list = torch.index_select(im_dim_list, 0, output[:,0].long())

scaling_factor = torch.min(416/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

torch.index_select：获取对应维度，对应索引的Tensor。前面我们专门保留了原始图像维度im_dim_list

现在，我们的坐标符合我们的图像在填充区域的尺寸。然而，在letter_box图像中，我们用缩放因子调整了图像的尺寸(请记住，两个维度都用一个公共因子进行了划分，以保持长宽比)。现在，我们撤消此重新调整以获得原始图像上边界框的坐标。

output[:,1:5] /= scaling_factor

还原为对应原始尺寸的边界框可能会超出原图边界范围。这里使用torch.clamp函数，对边界框的坐标上下线进行定义（下限为0，上限为原图像宽/高）

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 文件，其中包含很多可以随机选择的颜色。

现在让我们编写一个函数来绘制方框。

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


def write(x, results):
    c1 = tuple(x[1:3].int())
    c2 = tuple(x[3:5].int())
    img = results[int(x[0])]
    cls = int(x[-1])
    color = random.choice(colors)  # pallete提供的colors
    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

上面的函数从颜色中随机选择一种颜色绘制一个矩形。它还在边框的左上角创建一个实心矩形，并写入在整个实心矩形上检测到的对象的类。使用 cv2.rectangle 函数的-1参数创建一个实心矩形。

我们定义了这个函数之后，现在就来在图像上画边界框吧。

list(map(lambda x: write(x, loaded_ims), output))

最后，最后，将带有检测结果的图像保存到 det_names 中的地址。

det_names = pd.Series(imlist).apply(lambda x: "{}/det_{}".format(args.det,x.split("/")[-1]))

list(map(cv2.imwrite, det_names, loaded_ims))

总结打印输出时间

在检测器工作结束时，我们会 print 一个总结，其中包含了哪部分代码用了多少执行时间的信息。当我们必须比较不同的超参数对检测器速度的影响方式时，这会很有用。可以在命令行上执行 detection.py 脚本时设置batch 大小、objectness 置信度和 NMS 阈值等超参数。

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

检测结果如下：

报错

cv2.rectangle参数格式不对

cv2.error: OpenCV(4.5.5) :-1: error: (-5:Bad argument) in function 'rectangle'
> Overload resolution failed:
>  - Can't parse 'pt1'. Sequence item with index 0 has a wrong type
>  - Can't parse 'pt1'. Sequence item with index 0 has a wrong type
>  - Can't parse 'rec'. Expected sequence length 4, got 2
>  - Can't parse 'rec'. Expected sequence length 4, got 2

发现cv2.rectangle(img, c1, c2,color, 1)中的c1,c2中的元素是tensor类型。修改如下：

def write(x, results):
    c1 = tuple(x[1:3].int())
    c1 = [x.item() for x in c1]  # add
    c2 = tuple(x[3:5].int())
    c2 = [x.item() for x in c2]  # add
    ...

输出路径没有报错画框后的文件

det_names = pd.Series(imlist).apply(lambda x: "{}/det_{}".format(args.det,x.split("/")[-1])) 这里的文件命名有问题
它的路径是F:\workspace\YOLO_v3_tutorial_from_scratch\det\det_dog-cycle-car.png而代码是根据/进行分割所以没能获取到正常的文件名。改为如下：

`det_names = pd.Series(imlist).apply(lambda x: "{}/det_{}".format(args.det,os.path.split(x)[-1]))`

原文还讲解了如何处理视频实时显示，因为只是与detect.py略有不同，这里就不再描述了，感兴趣的可以自己去看原文。

这篇教程到此就告一段落了，既是教程，也是自己的一个学习记录。本人也在运行复现过程中对代码进行了一些修改。（主要是解决遇到的报错的问题）
欢迎大家批评指正。