利用detectron2快速使用faster RCNN

Source: Deep Learning on Medium

接著利用detectron2的model zoo中有的pretrain model去進行training,

一開始我認為使用coco的pretrain model去進行這種PCB板的訓練效果不會很好,不過這邊可以看到我只train 3000 epochs,但是最終的AP value還是達到64,若是train 30000 epochs則可以達到75左右,整體來說超乎我的預期,而我還不確定在增加training 的次數會不會再增加AP。

先給大家看predict的結果

tensor([0, 4, 3, 2, 1, 5], device='cuda:0')
Boxes(tensor([[289.1584, 326.0847, 325.2099, 358.8866],
[233.4901, 118.8735, 262.8138, 149.5809],
[106.2966, 181.0337, 140.5350, 211.7205],
[259.9618, 240.3079, 288.6585, 271.2735],
[500.9987, 334.7750, 527.8199, 374.0154],
[347.2850, 146.7405, 395.6854, 180.7809]], device='cuda:0'))

雖然畫質有點差,不過還是可以看見他有準確的偵測到PCB板上的錯誤,並且正確的分類,不過這樣看並不能說出他好的程度,所以我們這邊使用coco的判斷標準,求出裡面的AP value,而detectron2也有設置簡單的evaluator去求出AP、AP50、AP75等等的值。

| AP | AP50 | AP75 | APs | APm | APl |
|:------:|:------:|:------:|:------:|:------:|:------:|
| 64.318 | 93.136 | 79.455 | 62.998 | 65.486 | 70.000 |
[01/16 06:41:56 d2.evaluation.coco_evaluation]: Per-category bbox AP:
| category | AP | category | AP | category | AP |
|:-----------|:-------|:-----------|:-------|:-----------|:-------|
| open | 59.408 | short | 50.631 | mousebite | 63.177 |
| spur | 65.522 | copper | 77.720 | pin-hole | 69.452 |
OrderedDict([('bbox',
{'AP': 64.31825163056139,
'AP-copper': 77.72030131795276,
'AP-mousebite': 63.177111687191676,
'AP-open': 59.407758192556805,
'AP-pin-hole': 69.45151146413728,
'AP-short': 50.63100208959249,
'AP-spur': 65.52182503193735,
'AP50': 93.13612566723208,
'AP75': 79.45463612126157,
'APl': 69.99999999999999,
'APm': 65.48557070350444,
'APs': 62.997574669449286})])

可以看到AP value達到64,與我之前看的A VGG-16 Based Faster RCNN Model for PCB Error Inspection in Industrial AOI Applications論文中的AP相比還要高,當然,使用的dataset不同,不過我經過30000次訓練可以達到AP value 75左右,所以可以感受到resnet相比較於VGG16的強大。

另外,detectron2的backbone model是可以客製化的,所以今天model zoo 裏面沒有VGG16,你也可以自己寫一個沒什麼問題,不過這部份我還沒有嘗試,就不誤導大/p>