重磅资源|Pytorch1.0版本的Mask R-CNN的Facebook的官方实现

【导读】Facebook刚刚放出的基于Pytorch1.0版本的Faster R-CNN，Mask R-CNN的benchmark，比detectron更快，准确率更高。

项目地址：

https://github.com/facebookresearch/maskrcnn-benchmark

Faster R-CNN and Mask R-CNN in PyTorch 1.0

这个项目目的是为使用Pytorch1.0的用户提供一个简单的搭建检测和分割模型的必要的底层模块。

要点

网络摄像头和 Jupyter notebook的demo

我们提供了简单的网络摄像头的demo，为你演示如何使用maskrcnn_benchmark进行推理：

cd demo# by default, it runs on the GPU# for best results, use min-image-size 800python webcam.py --min-image-size 800# can also run it on the CPUpython webcam.py --min-image-size 300 MODEL.DEVICE cpu# or change the model that you want to usepython webcam.py --config-file ../configs/caffe2/e2e_mask_rcnn_R_101_FPN_1x_caffe2.py --min-image-size 300 MODEL.DEVICE cpu# in order to see the probability heatmaps, pass --show-mask-heatmapspython webcam.py --min-image-size 300 --show-mask-heatmaps MODEL.DEVICE cpu# by default, it runs on the GPU
# for best results, use min-image-size 800
python webcam.py --min-image-size 800
# can also run it on the CPU
python webcam.py --min-image-size 300 MODEL.DEVICE cpu
# or change the model that you want to use
python webcam.py --config-file ../configs/caffe2/e2e_mask_rcnn_R_101_FPN_1x_caffe2.py --min-image-size 300 MODEL.DEVICE cpu
# in order to see the probability heatmaps, pass --show-mask-heatmaps
python webcam.py --min-image-size 300 --show-mask-heatmaps MODEL.DEVICE cpu

这个demo在 demo/Mask_R-CNN_demo.ipynb可以找到。

安装

查看 INSTALL.md 的安装指令。

模型库和基线版本

预训练模型，基线版本以及和Detectron和mmdetection的对比在MODEL_ZOO.md中。

几行代码实现推理

我们提供了一个帮助类来简化使用预训练模型实现推理的pipline。下面是如何实现，可以在 demo 文件夹中运行：

from maskrcnn_benchmark.config import cfgfrom predictor import COCODemoconfig_file = "../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml"# update the config options with the config filecfg.merge_from_file(config_file)# manual override some optionscfg.merge_from_list(["MODEL.DEVICE", "cpu"])coco_demo = COCODemo(    cfg,    min_image_size=800,    confidence_threshold=0.7,)# load image and then run predictionimage = ...predictions = coco_demo.run_on_opencv_image(image)

config_file = "../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml"

# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cpu"])

coco_demo = COCODemo(
    cfg,
    min_image_size=800,
    confidence_threshold=0.7,
)
# load image and then run prediction
image = ...
predictions = coco_demo.run_on_opencv_image(image)

在COCO数据集上训练的表现

需要先安装 maskrcnn_benchmark，后面的例子才能运行。

需要先下载COCO数据集。我们推荐将coco数据集的软链接放到datasets/里，如下：

我们使用 Detectron中的 minival and valminusminival设置

# symlink the coco datasetcd ~/github/maskrcnn-benchmarkmkdir -p datasets/cocoln -s /path_to_coco_dataset/annotations datasets/coco/annotationsln -s /path_to_coco_dataset/train2014 datasets/coco/train2014ln -s /path_to_coco_dataset/test2014 datasets/coco/test2014ln -s /path_to_coco_dataset/val2014 datasets/coco/val2014
cd ~/github/maskrcnn-benchmark
mkdir -p datasets/coco
ln -s /path_to_coco_dataset/annotations datasets/coco/annotations
ln -s /path_to_coco_dataset/train2014 datasets/coco/train2014
ln -s /path_to_coco_dataset/test2014 datasets/coco/test2014
ln -s /path_to_coco_dataset/val2014 datasets/coco/val2014

你也可以自己设置数据集的路径。那样的话，你需要修改maskrcnn_benchmark/config/paths_catalog.py 指向你自己的数据集. 你也可以创建一个新的 paths_catalog.py 实现同样的两个类，然后在训练的时候作为一个配置参数传到 PATHS_CATALOG 中。

单GPU训练

python /path_to_maskrnn_benchmark/tools/train_net.py --config-file "/path/to/config/file.yaml""/path/to/config/file.yaml"

多GPU训练

我们使用内置的 torch.distributed.launch 实现多GPU的训练。.这个Pytorch中的函数会创建和我们需要使用的GPU数量一样多的进程，每个Python进程使用一个GPU。

export NGPUS=8python -m torch.distributed.launch --nproc_per_node=$NGPUS /path_to_maskrcnn_benchmark/tools/train_net.py --config-file "path/to/config/file.yaml"
python -m torch.distributed.launch --nproc_per_node=$NGPUS /path_to_maskrcnn_benchmark/tools/train_net.py --config-file "path/to/config/file.yaml"

摘要

更多的我们的实现相关的摘要信息，见 ABSTRACTIONS.md.

增加你自己的数据集

这个实现增加了对COCO风格的数据集的支持。在训练时增加新的数据集的支持需要像下面这样做：

from maskrcnn_benchmark.structures.bounding_box import BoxListclass MyDataset(object):    def __init__(self, ...):        # as you would do normally    def __getitem__(self, idx):        # load the image as a PIL Image        image = ...        # load the bounding boxes as a list of list of boxes        # in this case, for illustrative purposes, we use        # x1, y1, x2, y2 order.        boxes = [[0, 0, 10, 10], [10, 20, 50, 50]]        # and labels        labels = torch.tensor([10, 20])        # create a BoxList from the boxes        boxlist = BoxList(boxes, size=image.size, mode="xyxy")        # add the labels to the boxlist        boxlist.add_field("labels", labels)        if self.transforms:            image, boxlist = self.transforms(image, boxlist)        # return the image, the boxlist and the idx in your dataset        return image, boxlist, idx    def get_img_info(self, idx):        # get img_height and img_width. This is used if        # we want to split the batches according to the aspect ratio        # of the image, as it can be more efficient than loading the        # image from disk        return {"height": img_height, "width": img_width}
class MyDataset(object):
    def __init__(self, ...):
        # as you would do normally

    def __getitem__(self, idx):
        # load the image as a PIL Image
        image = ...

        # load the bounding boxes as a list of list of boxes
        # in this case, for illustrative purposes, we use
        # x1, y1, x2, y2 order.
        boxes = [[0, 0, 10, 10], [10, 20, 50, 50]]
        # and labels
        labels = torch.tensor([10, 20])

        # create a BoxList from the boxes
        boxlist = BoxList(boxes, size=image.size, mode="xyxy")
        # add the labels to the boxlist
        boxlist.add_field("labels", labels)

        if self.transforms:
            image, boxlist = self.transforms(image, boxlist)

        # return the image, the boxlist and the idx in your dataset
        return image, boxlist, idx

    def get_img_info(self, idx):
        # get img_height and img_width. This is used if
        # we want to split the batches according to the aspect ratio
        # of the image, as it can be more efficient than loading the
        # image from disk
        return {"height": img_height, "width": img_width}

内容就是这些。你也可以在boxlist中增加额外的字段，例如分割的掩模（使用structures.segmentation_mask.SegmentationMask）, 或者你自己的实例类型。

完整的如何实现 COCODataset 的例子，参见 maskrcnn_benchmark/data/datasets/coco.py.

注意:

上面提到的例子是训练用的，我们还利用cocoApi来计算测试时的准确率，测试数据集目前需要按照cocoApi的要求。

License

maskrcnn-benchmark is released under the MIT license. See LICENSE for additional details.