這是關於從頭實現YOLO v3檢測器教程的第4部分。最後,我們實現了網絡的轉發。在這一部分中,我們使用對象置信度作爲檢測的閾值,然後使用非最大抑制。
本教程的代碼設計爲在Python 3.5和PyTorch 0.4上運行。教程完整代碼可以在這裏找到 Github repo.
本教程分爲5個部分:
- Part 1 : Understanding How YOLO works
- Part 2 : Creating the layers of the network architecture
- Part 3 : Implementing the the forward pass of the network
- Part 4 (This one): Confidence Thresholding and Non-maximum Suppression
- Part 5 : Designing the input and the output pipelines
先修課程
- 課程1-3部分
- 基本的PyTorch工作知識,包括如何使用nn.Module、nn.sequence和torch.nn.parameter類創建自定義架構。
- NumPy基本知識
在前面的部分中,我們構建了一個模型,該模型在給定輸入圖像的情況下輸出多個目標檢測。準確地說,我們的輸出是一個B×10647×85的張量。B爲批處理圖像的數量,10647爲每幅圖像預測的包圍框的數量,85爲包圍框屬性的數量。
但是,如第1部分所述,我們必須使輸出符合對象分數閾值和非最大抑制,以獲得我將在本文的其餘部分中稱爲真實檢測的內容。 爲此,我們將在文件util.py中創建一個名爲write_results的函數
def write_results(prediction, confidence, num_classes, nms_conf = 0.4):
這些函數將預測、置信度(對象得分閾值)、num_classes(在我們的例子中是80)和nms_conf (NMS IoU閾值)作爲輸入。
對象置信度閾值
我們的預測張量包含有關B x 10647邊界框的信息。 對於每個對象得分低於閾值的邊界框,我們將其每個屬性(表示邊界框的整行)的值設置爲零。
conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)
prediction = prediction*conf_mask
執行非最大抑制
Note: I assume you understand what IoU (Intersection over union) is, and what Non-maximum suppression is. If that is not the case, refer to links at the end of the post).
我們現在擁有的邊界框屬性是由中心座標以及邊界框的高度和寬度描述的。然而,使用每個盒子對角的座標計算兩個盒子的IoU更容易。因此,我們將box的(center x, center y, height, width)轉換爲(top-left corner x, top-left corner y, right-bottom corner x, right-bottom corner y)
box_corner = prediction.new(prediction.shape)
box_corner[:,:,0] = (prediction[:,:,0] - prediction[:,:,2]/2)
box_corner[:,:,1] = (prediction[:,:,1] - prediction[:,:,3]/2)
box_corner[:,:,2] = (prediction[:,:,0] + prediction[:,:,2]/2)
box_corner[:,:,3] = (prediction[:,:,1] + prediction[:,:,3]/2)
prediction[:,:,:4] = box_corner[:,:,:4]
每個圖像中的真實檢測數可能不同。 例如,批量爲3的批次,其中圖像1,2和3分別具有5,2,4個真實檢測。 因此,必須同時對一個圖像進行置信度閾值處理和NMS。 這意味着,我們無法對所涉及的操作進行矢量化,並且必須遍歷prediction的第一維(包含批處理中的圖像索引)。
batch_size = prediction.size(0)
write = False
for ind in range(batch_size):
image_pred = prediction[ind] #圖片張量
#置信度閾值
#NMS
正如前面所描述的,write標誌用於指示我們還沒有初始化output,我們將使用這個張量在整個批處理中收集true的檢測結果。
一旦進入循環,讓我們清理一下。注意,每個邊界框行有85個屬性,其中80個是類得分。此時,我們只關心具有最大值的類的分數。因此,我們從每行中刪除80個類得分,並添加具有最大值的類的索引,以及該類的類得分。
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[:,:5], max_conf, max_conf_score)
image_pred = torch.cat(seq, 1)
還記得我們已經將對象置信度小於閾值的邊界框行設置爲0了嗎?我們把它們去掉。
non_zero_ind = (torch.nonzero(image_pred[:,4]))
try:
image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7)
except:
continue
#PyTorch 0.4 兼容
#由於上述代碼沒有引發異常,所以沒有檢測
#PyTorch 0.4中支持標量
if image_pred_.shape[0] == 0:
continue
try-except塊用於處理 我們沒有檢測到的情況。在這種情況下,我們使用continue來跳過此映像的循環主體的其餘部分。
現在,讓我們在圖像中檢測類。
#獲取圖像中檢測到的各種類
img_classes = unique(image_pred_[:,-1]) #-1索引保存類索引
因爲同一個類可以有多個true的檢測,所以我們使用一個名爲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表示)。
以下代碼在原始代碼文件中縮進了三個塊,但我沒有在此處縮進,因爲此頁面上的空間有限。
#得到一個特定類的檢測
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)
#對探測結果進行排序,使得具有最大對象置信度的條目位於頂部
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) #檢測數量
現在,我們執行NMS。
for i in range(idx):
#得到我們在循環中看到的邊框與後面的所有邊框的IoU
try:
ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
except ValueError:
break
except IndexError:
break
#將所有IoU > treshhold的檢測值歸零
iou_mask = (ious < nms_conf).float().unsqueeze(1)
image_pred_class[i+1:] *= iou_mask
#刪除非零項
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,第二個輸入包含每個邊界框。
如果我們有兩個具有大於閾值的IoU的同一類的邊界框,則消除具有較低類置信度的邊界框。 我們已經整理出了具有更高置信度的邊界框。
在循環體中,下面的行給出了邊界框的IoU,由i索引,所有邊界框的索引都高於i。
ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
每次迭代中,如果任何索引值大於i的邊界框具有大於閾值nms_thresh的IoU(框由i索引),則消除該特定框。
#將所有IoU > treshhold的檢測值歸零
iou_mask = (ious < nms_conf).float().unsqueeze(1)
image_pred_class[i+1:] *= iou_mask
#刪除非零項
non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
image_pred_class = image_pred_class[non_zero_ind]
還要注意,我們已經在try-catch塊中放置了代碼行來計算ious。 這是因爲循環被設計成運行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
"""
#獲取邊界框座標
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]
#獲取交叉矩形座標
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)
#交叉區域
inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * torch.clamp(inter_rect_y2 - inter_rect_y1 + 1, min=0)
#聯合區域
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
Writing the predictions
函數write_results輸出一個形狀爲Dx8的張量。這裏D是所有圖像中的真實檢測值,每一個都用一行表示。每個檢測有8個屬性,即檢測所屬批次中的圖像的索引、4個角座標、物體度得分、最大置信度類的得分、該類的索引。
和以前一樣,我們不初始化輸出張量,除非我們有一個檢測要分配給它。一旦初始化,我們將後續的檢測連接到它。我們使用寫標記來指示張量是否已經初始化。在遍歷類的循環結束時,我們將檢測結果添加到張量output。
batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
#對於圖像中類cls儘可能多的檢測,重複batch_id
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))
在函數結束時,我們檢查output是否已經初始化。如果沒有,那就意味着這批次的任何圖像中都沒有檢測到。在這種情況下,返回0。
try:
return output
except:
return 0
這就是本文的內容。在這篇文章的最後,我們終於有了一個張量形式的預測,它列出了每一個預測作爲它的行。現在剩下的惟一一件事就是創建一個輸入管道,從磁盤讀取圖像、計算預測、在圖像上繪製邊框,然後顯示/寫入這些圖像。這就是我們下一部分要做的 part.
擴展閱讀
- PyTorch tutorial
- IoU
- Non maximum suppresion
- Non-maximum Suppression