點擊上方“AI算法與圖像處理”,選擇加"星標"或“置頂”
重磅乾貨,第一時間送達
導讀
上一篇講到如何安裝MMDetection,今天要分享如何使用 MMDetection 訓練自定義數據集,其實非常簡單!
前言
深度學習發展到現在已經有很多優秀的模型,同時很多大公司也會在內部開發自己的框架,快速的實現業務,從而產生實際價值。
如下面的招聘要求一樣,市場需要這些能熟練使用現有工具快速實現,MMDetection 是一個非常好的選擇。
在本文中,你將知道如何使用定製的數據集推斷、測試和訓練預定義的模型。我們以ballon數據集爲例來描述整個過程。
氣球數據集:https://github.com/matterport/Mask_RCNN/tree/master/samples/balloon
https://github.com/matterport/Mask_RCNN/releases
1、準備自定義數據集
官方教程:https://mmdetection.readthedocs.io/en/latest/2_new_data_model.html
有三種方法在MMDetection中支持新的數據集:
將數據集重新組織爲COCO格式。
將數據集重新組織爲中間格式。
實現一個新的數據集。
官方建議使用前兩種方法,這兩種方法通常比第三種方法簡單。
在本文中,我們給出了一個將數據轉換爲COCO格式的示例。
注意:MMDetection目前只支持評估COCO格式數據集的mask AP。因此,例如實例分割任務,用戶應該將數據轉換爲coco格式。
COCO 標註格式
以下是實例分割所需的COCO格式所需的關鍵,完整的細節請參考這裏。
https://cocodataset.org/#format-data
{
"images": [image],
"annotations": [annotation],
"categories": [category]
}
image = {
"id": int,
"width": int,
"height": int,
"file_name": str,
}
annotation = {
"id": int,
"image_id": int,
"category_id": int,
"segmentation": RLE or [polygon],
"area": float,
"bbox": [x,y,width,height],
"iscrowd": 0 or 1,
}
categories = [{
"id": int,
"name": str,
"supercategory": str,
}]
假設我們使用ballon數據集。下載數據之後,我們需要實現一個函數來將註釋格式轉換爲COCO格式。然後我們可以使用實現的COCODataset加載數據,並執行訓練和評估。
如果你看一下數據集,你會發現數據集的格式如下:
{'base64_img_data': '',
'file_attributes': {},
'filename': '34020010494_e5cb88e1c4_k.jpg',
'fileref': '',
'regions': {'0': {'region_attributes': {},
'shape_attributes': {'all_points_x': [1020,
1000,
994,
1003,
1023,
1050,
1089,
1134,
1190,
1265,
1321,
1361,
1403,
1428,
1442,
1445,
1441,
1427,
1400,
1361,
1316,
1269,
1228,
1198,
1207,
1210,
1190,
1177,
1172,
1174,
1170,
1153,
1127,
1104,
1061,
1032,
1020],
'all_points_y': [963,
899,
841,
787,
738,
700,
663,
638,
621,
619,
643,
672,
720,
765,
800,
860,
896,
942,
990,
1035,
1079,
1112,
1129,
1134,
1144,
1153,
1166,
1166,
1150,
1136,
1129,
1122,
1112,
1084,
1037,
989,
963],
'name': 'polygon'}}},
'size': 1115004}
annotation 是一個JSON文件,其中每個 key 都表示圖像的所有註釋。將ballon數據集轉換爲coco格式的代碼如下所示。
import os.path as osp
import mmcv
def convert_balloon_to_coco(ann_file, out_file, image_prefix):
data_infos = mmcv.load(ann_file)
annotations = []
images = []
obj_count = 0
for idx, v in enumerate(mmcv.track_iter_progress(data_infos.values())):
filename = v['filename']
img_path = osp.join(image_prefix, filename)
width = mmcv.imread(img_path).shape[:2]
images.append(dict(
id = idx,
file_name = filename,
height=height,
width = width))
bboxes = []
labels = []
masks = []
for _, obj in v['regions'].items():
assert not obj['region_attributes']
obj = obj['shape_attributes']
px = obj['all_points_x']
py = obj['all_points_y']
poly = [(x+0.5, y+0.5) for x,y in zip(px,py)]
poly = [p for x in poly for p in x]
y_min, x_max, y_max = (
min(py), max(px),max(py))
data_anno = dict(
image_id = idx,
id = obj_count,
category_id = 0,
bbox = [x_min, y_min, x_max-x_min, y_max-y_min],
area = (x_max - x_min)*(y_max - y_min),
segmentation = [poly],
iscrowd =0)
annotations.append(data_anno)
obj_count += 1
coco_format_json = dict(
images = images,
annotations = annotations,
categories=[{'id':0, 'name':'balloon'}]
)
out_file)
# 對驗證集數據進行處理是,將下面路徑中的train 替換成val 即可
# 注意數據集 balloon 的路徑自行調整
ann_file = './balloon/train/via_region_data.json'
'./balloon/train/annotation_coco.json' =
image_prefix = './balloon/train'
out_file, image_prefix)
註釋:
# 可以加載 json, yaml, pkl 文件
import mmcv
mmcv.load('test.json')
# 刷新位置的進度條方式
mmcv.track_iter_progress(tasks)
參考資料:https://zhuanlan.zhihu.com/p/126725557
https://mmcv.readthedocs.io/en/stable/
通過上面的函數,用戶可以成功地將標註文件轉換成json格式,然後我們可以使用CocoDataset
對模型進行訓練和評估。
2、config文件配置
第二步是準備一個 config,這樣數據集就可以成功加載。假設我們想使用帶有FPN的Mask R-CNN,在balloon數據集上訓練檢測器的配置如下。假設配置在configs/balloon/
目錄下,命名爲mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_balloon.py
。配置如下所示。
# The new config inherits a base config to highlight the necessary modification
_base_ = '../mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_coco.py'
# We also need to change the num_classes in head to match the dataset's annotation
model = dict(
roi_head=dict(
bbox_head=dict(num_classes=1),
mask_head=dict(num_classes=1)))
# Modify dataset related settings
dataset_type = 'COCODataset'
classes = ('balloon',)
data = dict(
train=dict(
img_prefix='balloon/train/',
classes=classes,
ann_file='balloon/train/annotation_coco.json'),
val=dict(
img_prefix='balloon/val/',
classes=classes,
ann_file='balloon/val/annotation_coco.json'),
test=dict(
img_prefix='balloon/val/',
classes=classes,
ann_file='balloon/val/annotation_coco.json'))
# We can use the pre-trained Mask RCNN model to obtain higher performance
load_from = 'checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'
注意:
這裏的_base_ 要修改成
_base_ = '../mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_coco.py'
官方提供的路徑有一點問題
3、自定義數據集上訓練、測試、推理模型
訓練一個新模型
使用新的config 訓練一個模型,直接運行下面的代碼即可:
python tools/train.py configs/balloon/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_balloon.py
如果報錯
raise IOError(f'{filename} is not a checkpoint file')
OSError: checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth is not a checkpoint file
建議去官方提供的預訓練模型下載地址去下載,並放置在checkpoints
文件夾下
https://mmdetection.readthedocs.io/en/latest/model_zoo.html
直接下載:http://download.openmmlab.com/mmdetection/v2.0/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth
注意:
大概需要 9 G 的現存才能跑的起來。。。
測試並推理
測試訓練好的模型,直接運行:
python tools/test.py configs/balloon/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_balloon.py work_dirs/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_balloon.py/latest.pth --eval bbox segm
好的今天的分享就到這裏,如果對你有所幫助,記得三連哈!筆芯,新年快樂
個人微信(如果沒有備註不拉羣!)
請註明:
地區+學校/企業+研究方向+暱稱
下載1:何愷明頂會分享
在「AI算法與圖像處理」公衆號後臺回覆:何愷明,即可下載。總共有6份PDF,涉及 ResNet、Mask RCNN等經典工作的總結分析
下載2:終身受益的編程指南:Google編程風格指南
在「AI算法與圖像處理」公衆號後臺回覆:c++,即可下載。歷經十年考驗,最權威的編程規範!
下載3 CVPR2020
在「AI算法與圖像處理」公衆號後臺回覆:
CVPR2020
,即可下載1467篇CVPR 2020論文
覺得不錯就點亮在看吧
本文分享自微信公衆號 - AI算法與圖像處理(AI_study)。
如有侵權,請聯繫 [email protected] 刪除。
本文參與“OSC源創計劃”,歡迎正在閱讀的你也加入,一起分享。