使用 Google Colab上的PyTorch YOLOv3

作者|Hervind Philipe 編譯|VK 來源|Towards Data Science

對於計算機視覺愛好者來說，YOLO (You Only Look Once)是一個非常流行的實時對象檢測概念，因爲它非常快，而且性能非常好。

在本文中，我將共享處理視頻的代碼，以獲取谷歌Colab內每個對象的邊框。 我們將不討論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/yolov3
import time
import glob
import torch
import os

import argparse
from sys import platform
%cd yolov3
from models import *
from utils.datasets import *
from utils.utils import *

from IPython.display import HTML
from 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.weights
img_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。一個函數名信箱正在調整圖像的大小，並給圖像填充，因此一個寬度或高度變成416，而另一個小於等於416，但仍然可以被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)

結果

左邊是原始視頻，右邊是使用代碼處理的

試試你自己的視頻

1. 轉到GitHub上的谷歌Colab文件(https://colab.research.google.com/github/vindruid/yolov3-in-colab/blob/master/yolov3_video.ipynb)
2. 上傳你的視頻在input_video文件夾，只需運行最後一個單元格(predict & show video)

原文鏈接：https://towardsdatascience.com/yolov3-pytorch-on-google-colab-c4a79eeecdea

歡迎關注磐創AI博客站： http://panchuang.net/

sklearn機器學習中文官方文檔： http://sklearn123.com/

歡迎關注磐創博客資源彙總站： http://docs.panchuang.net/