Detectron2快速入门(1)

滴滴云技术支持发表于:2020年04月29日 15:35:56

Detectron2 是 Facebook AI Research 推出的基于Pytorch的CV库,它不但进一步集成了最新的目标检测算法,与 mmdetection 一样,Detectron2 也是通过配置文件来设置各种参数,从而实现当前最优的目标检测效果。开发者可以更快的进行迭代模型设计和实验。

  • 包含更多的功能:支持panoptic segmentation(Kaiming He et.al, CVPR2019),densepose,Cascade R-CNN,rotated bounding boxes等等。

  • 可扩展性强:从Detectron2开始,Facebook引入了自定义设计,允许用户更加方便地定制适合自己任务的目标检测器。这种可扩展性使得Detectron2更加灵活。更及时与全面的支持语义分割和全景分割的最新学术成果,而且将一直更新下去。

  • 实现质量:从头开始重写推出的Detectron2,解决了原始Detectron中的几个实现问题,比原始Detectron更快。

在本教程中,您将学习如何安装Detectron2, 并快速运行一个默认的检测模型预测图片

点击左侧+,新建terminal,输入以下命令安装环境:

# 查看当前安装的所有环境
conda info -e 
# 切换到目标环境,XXX为要切换到的pytorch环境,比如:ipykernel_py3_pytorch1.4
conda activate XXX 
# 安装detectron2依赖库
pip install cython pyyaml==5.1
pip install pycocotools
# 安装detectron2
pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/index.html
# 建议下载一份detectron2源码方便阅读,可以直接从官网clone,也可以从滴滴云内网下载(代码更新至2020年4月22日),内网下载方式:
wget https://dataset-public.s3-internal.didiyunapi.com/DAI教程/detectron2.zip
unzip detectron2.zip
import torch, torchvision
print(torch.__version__, torch.cuda.is_available())
# 确认torch为1.4.0,可用GPU为True
1.4.0 True
# 设置 detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

import numpy as np
import cv2
import random

from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog

import matplotlib.pyplot as plt
def cv2_imshow(img_bgr):
    plt.rcParams['figure.figsize'] = (18, 36)
    plt.axis('off')
    plt.imshow(img_bgr[...,::-1])
# 首先从COCO数据集随机下载一张图片
# !wget http://images.cocodataset.org/val2017/000000439715.jpg -O input.jpg
im = cv2.imread("./input.jpg")
cv2_imshow(im)

0015ea92e46f010c93deaf4e87f19dc

# 需提前下载预训练模型,可从滴滴云内网下载:
wget https://dataset-public.s3-internal.didiyunapi.com/DAI教程/pretrained_models/model_final_b275ba.pkl
# 记录当前模型文件路径,修改到cfg.MODEL.WEIGHTS
pwd
# 创建detectron2 默认config和默认Predictor,预测这张图
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # 设定阈值
cfg.MODEL.WEIGHTS = '/home/dc2-user/Detectron2/model_final_b275ba.pkl'
predictor = DefaultPredictor(cfg)
outputs = predictor(im)
# 查看输出
outputs["instances"].pred_classes
outputs["instances"].pred_boxes
Boxes(tensor([[123.2132, 248.8964, 461.1610, 480.0000],
        [254.3527, 156.5582, 342.6227, 410.8582],
        [  0.6390, 282.8371,  76.0437, 480.0000],
        [ 49.8697, 274.1834,  80.6168, 347.8782],
        [561.1337, 271.8238, 595.1155, 379.5087],
        [112.6202, 265.1960, 151.4431, 396.9056],
        [386.0321, 271.5787, 412.9495, 305.1869],
        [344.9710, 269.1204, 391.0746, 298.7343],
        [335.6159, 252.1648, 415.5287, 277.3849],
        [508.7762, 267.0611, 571.7914, 295.6371],
        [329.6030, 230.0288, 394.8186, 258.0396],
        [522.6676, 279.2706, 562.2932, 367.9320],
        [406.3995, 272.7451, 460.7343, 344.4586],
        [328.6352, 230.9069, 406.6131, 282.3652]], device='cuda:0'))
# 可视化预测结果
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2_imshow(v.get_image()[:, :, ::-1])

0015ea92e547c1f030e113bdd7727bb