Detectron2的安装与简单使用

Nov 6,2019   2723 words   10 min

Tags: DeepLearning

本篇笔记主要记录Facebook AI Research的Detectron2框架的安装与使用,如Mask R-CNN等卷积网络。

1.安装依赖

(1)Python

Detectron2只支持大于等于3.6的Python版本,建议直接安装对应版本的Anaconda环境即可。Anaconda官网下载地址是这里

这里提醒一下,如果你本机已经有了之前装的一些Python环境,建议在安装时不要把这里的Python路径加入系统环境变量。这样在运行时手动指定下路径即可,可以不破坏现有的环境。一般Python可执行文件以及pip和conda,都在anaconda3/bin目录下面,直接运行即可。

(2)CUDA

Detectron2运行需要GPU环境,因此还需要装好CUDA。关于Ubuntu下安装CUDA开发环境的内容见这篇博客

安装是否成功可以通过在终端中输入nvcc --version判断,如果正确显示了版本则说明安装成功。

(2)PyTorch&torchvision

在PyTorch官网里选择适合你系统的选项,然后会自动生成命令,直接运行即可。例如我的命令是conda install pytorch torchvision cudatoolkit=9.2 -c pytorch

这里依然需要注意的是,如果你没有把Python路径添加到系统变量中,那么conda命令需要指明路径,如/root/anaconda3/bin/conda

(3)OpenCV

可以直接使用pip安装。pip install opencv-python或者pip install opencv-contrib-python即可。

(4)pycocotools

输入以下命令安装即可:

pip install cython; pip install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
(5)GCC

一般已经有了,如果没有直接apt安装一下即可。

2.编译

在终端中输入以下命令即可编译:

git clone git@github.com:facebookresearch/detectron2.git
cd detectron2
python setup.py build develop

除此之外可能还有一些fvcoredetectron2等库,直接用pip即可安装。

3.简单使用

下面代码演示了利用了一个在COCO数据集上训练好的Mask R-CNN网络的使用,需要注意的地方都写在注释里了。

# coding=utf-8

# 加载一些基础包以及设置logger
import detectron2
from detectron2.utils.logger import setup_logger

setup_logger()

# 加载其它一些库
import numpy as np
import cv2

# 加载相关工具
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog

if __name__ == '__main__':
    input_path = "input1.jpg"
    output_path = "result_instance_segmentation.jpg"

    # 指定模型的配置配置文件路径及网络参数文件的路径
    # 对于像下面这样写法的网络参数文件路径,程序在运行的时候就自动寻找,如果没有则下载。
    # Instance segmentation model
    model_file_path = "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
    model_weights = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"

    # 加载图片
    img = cv2.imread(input_path)

    # 创建一个detectron2配置
    cfg = get_cfg()
    # 要创建的模型的名称
    cfg.merge_from_file(model_file_path)
    # 为模型设置阈值
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
    # 加载模型需要的数据
    cfg.MODEL.WEIGHTS = model_weights
    # 基于配置创建一个默认推断
    predictor = DefaultPredictor(cfg)
    # 利用这个推断对加载的影像进行分析并得到结果
    # 对于输出结果格式可以参考这里https://detectron2.readthedocs.io/tutorials/models.html#model-output-format
    outputs = predictor(img)

    # 控制台中输出一些结果
    print(outputs["instances"].pred_classes)
    print(outputs["instances"].pred_boxes)

    # 得到结果后可以使用Visualizer对结果可视化
    # img[:, :, ::-1]表示将BGR波段顺序变成RGB
    # scale表示输出影像的缩放尺度,太小了会导致看不清
    v = Visualizer(img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    # 获得绘制的影像
    result = v.get_image()[:, :, ::-1]
    # 将影像保存到文件
    cv2.imwrite(output_path, result)

输入图像如下。 输出的语义分割结果如下。

此外还有人体关键点检测和全图的语义分割示例,完整代码见Github,点击查看。关于不同网络的名称和设置可以参考Detectron2的Model Zoo,网页是这里

以下展示了不同方法的分割效果。 原图 语义分割结果 人体关键点检测结果 全图分割结果

本文作者原创,未经许可不得转载,谢谢配合

返回顶部