對于計(jì)算機(jī)視覺愛好者來說，YOLO （You Only Look Once）是一個(gè)非常流行的實(shí)時(shí)目標(biāo)檢測算法，因?yàn)樗�非�？�，同時(shí)性能非常好。在本文中，我將共享一個(gè)視頻處理的代碼，以獲取視頻中每個(gè)對象目標(biāo)的邊框。我們將不討論YOLO的概念或架構(gòu)，因?yàn)楹芏嗪玫奈恼乱呀?jīng)在媒體中詳細(xì)闡述了這些知識點(diǎn)。這里我們只討論函數(shù)代碼。

開始

谷歌Colab地址：https：／／colab．research．google．com／github／vindruid／yolov3－in－colab／blob／master／yolov3＿video．ipynb。yolo的git倉庫：https：／／github．com／ultralytics／yolov3。盡管倉庫已經(jīng)包含如何使用YOLOv3的教程，教程只需要運(yùn)行python detect．py －－source file．mp4，但是我簡化了代碼，具體在谷歌Colab ／ Jupyter筆記本中。準(zhǔn)備YoloV3和LoadModel首先克隆YoloV3倉庫，然后導(dǎo)入通用包和repo函數(shù)！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設(shè)置參數(shù)解析器，初始化設(shè)備（CPU ／ CUDA），初始化YOLO模型，然后加載權(quán)重。parser ＝ argparse．ArgumentParser（）parser．a(chǎn)dd＿argument（＇－－cfg＇， type＝str， default＝＇cfg／yolov3－spp．cfg＇， help＝＇＊．cfg path＇）parser．a(chǎn)dd＿argument（＇－－names＇， type＝str， default＝＇data／coco．names＇， help＝＇＊．names path＇）parser．a(chǎn)dd＿argument（＇－－weights＇， type＝str， default＝＇weights／yolov3－spp－ultralytics．pt＇， help＝＇weights path＇）parser．a(chǎn)dd＿argument（＇－－img－size＇， type＝int， default＝416， help＝＇inference size （pixels）＇）parser．a(chǎn)dd＿argument（＇－－conf－thres＇， type＝float， default＝0．3， help＝＇object confidence threshold＇）parser．a(chǎn)dd＿argument（＇－－iou－thres＇， type＝float， default＝0．6， help＝＇IOU threshold for NMS＇）parser．a(chǎn)dd＿argument（＇－－device＇， default＝＇＇， help＝＇device id （i．e． 0 or 0，1） or cpu＇）parser．a(chǎn)dd＿argument（＇－－classes＇， nargs＝＇＋＇， type＝int， help＝＇filter by class＇）parser．a(chǎn)dd＿argument（＇－－agnostic－nms＇， action＝＇store＿true＇， help＝＇class－agnostic NMS＇）opt ＝ parser．parse＿args（args ＝ ［］）
weights ＝ opt．weightsimg＿size ＝  opt．img＿size

＃ 初始化設(shè)備device ＝ torch＿utils．select＿device（opt．device）

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

＃ 加載權(quán)重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權(quán)值，該算法稱其在平均精度上遠(yuǎn)遠(yuǎn)優(yōu)于其他YOLOv3模型functiontorch＿utils．select＿device（）將自動找到可用的GPU，除非輸入是“cpu”對象Darknet在PyTorch上使用預(yù)訓(xùn)練的權(quán)重來初始化了YOLOv3架構(gòu)（此時(shí)我們不希望訓(xùn)練模型）預(yù)測視頻中的目標(biāo)檢測接下來，我們將讀取視頻文件并使用矩陣框重寫視頻。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，而寬度和高度則根據(jù)原始視頻來設(shè)置。開始對視頻中的每一幀進(jìn)行循環(huán)以獲得預(yù)測。while img0 is not None：

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

＃ 轉(zhuǎn)換        img ＝ img［：， ：， ：：－1］．transpose（2， 0， 1）  ＃ BGR 到 RGB， 到 3xHxW        img ＝ np．a(chǎn)scontiguousarray（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］        ＃ 應(yīng)用 NMS        pred ＝ non＿max＿suppression（pred， opt．conf＿thres， opt．iou＿thres， classes＝opt．classes， agnostic＝opt．a(chǎn)gnostic＿nms）這個(gè)模型的圖像大小是416，letterbox函數(shù)調(diào)整了圖像的大小，并給圖像進(jìn)行填充，使得圖像可以被32整除。第二部分是將圖像轉(zhuǎn)換為RGB格式，并將通道設(shè)置到第一維，即（C，H，W），然后將圖像數(shù)據(jù)放入設(shè)備（GPU或CPU）中，將像素從0－255縮放到0－1。在我們將圖像放入模型之前，我們需要使用img．unsqeeze（0）函數(shù)，因?yàn)槲覀儽仨殞�圖像重新格式化為4維（N，C，H，W）， N是圖像的數(shù)量，在本例中為1。對圖像進(jìn)行預(yù)處理后，將其放入模型中得到預(yù)測框，但是預(yù)測有很多的框，所以我們需要非最大抑制方法來過濾和合并框。

畫邊界框和標(biāo)簽，然后寫入視頻我們在NMS之后循環(huán)所有的預(yù)測（pred）來繪制邊界框，但是圖像已經(jīng)被調(diào)整為416像素大小了，我們需要使用scale＿coords函數(shù)將其縮放為原始大小，然后使用plot＿one＿box函數(shù)來繪制框 ＃ 檢測        for i， det in enumerate（pred）：  ＃檢測每個(gè)圖片            im0 ＝ img0

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

＃ 寫入結(jié)果                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的視頻視頻在函數(shù)predict＿one＿video被寫入為Mp4格式，我們壓縮成h264格式，所以視頻可以在谷歌Colab ／ Jupyter上直接播放。顯示原始視頻我們使用IPython．display．HTML來顯示視頻，其寬度為400像素，視頻是用二進(jìn)制讀取的。path＿video ＝ os．path．join（＂input＿video＂，＂opera＿h(yuǎn)ouse．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視頻寫入器的輸出是一個(gè)比原始視頻大3倍的Mp4視頻，它不能在谷歌Colab上顯示，解決方案之一是我們對視頻進(jìn)行壓縮。我們使用ffmpeg －i ｛save＿path｝ －vcodec libx264 ｛compressed＿path｝path＿video ＝ os．path．join（＂input＿video＂，＂opera＿h(yuǎn)ouse．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）結(jié)果

左邊是原始視頻，右邊是使用代碼處理過的視頻試試你自己的視頻轉(zhuǎn)到GitHub上的谷歌Colab文件上傳你的視頻在input＿video文件夾中，只需運(yùn)行最后一個(gè)單元格即可（predict ＆ show video）。