基於pytorch的yolo v3的理解(下篇)
參考連結:
原文連結:https://blog.paperspace.com/how-to-implement-a-yolo-object-detector-in-pytorch/
機器之心翻譯連結:https://mp.weixin.qq.com/s/jOcpMR2B3x-Nt2FvM6ThrQ
本文是基於這兩篇文章,再加上我自己的理解寫成的,註釋用紅色標出。
------------------------------------------------------------------------------------
前幾日,機器之心編譯介紹了《從零開始 PyTorch 專案:YOLO v3 目標檢測實現》的前 3 部分,介紹了 YOLO 的工作原理、建立 YOLO 網路層級和實現網路的前向傳播的方法。本文包含了該教程的後面兩個部分,將介紹「置信度閾值設定和非極大值抑制」以及「設計輸入和輸出流程」的方法。總體而言,本教程的目的是使用 PyTorch 實現基於 YOLO v3 的目標檢測器,後者是一種快速的目標檢測演算法。
本教程使用的程式碼需要執行在 Python 3.5 和 PyTorch 0.3 版本之上。你可以在以下連結中找到所有程式碼:https://github.com/ayooshkathuria/YOLO_v3_tutorial_from_scratch
所需背景知識
1. 本教程 1-3 部分
2. 瞭解 PyTorch 基本工作方式,包括使用 nn.Module、nn.Sequential 和 torch.nn.parameter 類建立自定義架構的方式
3. NumPy 基本知識
4. OpenCV 基本知識
如果你缺少這些預備知識,可參閱文末擴充套件閱讀部分了解。
置信度閾值設定和非極大值抑制
在前面 3 部分中,我們已經構建了一個能為給定輸入影象輸出多個目標檢測結果的模型。具體來說,我們的輸出是一個形狀為 B x 10647 x 85 的張量;其中 B 是指一批(batch)中影象的數量,10647 是每個影象中所預測的邊界框的數量
但是,正如第 1 部分所述,我們必須使我們的輸出滿足 objectness 分數閾值和非極大值抑制(NMS),以得到後文所說的「真實(true)」檢測結果。要做到這一點,我們將在 util.py 檔案中建立一個名為 write_results 的函式。
def write_results(prediction, confidence, num_classes, nms_conf = 0.4):
該函式的輸入為預測結果、置信度(objectness 分數閾值)、num_classes(我們這裡是 80)和 nms_conf(NMS IoU 閾值)。
目標置信度閾值
我們的預測張量包含有關 B x 10647 邊界框的資訊。對於有低於一個閾值的 objectness 分數的每個邊界框,我們將其每個屬性的值(表示該邊界框的一整行)都設為零。
conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)#低於置信值的置為0,否則就是本身
prediction = prediction*conf_mask
執行非極大值抑制
注:我假設你已經理解 IoU(Intersection over union)和非極大值抑制(Non-maximum suppression)的含義了。如果你還不理解,請參閱文末提供的連結。
我們現在擁有的邊界框屬性是由中心座標以及邊界框的高度和寬度決定的。但是,使用每個框的兩個對角座標能更輕鬆地計算兩個框的 IoU。所以,我們可以將我們的框的 (中心 x, 中心 y, 高度, 寬度) 屬性轉換成 (左上角 x, 左上角 y, 右下角 x, 右下角 y)。
prediction的結構是:d(batch_size, grid_size*grid_size*num_anchors, bbox_attrs)
box_corner = prediction.new(prediction.shape)
box_corner[:,:,0] = (prediction[:,:,0] - prediction[:,:,2]/2)#xmin
box_corner[:,:,1] = (prediction[:,:,1] - prediction[:,:,3]/2)#ymin
box_corner[:,:,2] = (prediction[:,:,0] + prediction[:,:,2]/2)#xmax
box_corner[:,:,3] = (prediction[:,:,1] + prediction[:,:,3]/2)#ymax
prediction[:,:,:4] = box_corner[:,:,:4]
每張影象中的「真實」檢測結果的數量可能存在差異。比如,一個大小為 3 的 batch 中有 1、2、3 這 3 張影象,它們各自有 5、2、4 個「真實」檢測結果。因此,一次只能完成一張影象的置信度閾值設定和 NMS。也就是說,我們不能將所涉及的操作向量化,而且必須在預測的第一個維度(包含一個 batch 中影象的索引)上迴圈。
batch_size = prediction.size(0)#獲取prediction的第一個維度,batch_size
write = False
for ind in range(batch_size):
image_pred = prediction[ind] #image Tensor
#confidence threshholding
#NMS
如前所述,write 標籤是用於指示我們尚未初始化輸出,我們將使用一個張量來收集整個 batch 的「真實」檢測結果。
進入迴圈後,我們再更清楚地說明一下。注意每個邊界框行都有 85 個屬性,其中 80 個是類別分數。此時,我們只關心有最大值的類別分數。所以,我們移除了每一行的這 80 個類別分數,並且轉而增加了有最大值的類別的索引以及那一類別的類別分數。
#找出每一行最大的max_conf以及max_conf_score的值
max_conf, max_conf_score = torch.max(image_pred[:,5:5+ num_classes], 1)
max_conf = max_conf.float().unsqueeze(1)
max_conf_score = max_conf_score.float().unsqueeze(1)#然後形成seq=預測值位置資訊+最大類別+最大類別的概率值
#其中image_pred的第一個維度是batch_size,第二個維度是(xmin,xmax,ymin,ymax)
seq = (image_pred[:,:5], max_conf, max_conf_score)
image_pred = torch.cat(seq, 1)#1表示橫向拼接
記得我們將 object 置信度小於閾值的邊界框行設為零了嗎?讓我們擺脫它們。
non_zero_ind = (torch.nonzero(image_pred[:,4]))
try:
image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7)#reshape成7,【xmin,ymin,xmax,ymax,confidence,max_conf,max_conf_score】
except:#處理無檢測結果
continue
#For PyTorch 0.4 compatibility
#Since the above code with not raise exception for no detection
#as scalars are supported in PyTorch 0.4
if image_pred_.shape[0] == 0:
continue
其中的 try-except 模組的目的是處理無檢測結果的情況。在這種情況下,我們使用 continue 來跳過對本影象的迴圈。
現在,讓我們獲取一張影象中所檢測到的類別。
#Get the various classes detected in the image
img_classes = unique(image_pred_[:,-1]) # -1 index holds the class index,這裡-1表示一張圖片檢測到的所有的類別,unique表示每個類別僅僅出現一次
因為同一類別可能會有多個「真實」檢測結果,所以我們使用一個名叫 unique 的函式來獲取任意給定影象中存在的類別。
def unique(tensor):
tensor_np = tensor.cpu().numpy()
unique_np = np.unique(tensor_np)
unique_tensor = torch.from_numpy(unique_np)
tensor_res = tensor.new(unique_tensor.shape)
tensor_res.copy_(unique_tensor)
return tensor_res
然後,我們按照類別執行 NMS。
for cls in img_classes:
#perform NMS
一旦我們進入迴圈,我們要做的第一件事就是提取特定類別(用變數 cls 表示)的檢測結果。
注意,以下程式碼在原始程式碼檔案中有 3 格縮排,但因為頁面空間有限,這裡沒有縮排。
#get the detections with one particular class,提取特定類別的檢測結果
cls_mask = image_pred_*(image_pred_[:,-1] == cls).float().unsqueeze(1)
class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze()
image_pred_class = image_pred_[class_mask_ind].view(-1,7)
#sort the detections such that the entry with the maximum objectness
s#confidence is at the top
conf_sort_index = torch.sort(image_pred_class[:,4], descending = True )[1]
image_pred_class = image_pred_class[conf_sort_index]
idx = image_pred_class.size(0) #Number of detections
現在,我們執行 NMS。
for i in range(idx):
#Get the IOUs of all boxes that come after the one we are looking at
#in the loop
try:
ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])#第一個輸入是邊界框行,這是由迴圈中的變數 i 索引的。bbox_iou 的第二個輸入是多個邊界框行構成的張量。
except ValueError:
break
except IndexError:
break
#Zero out all the detections that have IoU > treshhold
iou_mask = (ious < nms_conf).float().unsqueeze(1)
image_pred_class[i+1:] *= iou_mask
#Remove the non-zero entries
non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
image_pred_class = image_pred_class[non_zero_ind].view(-1,7)
這裡,我們使用了函式 bbox_iou。第一個輸入是邊界框行,這是由迴圈中的變數 i 索引的。bbox_iou 的第二個輸入是多個邊界框行構成的張量。bbox_iou 函式的輸出是一個張量,其中包含通過第一個輸入代表的邊界框與第二個輸入中的每個邊界框的 IoU。
如果我們有 2 個同樣類別的邊界框且它們的 IoU 大於一個閾值,那麼就去掉其中類別置信度較低的那個。我們已經對邊界框進行了排序,其中有更高置信度的在上面。
在迴圈部分,下面的程式碼給出了框的 IoU,其中通過 i 索引所有索引排序高於 i 的邊界框。
ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
每次迭代時,如果有邊界框的索引大於 i 且又大於閾值 nms_thresh 的 IoU(與索引為 i 的框),那麼就去掉那個特定的框。[因為iou>nms_conf的話是重複的框]僅僅保留iou<=nms_conf的ious
#Zero out all the detections that have IoU > treshhold
iou_mask = (ious < nms_conf).float().unsqueeze(1)#僅僅保留ious<nms_conf
image_pred_class[i+1:] *= iou_mask
#Remove the non-zero entries
non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
image_pred_class = image_pred_class[non_zero_ind]
還要注意,我們已經將用於計算 ious 的程式碼放在了一個 try-catch 模組中。這是因為這個迴圈在設計上是為了執行 idx 次迭代(image_pred_class 中的行數)。但是,當我們繼續迴圈時,一些邊界框可能會從 image_pred_class 移除。這意味著,即使只從 image_pred_class 中移除了一個值,我們也不能有 idx 次迭代。因此,我們可能會嘗試索引一個邊界之外的值(IndexError),片狀的 image_pred_class[i+1:] 可能會返回一個空張量,從而指定觸發 ValueError 的量。此時,我們可以確定 NMS 不能進一步移除邊界框,然後跳出迴圈。
計算 IoU
這裡是 bbox_iou 函式。
def bbox_iou(box1, box2):
"""
Returns the IoU of two bounding boxes
"""
#Get the coordinates of bounding boxes
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:,0], box1[:,1], box1[:,2], box1[:,3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:,0], box2[:,1], box2[:,2], box2[:,3]
#get the corrdinates of the intersection rectangle
inter_rect_x1 = torch.max(b1_x1, b2_x1)
inter_rect_y1 = torch.max(b1_y1, b2_y1)
inter_rect_x2 = torch.min(b1_x2, b2_x2)
inter_rect_y2 = torch.min(b1_y2, b2_y2)
#Intersection area
inter_area = (inter_rect_x2 - inter_rect_x1 + 1)*(inter_rect_y2 - inter_rect_y1 + 1)
#Union Area
b1_area = (b1_x2 - b1_x1 + 1)*(b1_y2 - b1_y1 + 1)
b2_area = (b2_x2 - b2_x1 + 1)*(b2_y2 - b2_y1 + 1)
iou = inter_area / (b1_area + b2_area - inter_area)
return iou
寫出預測
write_results 函式輸出一個形狀為 Dx8 的張量;其中 D 是所有影象中的「真實」檢測結果,每個都用一行表示。每一個檢測結果都有 8 個屬性,即:該檢測結果所屬的 batch 中影象的索引、4 個角的座標、objectness 分數、有最大置信度的類別的分數、該類別的索引。
如之前一樣,我們沒有初始化我們的輸出張量,除非我們有要分配給它的檢測結果。一旦其被初始化,我們就將後續的檢測結果與它連線起來。我們使用 write 標籤來表示張量是否初始化了。在類別上迭代的迴圈結束時,我們將所得到的檢測結果加入到張量輸出中。
batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
#Repeat the batch_id for as many detections of the class cls in the image
seq = batch_ind, image_pred_class
if not write:
output = torch.cat(seq,1)
write = True
else:
out = torch.cat(seq,1)
output = torch.cat((output,out))
在該函式結束時,我們會檢查輸出是否已被初始化。如果沒有,就意味著在該 batch 的任意影象中都沒有單個檢測結果。在這種情況下,我們返回 0。
try:
return output
except:
return 0
這部分就到此為止了。在這部分結束時,我們終於有了一個張量形式的預測結果,其中以行的形式列出了每個預測。現在還剩下:創造一個從磁碟讀取影象的輸入流程,計算預測結果,在影象上繪製邊界框,然後展示/寫入這些影象。這是下一部分要介紹的內容。
設計輸入和輸出流程
在這一部分,我們將為我們的檢測器構建輸入和輸出流程。這涉及到從磁碟讀取影象,做出預測,使用預測結果在影象上繪製邊界框,然後將它們儲存到磁碟上。我們也會介紹如何讓檢測器在相機饋送或視訊上實時工作。我們將引入一些命令列標籤,以便能使用該網路的各種超引數進行一些實驗。接下來就開始吧。
注:這部分需要安裝 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 successfull