【Mask scoring RCNN】实现目标检测

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.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: 所有labelme的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标注得到的，只有两个点，需要转成四个点
                    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'])  # 解析原图片数据
        # img=io.imread(data['imagePath']) # 通过图片路径打开图片
        # 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 默认为背景
        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文件时报错（不知道为什么）
        # list(map(int,a[1:-1].split(','))) a=annotation['bbox'] 使用该方式转成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)  # 注意，源代码默认为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)  # 画边界线
        # cv2.fillPoly(img, [np.asarray(points)], 1)  # 画多边形 内部像素值为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对应对象，只需计算1对应的行列号（左上角行列号，右下角行列号，就可以算出其边框）
        '''
        # np.where(mask==1)
        index = np.argwhere(mask == 1)
        rows = index[:, 0]
        clos = index[:, 1]
        # 解析左上角行列号
        left_top_r = np.min(rows)  # y
        left_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 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])

代码直接使用

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. 数据和预训练模型

├──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格式文件的代码
下载预训练模型：官网下载链接将下载的预训练模型放在maskscoring_rcnn/pretrained_models文件夹下面

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 print
    CATEGORIES = [
        "__background",
        "object1",#根据自己的数据集修改类别信息
        "object2",
        "object3",
    ]

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

#!/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]) # 打开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 file
        img_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)

        # 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)

运行测试

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 #预测代码