Mask scoring RCNN

数据集准备

1. 标注数据集

使用labelme标注原始数据，每一张图片都会生成json文件。labelme标注工具的使用可以查看笔者的博客【labelme】数据标注工具

2. json文件转换为coco格式文件

转换代码如下：

# -*- coding:utf-8 -*-
"""
2020.09.27：alian
将labelme标注的json文件转换成coco格式
"""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.Imageclass 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: 所有labelme的json文件路径组成的列表:param save_json_path: json保存位置'''self.labelme_json = labelme_jsonself.save_json_path = save_json_pathself.images = []self.categories = []self.annotations = []# self.data_coco = {}self.label = []self.annID = 1self.height = 0self.width = 0self.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标注得到的，只有两个点，需要转成四个点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 += 1def image(self, data, num):image = {}#img = utils.img_b64_to_arr(data['imageData'])  # 解析原图片数据# img=io.imread(data['imagePath']) # 通过图片路径打开图片# img = cv2.imread(data['imagePath'], 0)# height, width = img.shape[:2]height = data['imageHeight']width = data['imageWidth']image['height'] = heightimage['width'] = widthimage['id'] = num + 1image['file_name'] = data['imagePath'].split('/')[-1]self.height = heightself.width = widthreturn imagedef categorie(self, label):categorie = {}categorie['supercategory'] = 'Cancer'categorie['id'] = len(self.label) + 1  # 0 默认为背景categorie['name'] = labelreturn categoriedef annotation(self, points, label, num):annotation = {}annotation['segmentation'] = [list(np.asarray(points).flatten())]annotation['iscrowd'] = 0annotation['image_id'] = num + 1# annotation['bbox'] = str(self.getbbox(points)) # 使用list保存json文件时报错（不知道为什么）# list(map(int,a[1:-1].split(','))) a=annotation['bbox'] 使用该方式转成listannotation['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)  # 注意，源代码默认为1annotation['id'] = self.annIDreturn annotationdef getcatid(self, label):for categorie in self.categories:if label == categorie['name']:return categorie['id']return 1def getbbox(self, points):# img = np.zeros([self.height,self.width],np.uint8)# cv2.polylines(img, [np.asarray(points)], True, 1, lineType=cv2.LINE_AA)  # 画边界线# cv2.fillPoly(img, [np.asarray(points)], 1)  # 画多边形 内部像素值为1polygons = pointsmask = self.polygons_to_mask([self.height, self.width], polygons)return self.mask2box(mask)def mask2box(self, mask):'''从mask反算出其边框mask：[h,w]  0、1组成的图片1对应对象，只需计算1对应的行列号（左上角行列号，右下角行列号，就可以算出其边框）'''# np.where(mask==1)index = np.argwhere(mask == 1)rows = index[:, 0]clos = index[:, 1]# 解析左上角行列号left_top_r = np.min(rows)  # yleft_top_c = np.min(clos)  # x# 解析右下角行列号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] 对应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 maskdef data2coco(self):data_coco = {}data_coco['images'] = self.imagesdata_coco['categories'] = self.categoriesdata_coco['annotations'] = self.annotationsreturn data_cocodef 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 jsonlabelme_json = glob.glob(src_folder + '/*.json')labelme2coco(labelme_json, sys.argv[2])

代码直接使用

conda activate labelme  # 激活labelme的运行虚拟环境
cd /maskscoring_rcnn/datasets  # 进入代码所在的文件夹
python  json-coco.py json_train annotations/coco_object_train.json  # json_train 为json的储存路径，annotations/coco_object_train.json为coco格式json文件的保存路径

环境搭建

1. 搭建pytorch环境

conda create -n maskscoring python=3.6
conda  activate maskscoring
conda install ipython
pip install ninja yacs matplotlib pyqt5
conda install pytorch==1.0.0 torchvision==0.2.1 cuda90 -c pytorch

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. 编译maskscoring文件

cd $INSTALL_DIR
git clone https://github.com/zjhuang22/maskscoring_rcnn
cd maskscoring_rcnn
python setup.py build develop

训练

1. 数据和预训练模型

2. 修改训练脚本

总共有三个需要修改的训练脚本
（1）修改 maskscoring_rcnn/configs 目录下的配置文件，选择其中的 e2e_ms_rcnn_R_50_FPN_1x.yaml训练脚本，修改如下：

MODEL:META_ARCHITECTURE: "GeneralizedRCNN"WEIGHT: "catalog://ImageNetPretrained/MSRA/R-50"PRETRAINED_MODELS: 'pretrained_models'
DATASETS:TRAIN: ("coco_train",) # 1.设置训练验证集,名字可以随意起，和其他配置文件对应即可。TEST: ("coco_val",)
……（省略数行）
SOLVER:  BASE_LR: 0.002 #设置基础学习率，原为0.02  WEIGHT_DECAY: 0.0001  STEPS: (60000, 80000)     MAX_ITER: 5000 #2.设置最大迭代次数，可根据图片数量酌情增减，改小也可以更快看到结果。原为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", ),#添加自己的数据集路径信息，在相应的代码段后面添加两行即可"coco_train":("coco_guardrai_train","/home/lianlirong/maskscoring_rcnn/datasets/annotations/coco_guardrai_train.json"),"coco_val":("coco_guardrai_test","/home/lianlirong/maskscoring_rcnn/datasets/annotations/coco_guardrai_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,)训练集中图片的最小边长，酌情修改# Maximum size of the side of the image during training
_C.INPUT.MAX_SIZE_TRAIN = 1333 #训练集中图片的最大边长，酌情修改# Size of the smallest side of the image during testing
_C.INPUT.MIN_SIZE_TEST = 800 #测试集中图片的最小边长，酌情修改# Maximum size of the side of the image during testing
_C.INPUT.MAX_SIZE_TEST = 1333 #测试集中图片的最大边长，酌情修改
……省略数行……
_C.MODEL.ROI_BOX_HEAD.NUM_CLASSES = 3 # 修改分类数量，coco对应81（80+1），注意1加的是背景
_C.SOLVER.BASE_LR = 0.005   # 修改学习率，默认为0.001
_C.SOLVER.CHECKPOINT_PERIOD = 1000  # 修改check point数量，根据需要自定义
_C.SOLVER.IMS_PER_BATCH = 1   # 修改batch size，默认16，这个值要能整除训练集的图片数量
_C.TEST.IMS_PER_BATCH = 1   # 修改test batch size，默认8，这个值要能整除测试集的图片数量
_C.OUTPUT_DIR = "models/"   # 设置模型保存路径（对应自定义文件夹）

开始训练

CUDA_VISIBLE_DEVICES=7 python tools/train_net.py --config-file configs/e2e_mask_rcnn_R_50_FPN_1x.yaml

同时在models文件夹下可以查看log.txt训练日志

测试

在测试阶段也有三处的文件需要修改
（1）修改maskscoring_rcnn/configs 路径下的对应的yaml文件的权重路径。

MODEL:META_ARCHITECTURE: "GeneralizedRCNN"WEIGHT: "models/model_0005000.pth"   # 训练好的模型路径BACKBONE:CONV_BODY: "R-50-FPN"OUT_CHANNELS: 256

（2）修改maskscoring_rcnn/demo 路径下的 predictor.py 文件，添加类别信息。这个文件在原来的demo目录下是没有的，从mask rcnn benchmark的demo文件下复制过来即可。

class COCODemo(object):# COCO categories for pretty printCATEGORIES = ["__background","object1",#根据自己的数据集修改类别信息"object2","object3",]

（3）在maskscoring_rcnn/demo 下新建 predict.py，用于预测。

#!/usr/bin/env python
# coding=UTF-8import 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]) # 打开mask开关
cfg.merge_from_list(["MODEL.DEVICE", "cuda"]) # or设置为CPU ["MODEL.DEVICE", "cpu"]
#cfg.merge_from_list(["MODEL.DEVICE", "cpu"])coco_demo = COCODemo(cfg,min_image_size=800,confidence_threshold=0.5, # 3.设置置信度
)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 fileimg_path = os.path.join(in_folder, file_name)image = cv2.imread(img_path)# method1. 直接得到opencv图片结果#predictions = coco_demo.run_on_opencv_image(image)#save_path = os.path.join(out_folder, file_name)#cv2.imwrite(save_path, predictions)# method2. 获取预测结果predictions = coco_demo.compute_prediction(image)top_predictions = coco_demo.select_top_predictions(predictions)# drawimg = 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 resultsboxes = 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)

运行测试

CUDA_VISIBLE_DEVICES=7 python demo/predict.py

文件夹树形结果如下所示

├──maskscoring_rcnn（项目文件）
　　　　├── datasets # 数据库
　　　　　　　　├──annotations # coco格式文件
　　　　　　　　　　　　├──coco_object_train.json # coco格式文件
　　　　　　　　　　　　├──coco_object_test.json # coco格式文件
　　　　　　　　├──coco_object_train # 训练数据集原始图像
　　　　　　　　├──json_train # 训练数据集json文件
　　　　　　　　├──coco_object_test # 测试数据集原始图像
　　　　　　　　├──json_test # 测试数据集json文件
　　　　　　　　├──json-coco.py # 将json文件转换为coco格式文件的代码
　　　　├──models # 保存训练模型的目录
　　　　　　　　├──log.txt # 训练日志
　　　　├──pretrained_models
　　　　　　　　├──R-50.pkl # 预训练模型
　　　　　　　　├──R-101.pkl # 预训练模型
　　　　├──configs # 配置文件夹
　　　　　　　　├──e2e_mask_rcnn_R_50_FPN_1x.yaml # 需要修改的训练脚本
　　　　　　　　├──e2e_mask_rcnn_R_101_FPN_1x.yaml
　　　　├──maskrcnn_benchmask
　　　　　　　　├──config
　　　　　　　　　　　　├──paths_catalog.py # 需要修改的训练脚本
　　　　　　　　　　　　├──defaults.py # 需要修改的训练脚本
　　　　├──tools # 指令代码
　　　　　　　　├──train_net.py # 训练代码
　　　　　　　　├──test_net.py #测试代码
　　　　├──demo
　　　　　　　　├──predictor.py # 从/maskrcnn-benchmark/demo复制过来的
　　　　　　　　├──predict.py #预测代码