使用mask scoring RCNN训练自己的数据集
生活随笔
收集整理的這篇文章主要介紹了
使用mask scoring RCNN训练自己的数据集
小編覺(jué)得挺不錯(cuò)的,現(xiàn)在分享給大家,幫大家做個(gè)參考.
本文主要參考下面兩篇博文，并在部分細(xì)節(jié)處做了修改。
https://blog.csdn.net/XX_123_1_RJ/article/details/102733175?depth_1-utm_source=distribute.pc_relevant.none-task&utm_source=distribute.pc_relevant.none-task
https://blog.csdn.net/linolzhang/article/details/97833354
一、數(shù)據(jù)集準(zhǔn)備
（訓(xùn)練集驗(yàn)證集測(cè)試集的數(shù)據(jù)分別準(zhǔn)備）
1、標(biāo)注數(shù)據(jù)集
大多數(shù)人會(huì)用labelme來(lái)標(biāo)注數(shù)據(jù)集，然后用labelme將每張標(biāo)注圖片都生成一個(gè)json文件。labelme教程網(wǎng)上很多，這里不再贅述。
本人由于原圖的標(biāo)注目標(biāo)很小，用labelme標(biāo)注未免不精確，所以先用PS手動(dòng)標(biāo)注后再寫代碼把標(biāo)注圖轉(zhuǎn)換成了labelme格式的json文件。
結(jié)果如圖：
2、將這些json文件轉(zhuǎn)換成coco格式
這一步我使用如下代碼可成功轉(zhuǎn)換。
# -*- coding:utf-8 -*-
import os, sys
import argparse
import json
import matplotlib.pyplot as plt
import skimage.io as io
from labelme import utils
import numpy as np
import glob
import PIL.Image
class MyEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, np.integer):
return int(obj)
elif isinstance(obj, np.floating):
return float(obj)
elif isinstance(obj, np.ndarray):
return obj.tolist()
else:
return super(MyEncoder, self).default(obj)
class labelme2coco(object):
def __init__(self, labelme_json=[], save_json_path='./tran.json'):
'''
:param labelme_json: 所有l(wèi)abelme的json文件路徑組成的列表
:param save_json_path: json保存位置
'''
self.labelme_json = labelme_json
self.save_json_path = save_json_path
self.images = []
self.categories = []
self.annotations = []
# self.data_coco = {}
self.label = []
self.annID = 1
self.height = 0
self.width = 0
self.save_json()
def data_transfer(self):
for num, json_file in enumerate(self.labelme_json):
with open(json_file, 'r') as fp:
data = json.load(fp) # 加載json文件
self.images.append(self.image(data, num))
for shapes in data['shapes']:
label = shapes['label']
if label not in self.label:
self.categories.append(self.categorie(label))
self.label.append(label)
points = shapes['points'] # 這里的point是用rectangle標(biāo)注得到的，只有兩個(gè)點(diǎn)，需要轉(zhuǎn)成四個(gè)點(diǎn)
points.append([points[0][0], points[1][1]])
points.append([points[1][0], points[0][1]])
self.annotations.append(self.annotation(points, label, num))
self.annID += 1
def image(self, data, num):
image = {}
#img = utils.img_b64_to_arr(data['imageData']) # 解析原圖片數(shù)據(jù)
# img=io.imread(data['imagePath']) # 通過(guò)圖片路徑打開(kāi)圖片
# img = cv2.imread(data['imagePath'], 0)
# height, width = img.shape[:2]
height = data['imageHeight']
width = data['imageWidth']
image['height'] = height
image['width'] = width
image['id'] = num + 1
image['file_name'] = data['imagePath'].split('/')[-1]
self.height = height
self.width = width
return image
def categorie(self, label):
categorie = {}
categorie['supercategory'] = 'Cancer'
categorie['id'] = len(self.label) + 1 # 0 默認(rèn)為背景
categorie['name'] = label
return categorie
def annotation(self, points, label, num):
annotation = {}
annotation['segmentation'] = [list(np.asarray(points).flatten())]
annotation['iscrowd'] = 0
annotation['image_id'] = num + 1
# annotation['bbox'] = str(self.getbbox(points)) # 使用list保存json文件時(shí)報(bào)錯(cuò)（不知道為什么）
# list(map(int,a[1:-1].split(','))) a=annotation['bbox'] 使用該方式轉(zhuǎn)成list
annotation['bbox'] = list(map(float, self.getbbox(points)))
annotation['area'] = annotation['bbox'][2] * annotation['bbox'][3]
# annotation['category_id'] = self.getcatid(label)
annotation['category_id'] = self.getcatid(label) # 注意，源代碼默認(rèn)為1
annotation['id'] = self.annID
return annotation
def getcatid(self, label):
for categorie in self.categories:
if label == categorie['name']:
return categorie['id']
return 1
def getbbox(self, points):
# img = np.zeros([self.height,self.width],np.uint8)
# cv2.polylines(img, [np.asarray(points)], True, 1, lineType=cv2.LINE_AA) # 畫(huà)邊界線
# cv2.fillPoly(img, [np.asarray(points)], 1) # 畫(huà)多邊形 內(nèi)部像素值為1
polygons = points
mask = self.polygons_to_mask([self.height, self.width], polygons)
return self.mask2box(mask)
def mask2box(self, mask):
'''從mask反算出其邊框
mask：[h,w] 0、1組成的圖片
1對(duì)應(yīng)對(duì)象，只需計(jì)算1對(duì)應(yīng)的行列號(hào)（左上角行列號(hào)，右下角行列號(hào)，就可以算出其邊框）
'''
# np.where(mask==1)
index = np.argwhere(mask == 1)
rows = index[:, 0]
clos = index[:, 1]
# 解析左上角行列號(hào)
left_top_r = np.min(rows) # y
left_top_c = np.min(clos) # x
# 解析右下角行列號(hào)
right_bottom_r = np.max(rows)
right_bottom_c = np.max(clos)
# return [(left_top_r,left_top_c),(right_bottom_r,right_bottom_c)]
# return [(left_top_c, left_top_r), (right_bottom_c, right_bottom_r)]
# return [left_top_c, left_top_r, right_bottom_c, right_bottom_r] # [x1,y1,x2,y2]
return [left_top_c, left_top_r, right_bottom_c - left_top_c,
right_bottom_r - left_top_r] # [x1,y1,w,h] 對(duì)應(yīng)COCO的bbox格式
def polygons_to_mask(self, img_shape, polygons):
mask = np.zeros(img_shape, dtype=np.uint8)
mask = PIL.Image.fromarray(mask)
xy = list(map(tuple, polygons))
PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
mask = np.array(mask, dtype=bool)
return mask
def data2coco(self):
data_coco = {}
data_coco['images'] = self.images
data_coco['categories'] = self.categories
data_coco['annotations'] = self.annotations
return data_coco
def save_json(self):
self.data_transfer()
self.data_coco = self.data2coco()
# 保存json文件
json.dump(self.data_coco, open(self.save_json_path, 'w'), indent=4, cls=MyEncoder) # indent=4 更加美觀顯示
if __name__ == '__main__':
src_folder = os.path.abspath(sys.argv[1])
# load src - join json
labelme_json = glob.glob(src_folder + '/*.json')
labelme2coco(labelme_json, sys.argv[2])
在運(yùn)行這個(gè)代碼時(shí)，只有把所有需要的模塊都安裝在anaconda當(dāng)時(shí)安裝labelme的那個(gè)虛擬環(huán)境下才能運(yùn)行成功。
二、環(huán)境搭建（linux）
1、創(chuàng)建pytorch環(huán)境
conda create --name maskrcnn_benchmark source activate maskrcnn_benchmark #所有模塊的安裝都在此虛擬環(huán)境下 conda install ipython pip install ninja yacs cython matplotlib pyqt5 conda install pytorch-nightly torchvision=0.2.1 cudatoolkit=9.0
上面的步驟執(zhí)行完之后還要離線安裝torch1.0.1。因?yàn)槟撤N墻的存在，在線下載torch不太容易實(shí)現(xiàn)，國(guó)內(nèi)鏡像源又沒(méi)有1.0.1這個(gè)版本。而經(jīng)過(guò)博主長(zhǎng)期的踩坑發(fā)現(xiàn)torch1.0.1和torchvision=0.2.1加上numpy1.17才是可用組合。這是torch1.0.1的下載鏈接：http://download.pytorch.org/whl/cu100/torch-1.0.1-cp36-cp36m-linux_x86_64.whl，建議直接迅雷下載。下載完成后，cd到模塊所在目錄然后pip installtorch-1.0.1-cp36-cp36m-linux_x86_64.whl即可。（本人的python是3.6，請(qǐng)酌情修改下載鏈接）
2、安裝cocoapi及apex
export INSTALL_DIR=$PWD # install pycocotools git clone https://github.com/cocodataset/cocoapi.git cd cocoapi/PythonAPI python setup.py build_ext install # install apex cd $INSTALL_DIR git clone https://github.com/NVIDIA/apex.git cd apex python setup.py install --cuda_ext --cpp_ext
3、編譯模型代碼
# install PyTorch Detection cd $INSTALL_DIR #maskrcnn-benchmark #git clone https://github.com/facebookresearch/maskrcnn-benchmark.git git clone https://github.com/zjhuang22/maskscoring_rcnn cd maskscoring_rcnn python setup.py build develop
三、訓(xùn)練前的準(zhǔn)備
1、數(shù)據(jù)和預(yù)訓(xùn)練模型準(zhǔn)備
在下載的maskscoring_rcnn中新建一個(gè)datasets目錄，可按如下結(jié)構(gòu)放置你的json文件和原始圖像
─ datasets
└── annotations
├── coco_train.json
└── coco_test.json
└── coco_train #該文件夾放置訓(xùn)練集的原始圖像
└── coco_test #該文件夾放置測(cè)試集的原始圖像
另外，maskscoring_rcnn的pretrained_models目錄下需要放置R-101.pkl和R-50.pkl這兩個(gè)預(yù)訓(xùn)練模型，如果服務(wù)器連了網(wǎng)，在開(kāi)始訓(xùn)練模型之前會(huì)自動(dòng)下載這兩個(gè)模型，如果服務(wù)器沒(méi)有網(wǎng)就需要手動(dòng)下載放到pretrained_models下了。作者在GitHub也放了有這些模型的百度網(wǎng)盤鏈接。
2、修改參數(shù)
（1）修改 maskscoring_rcnn/configs 目錄下的配置文件，選擇其中的 e2e_ms_rcnn_R_50_FPN_1x.yaml訓(xùn)練腳本，修改如下：
MODEL:
META_ARCHITECTURE: "GeneralizedRCNN"
WEIGHT: "catalog://ImageNetPretrained/MSRA/R-50"
PRETRAINED_MODELS: 'pretrained_models'
DATASETS:
TRAIN: ("coco_train_xxx",) # 1.設(shè)置訓(xùn)練驗(yàn)證集,名字可以隨意起，和其他配置文件對(duì)應(yīng)即可。
TEST: ("coco_val_xxx",)
……（省略數(shù)行）
SOLVER:
BASE_LR: 0.002 #設(shè)置基礎(chǔ)學(xué)習(xí)率，原為0.02
WEIGHT_DECAY: 0.0001
STEPS: (60000, 80000)
MAX_ITER: 5000 #2.設(shè)置最大迭代次數(shù)，可根據(jù)圖片數(shù)量酌情增減，改小也可以更快看到結(jié)果。原為90000
（2）修改 maskscoring_rcnn/maskrcnn_benchmark/config 下的 paths_catalog.py 文件：
DATASETS = {
"coco_2014_train": ( "coco/train2014", "coco/annotations/instances_train2014.json",),
"coco_2014_val": ("coco/val2014", "coco/annotations/instances_val2014.json"),
"coco_2014_minival": ( "coco/val2014", "coco/annotations/instances_minival2014.json", ),
"coco_2014_valminusminival": (
"coco/val2014", "coco/annotations/instances_valminusminival2014.json", ),
#添加自己的數(shù)據(jù)集路徑信息，在相應(yīng)的代碼段后面添加兩行即可
"coco_train_xxx": ("coco_mydata_train", "annotations/coco_mydata_train.json"),
"coco_val_xxx": ("coco_mydata_test", "annotations/coco_mydata_test.json"),
}
（3）修改 maskscoring_rcnn/maskrcnn_benchmark/config 下的 defaults.py 配置文件：
# Size of the smallest side of the image during training
_C.INPUT.MIN_SIZE_TRAIN = 800 # (800,)訓(xùn)練集中圖片的最小邊長(zhǎng)，酌情修改
# Maximum size of the side of the image during training
_C.INPUT.MAX_SIZE_TRAIN = 1333 #訓(xùn)練集中圖片的最大邊長(zhǎng)，酌情修改
# Size of the smallest side of the image during testing
_C.INPUT.MIN_SIZE_TEST = 800 #測(cè)試集中圖片的最小邊長(zhǎng)，酌情修改
# Maximum size of the side of the image during testing
_C.INPUT.MAX_SIZE_TEST = 1333 #測(cè)試集中圖片的最大邊長(zhǎng)，酌情修改
……省略數(shù)行……
_C.MODEL.ROI_BOX_HEAD.NUM_CLASSES = 3 # 修改分類數(shù)量，coco對(duì)應(yīng)81（80+1），注意1加的是背景 _C.SOLVER.BASE_LR = 0.005 # 修改學(xué)習(xí)率，默認(rèn)為0.001 _C.SOLVER.CHECKPOINT_PERIOD = 1000 # 修改check point數(shù)量，根據(jù)需要自定義 _C.SOLVER.IMS_PER_BATCH = 1 # 修改batch size，默認(rèn)16，這個(gè)值要能整除訓(xùn)練集的圖片數(shù)量 _C.TEST.IMS_PER_BATCH = 1 # 修改test batch size，默認(rèn)8，這個(gè)值要能整除測(cè)試集的圖片數(shù)量 _C.OUTPUT_DIR = "models/" # 設(shè)置模型保存路徑（對(duì)應(yīng)自定義文件夾）
四、開(kāi)始訓(xùn)練
到maskscoring_rcnn所在目錄下執(zhí)行：
python tools/train_net.py --config-file configs/e2e_ms_rcnn_R_50_FPN_1x.yaml python tools/test_net.py --config-file configs/e2e_ms_rcnn_R_50_FPN_1x.yaml
在models里面可以查看訓(xùn)練日志。
五、模型預(yù)測(cè)
1、修改maskscoring_rcnn/configs 路徑下的對(duì)應(yīng)的yaml文件的權(quán)重路徑。
MODEL:
META_ARCHITECTURE: "GeneralizedRCNN"
WEIGHT: "models/model_0005000.pth" # 訓(xùn)練好的模型路徑
BACKBONE:
CONV_BODY: "R-50-FPN"
OUT_CHANNELS: 256
2、修改maskscoring_rcnn/demo 路徑下的 predictor.py 文件，添加類別信息。這個(gè)文件在原來(lái)的demo目錄下是沒(méi)有的，從mask rcnn benchmark的demo文件下復(fù)制過(guò)來(lái)即可。
class COCODemo(object):
# COCO categories for pretty print
CATEGORIES = [
"__background",
"cla_a",#根據(jù)自己的數(shù)據(jù)集修改類別信息
"cla_b",
"cla_c",
]
3、在maskscoring_rcnn/demo 下新建 predict.py，用于預(yù)測(cè)。
#!/usr/bin/env python
# coding=UTF-8
import os, sys
import numpy as np
import cv2
from maskrcnn_benchmark.config import cfg
from predictor import COCODemo
# 1.修改后的配置文件
config_file = "configs/e2e_ms_rcnn_R_50_FPN_1x.yaml"
# 2.配置
cfg.merge_from_file(config_file) # merge配置文件
cfg.merge_from_list(["MODEL.MASK_ON", True]) # 打開(kāi)mask開(kāi)關(guān)
cfg.merge_from_list(["MODEL.DEVICE", "cuda"]) # or設(shè)置為CPU ["MODEL.DEVICE", "cpu"]
#cfg.merge_from_list(["MODEL.DEVICE", "cpu"])
coco_demo = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.5, # 3.設(shè)置置信度
)
if __name__ == '__main__':
in_folder = './datasets/test_images/'
out_folder = './datasets/test_images_out/'
if not os.path.exists(out_folder):
os.makedirs(out_folder)
for file_name in os.listdir(in_folder):
if not file_name.endswith(('jpg', 'png')):
continue
# load file
img_path = os.path.join(in_folder, file_name)
image = cv2.imread(img_path)
# method1. 直接得到opencv圖片結(jié)果
#predictions = coco_demo.run_on_opencv_image(image)
#save_path = os.path.join(out_folder, file_name)
#cv2.imwrite(save_path, predictions)
# method2. 獲取預(yù)測(cè)結(jié)果
predictions = coco_demo.compute_prediction(image)
top_predictions = coco_demo.select_top_predictions(predictions)
# draw
img = coco_demo.overlay_boxes(image, top_predictions)
img = coco_demo.overlay_mask(img, predictions)
img = coco_demo.overlay_class_names(img, top_predictions)
save_path = os.path.join(out_folder, file_name)
cv2.imwrite(save_path, img)
# print results
boxes = top_predictions.bbox.numpy()
labels = top_predictions.get_field("labels").numpy() #label = labelList[np.argmax(scores)]
scores = top_predictions.get_field("scores").numpy()
masks = top_predictions.get_field("mask").numpy()
for i in range(len(boxes)):
print('box:', i, ' label:', labels[i])
x1,y1,x2,y2 = [round(x) for x in boxes[i]] # = map(int, boxes[i])
print('x1,y1,x2,y2:', x1,y1,x2,y2)
4、運(yùn)行程序。
python demo/predict.py
在運(yùn)行的過(guò)程中會(huì)報(bào)錯(cuò)找不到文件或者無(wú)法導(dǎo)入相關(guān)的庫(kù)，此時(shí)把相應(yīng)的文件從 mask rcnn benchmark 對(duì)應(yīng)的文件夾復(fù)制過(guò)來(lái)即可。具體操作可參考：https://www.cnblogs.com/littleLittleTiger/p/12582747.html
成功截圖如下
總結(jié)
以上是生活随笔為你收集整理的使用mask scoring RCNN训练自己的数据集的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。
- 上一篇： 一个拥有12年SAP CRM WebCl
- 下一篇： python的基础，特点