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)
# 需提前下载预训练模型,可从滴滴云内网下载: 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])