实战基于PyTorch的yolo-v3模型

对于计算机视觉爱好者来说，YOLO （You Only Look Once）是一个非常流行的实时目标检测算法，因为它非常快，同时性能非常好。在本文中，我将共享一个视频处理的代码，以获取视频中每个对象目标的边框。我们将不讨论YOLO的概念或架构，因为很多好的文章已经在媒体中详细阐述了这些知识点。这里我们只讨论函数代码。

开始

谷歌Colab地址：https：／／colab．research．google．com／github／vindruid／yolov3－in－colab／blob／master／yolov3＿video．ipynb。yolo的git仓库：https：／／github．com／ultralytics／yolov3。尽管仓库已经包含如何使用YOLOv3的教程，教程只需要运行python detect．py －－source file．mp4，但是我简化了代码，具体在谷歌Colab ／ Jupyter笔记本中。准备YoloV3和LoadModel首先克隆YoloV3仓库，然后导入通用包和repo函数！git clone https：／／github．com／ultralytics／yolov3import timeimport globimport torchimport os

import argparsefrom sys import platform％cd yolov3from models import ＊from utils．datasets import ＊from utils．utils import ＊

from IPython．display import HTMLfrom base64 import b64encode设置参数解析器，初始化设备（CPU ／ CUDA），初始化YOLO模型，然后加载权重。parser ＝ argparse．ArgumentParser（）parser．add＿argument（＇－－cfg＇， type＝str， default＝＇cfg／yolov3－spp．cfg＇， help＝＇＊．cfg path＇）parser．add＿argument（＇－－names＇， type＝str， default＝＇data／coco．names＇， help＝＇＊．names path＇）parser．add＿argument（＇－－weights＇， type＝str， default＝＇weights／yolov3－spp－ultralytics．pt＇， help＝＇weights path＇）parser．add＿argument（＇－－img－size＇， type＝int， default＝416， help＝＇inference size （pixels）＇）parser．add＿argument（＇－－conf－thres＇， type＝float， default＝0．3， help＝＇object confidence threshold＇）parser．add＿argument（＇－－iou－thres＇， type＝float， default＝0．6， help＝＇IOU threshold for NMS＇）parser．add＿argument（＇－－device＇， default＝＇＇， help＝＇device id （i．e． 0 or 0，1） or cpu＇）parser．add＿argument（＇－－classes＇， nargs＝＇＋＇， type＝int， help＝＇filter by class＇）parser．add＿argument（＇－－agnostic－nms＇， action＝＇store＿true＇， help＝＇class－agnostic NMS＇）opt ＝ parser．parse＿args（args ＝ ［］）
weights ＝ opt．weightsimg＿size ＝  opt．img＿size

＃ 初始化设备device ＝ torch＿utils．select＿device（opt．device）

＃ 初始化模型model ＝ Darknet（opt．cfg， img＿size）

＃ 加载权重attempt＿download（weights）if weights．endswith（＇．pt＇）：  ＃ pytorch格式    model．load＿state＿dict（torch．load（weights， map＿location＝device）［＇model＇］）else：  ＃ darknet 格式    load＿darknet＿weights（model， weights）

model．to（device）．eval（）；＃ 获取名字和颜色names ＝ load＿classes（opt．names）colors ＝ ［［random．randint（0， 255） for ＿ in range（3）］ for ＿ in range（len（names））］

％cd ．．我们正在使用YOLOv3－sp－ultralytics权值，该算法称其在平均精度上远远优于其他YOLOv3模型functiontorch＿utils．select＿device（）将自动找到可用的GPU，除非输入是“cpu”对象Darknet在PyTorch上使用预训练的权重来初始化了YOLOv3架构（此时我们不希望训练模型）预测视频中的目标检测接下来，我们将读取视频文件并使用矩阵框重写视频。def predict＿one＿video（path＿video）：    cap  ＝ cv2．VideoCapture（path＿video）    ＿， img0 ＝ cap．read（）

save＿path ＝ os．path．join（output＿dir， os．path．split（path＿video）［－1］）     fps ＝ cap．get（cv2．CAP＿PROP＿FPS）    w ＝ int（cap．get（cv2．CAP＿PROP＿FRAME＿WIDTH））    h ＝ int（cap．get（cv2．CAP＿PROP＿FRAME＿HEIGHT））    vid＿writer ＝ cv2．VideoWriter（save＿path， cv2．VideoWriter＿fourcc（＊＇MP4V＇）， fps， （w， h））我们使用MP4格式写入新的视频，变量为vid＿writer，而宽度和高度则根据原始视频来设置。开始对视频中的每一帧进行循环以获得预测。while img0 is not None：

img ＝ letterbox（img0， new＿shape＝opt．img＿size）［0］

＃ 转换        img ＝ img［：， ：， ：：－1］．transpose（2， 0， 1）  ＃ BGR 到 RGB， 到 3xHxW        img ＝ np．ascontiguousarray（img）        img ＝ torch．from＿numpy（img）．to（device）        img ＝ img．float（）  ＃ uint8 到 fp16／32        img ／＝ 255．0  ＃ 0 － 255 到 0．0 － 1．0        if img．ndimension（） ＝＝ 3：            img ＝ img．unsqueeze（0）

pred ＝ model（img）［0］        ＃ 应用 NMS        pred ＝ non＿max＿suppression（pred， opt．conf＿thres， opt．iou＿thres， classes＝opt．classes， agnostic＝opt．agnostic＿nms）这个模型的图像大小是416，letterbox函数调整了图像的大小，并给图像进行填充，使得图像可以被32整除。第二部分是将图像转换为RGB格式，并将通道设置到第一维，即（C，H，W），然后将图像数据放入设备（GPU或CPU）中，将像素从0－255缩放到0－1。在我们将图像放入模型之前，我们需要使用img．unsqeeze（0）函数，因为我们必须将图像重新格式化为4维（N，C，H，W）， N是图像的数量，在本例中为1。对图像进行预处理后，将其放入模型中得到预测框，但是预测有很多的框，所以我们需要非最大抑制方法来过滤和合并框。

画边界框和标签，然后写入视频我们在NMS之后循环所有的预测（pred）来绘制边界框，但是图像已经被调整为416像素大小了，我们需要使用scale＿coords函数将其缩放为原始大小，然后使用plot＿one＿box函数来绘制框 ＃ 检测        for i， det in enumerate（pred）：  ＃检测每个图片            im0 ＝ img0

if det is not None and len（det）：                ＃ 更改框的大小                det［：， ：4］ ＝ scale＿coords（img．shape［2：］， det［：， ：4］， im0．shape）．round（）

＃ 写入结果                for ＊xyxy， conf， cls in det：                    label ＝ ＇％s ％．2f＇ ％ （names［int（cls）］， conf）                    plot＿one＿box（xyxy， im0， label＝label， color＝colors［int（cls）］）        vid＿writer．write（im0）        ＿， img0 ＝ cap．read（）播放Colab的视频视频在函数predict＿one＿video被写入为Mp4格式，我们压缩成h264格式，所以视频可以在谷歌Colab ／ Jupyter上直接播放。显示原始视频我们使用IPython．display．HTML来显示视频，其宽度为400像素，视频是用二进制读取的。path＿video ＝ os．path．join（＂input＿video＂，＂opera＿house．mp4＂）save＿path ＝ predict＿one＿video（path＿video）

＃ 显示视频mp4 ＝ open（path＿video，＇rb＇）．read（）data＿url ＝ ＂data：video／mp4；base64，＂ ＋ b64encode（mp4）．decode（）HTML（＂＂＂＜video width＝400 controls＞      ＜source src＝＂％s＂ type＝＂video／mp4＂＞＜／video＞＂＂＂ ％ data＿url）压缩和显示处理过的视频OpenCV视频写入器的输出是一个比原始视频大3倍的Mp4视频，它不能在谷歌Colab上显示，解决方案之一是我们对视频进行压缩。我们使用ffmpeg －i ｛save＿path｝ －vcodec libx264 ｛compressed＿path｝path＿video ＝ os．path．join（＂input＿video＂，＂opera＿house．mp4＂）save＿path ＝ predict＿one＿video（path＿video）＃ 压缩视频compressed＿path ＝ os．path．join（＂output＿compressed＂， os．path．split（save＿path）［－1］）os．system（f＂ffmpeg －i ｛save＿path｝ －vcodec libx264 ｛compressed＿path｝＂）

＃显示视频mp4 ＝ open（compressed＿path，＇rb＇）．read（）data＿url ＝ ＂data：video／mp4；base64，＂ ＋ b64encode（mp4）．decode（）HTML（＂＂＂＜video width＝400 controls＞      ＜source src＝＂％s＂ type＝＂video／mp4＂＞＜／video＞＂＂＂ ％ data＿url）结果

左边是原始视频，右边是使用代码处理过的视频试试你自己的视频转到GitHub上的谷歌Colab文件上传你的视频在input＿video文件夹中，只需运行最后一个单元格即可（predict ＆ show video）。