YOLOV3--pytorch(5.译)

原文地址:https://blog.paperspace.com/how-to-implement-a-yolo-v3-object-detector-from-scratch-in-pytorch-part-5/

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

Note: You will need to install OpenCV 3 for this part.

在 tour detector file中创建文件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是我们执行检测器( to run our detector)的文件,因此我们可以很好的传递给它命令行参数。 我已经使用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数据集中对象的名称。 在检测器目录中创建文件夹数据。

dir 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

初始化network,并加载权重。

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

Create the 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返回的输出中,其中一个属性是批量映像的索引(the index of the image in batch)。 我们以这样一种方式转换该特定属性,即它现在表示imlist中图像的索引,该列表包含所有图像的地址。

之后,我们打印每次检测所花费的时间以及每个图像中检测到的对象。

如果batch的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(异步调用)。 如果在GPU作业实际结束之前打印end = time.time(),这可能会导致误导时间。

现在,我们在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

上面的函数绘制一个矩形,颜色随机选择颜色。 它还在边界框的左上角创建一个填充矩形,并在填充矩形中写入检测到的对象的类。 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,置信度,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

produces the output。

The following code is run on CPU. Expect detection times to be much much faster on GPU. It's around 0.1 sec / image on a Tesla K80.

An image with name det_dog-cycle-car.png is saved in the det directory.

 

在视频/网络摄像头上运行检测器

为了在视频或网络摄像头上运行探测器,代码几乎保持不变,除了我们不必迭代批量,而是迭代视频帧。

用于在视频上运行检测器的代码可以在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

我们以类似于迭代图像的方式迭代帧。

许多地方已经简化了许多代码,因为我们不再需要处理批处理,而是一次只处理一个图像。 这是因为一次只能有一个帧。 这包括使用元组代替im_dim_list的张量和写函数的分钟更改。

每次迭代,我们都会跟踪一个名为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     

 

结论

在这一系列的教程中,我们从头开始实现了一个对象检测器,并为达到这个目标而欢呼。 我仍然认为能够生成有效的代码是深度学习从业者可以拥有的最被低估的技能之一。 无论你有什么革命性的想法,除非你能测试它,否则没用。 为此,您需要具备强大的编码技能。

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

猜你喜欢

转载自blog.csdn.net/qq_38096989/article/details/84784331
今日推荐