本篇笔记主要记录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
除此之外可能还有一些fvcore
、detectron2
等库,直接用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,网页是这里。
以下展示了不同方法的分割效果。 原图 语义分割结果 人体关键点检测结果 全图分割结果
本文作者原创,未经许可不得转载,谢谢配合