linux+(RTX3090)cuda11环境下玩转detectron2

detectron2安装

安装环境

打开detectron2的detectron2/INSTALL.md可以看到

一般用这里推荐的版本安装不会出现问题，所以这里就不安装py和pytorch的最新版本了。

笔者使用的linux环境以及RTX3090GPU，（试着安装了py3.7+pytorch1.8.0+cuda11.1但会报错)，测试成功的环境为：
py3.6+pytorch1.7.1+cuda11.0

这里首先给大家介绍一下pytorch的安装流程（怕大家安装pytorch出现问题），后面再说咱们detectron2的安装流程

pytorch安装流程

为了避免大家在后续的安装流程里pytorch出现问题，这里把这单独拿出来

注意，不同版本的python和cuda，还有不同的linux，win环境安装的pytorch是不一样的，大家需要参考pytorch官网来进行安装。）

点击官网的install按钮，进入安装匹配版本的网页。

我们选择自己的环境，官网就会自动给出一个安装指令。比如这里我们选择用稳定版本stable，linux环境，使用pip安装，cuda是11.1，语言是python，官网就会自动给出一个安装的command

这个command要安装的pytorch版本是1.9.0。但是一般安装最新的pytorch一般会出现兼容性问题（笔者安装1.9.0运行detectron2确实出现了兼容性问题，所以我们安装1.7.1）
点击pytorch官网此页面上的按钮install previous versions of PyTorch

寻找和我们的cuda11，linux环境匹配的1.7.1版本的pytorch

安装时输入这个就可以安装pytorch1.7.1和对应版本的torch附件了

conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=11.0 -c pytorch

好，现在让我们正式进入detectron2的安装流程

detectron2安装流程

笔者的安装环境是linux，主要使用anaconda进行安装和训练

首先用conda创建一个py3.6的环境

conda create -n detect2 python=3.6

再安装pytorch

安装1.7.1版本的pytorch （看上面的pytorch安装流程章节，安装对应的）

conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=11.0 -c pytorch

查看pytorch是否成功安装的方法：服务器中输入

python
import torch    # 如正常则静默
a = torch.Tensor([1.])    # 如正常则静默
a.cuda()    # 如正常则返回"tensor([ 1.], device='cuda:0')"
from torch.backends import cudnn # 如正常则静默
cudnn.is_acceptable(a.cuda())    # 如正常则返回 "True"
exit()

接下来安装opencv-python

pip install opencv-python

安装fvcore

pip install fvcore

然后安装pycocotools
安装pycocotools需要的cython

pip install cython

然后安装pycocotools

pip install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

按照项目提供的 requirements，我们还需要GCC 编译器版本>=4.9

linux查看gcc版本：

gcc -v 

接下来：安装detectron2（可以直接安装，但这里我们还是git拉到本地以后，再从源代码构建比较好，不然环境可能会出现问题）

git clone https://github.com/facebookresearch/detectron2.git
cd detectron2
pip install -e .

（注意上面最后一段结尾不要少那一点）
注意，这时候我们还没有安装好，因为（我也不知道的）一些原因，pip install -e .安装出来的tensorflow和tensorflow-GPU，tensorboard是不兼容的，所以我们再次看向detectron2/INSTALL.md
安装一下官方已经配好的那些环境

python -m pip install detectron2 -f \
https://dl.fbaipublicfiles.com/detectron2/wheels/cu110/torch1.7/index.html

查看python和pytorch版本号

python
import torch
print(torch.__version__)
exit()

此时我们就全部安装好了，我们可以运行demo了。

(注，如果大家的matplolib有报错，可以再安装一下matplotlib==2.2.3)

pip install matplotlib==2.2.3

demo运行测试

cd到detectron2/demo文件夹中，输入这三行代码：
这行大家要把路径改成自己的detectron2路径

python demo.py --config-file /home/fengsiyuan/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \

同样，这里改成自己的测试路径

--input /home/fengsiyuan/ceshi/input/input1.jpg /home/fengsiyuan/ceshi/input/input2.jpg --output /home/fengsiyuan/ceshi/output/ \

这里这行我们选择maskrcnn实例分割模型即可

--opts MODEL.WEIGHTS detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl

实际上我们可以选择很多个模型，比如这里我输入全景分割的模型进行测试的话，也有很好的效果（大家把路径改成自己的，试试全景分割）：
（这同样也需要三行代码）

python demo.py --config-file /home/fengsiyuan/detectron2/configs/COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml \
--input /home/fengsiyuan/ceshi/input/input1.jpg /home/fengsiyuan/ceshi/input/input2.jpg --output /home/fengsiyuan/ceshi/output/ \
--opts MODEL.WEIGHTS detectron2://COCO-PanopticSegmentation/panoptic_fpn_R_101_3x/139514519/model_final_cafdb1.pkl

事实上，参考detectron2/GETTING_STARTED.md，我们可以从model zoo里选择很多个不同的模型进行测试，也可以输入网页图片，视频等进行测试，还可以进行CPU上的测试（具体内容看detectron2的md文件）

detectron2的运行流程就是一个又一个不同的readme.md文件组成的，写的比它的官方文档详细，大家可以仔细看来学习

coco数据集上训练及评估

我们阅读datasets/README.md，这里面给了我们整理自己数据集路径的方法

我们可以在detectron2根目录下运行

./datasets/prepare_for_tests.sh

官方会在datasets文件夹内下一个微小的数据集便于我们测试。此时我们在dev文件夹里找一个sh文件，然后运行

dev/run_*_tests.sh

就可以完成这个微笑数据集的测试（注意，这里run和test之间的*是要你在dev文件夹内选一个sh文件，然后自己填名字的）

当然我们也可以在COCO数据集官网下载完整的COCO2017数据集进行测试（2014看提示好像也可以）（我就下载的完整数据集，因为那个微型数据集我测试的时候出现了问题。）

下载完成后，我在datasets文件夹内按如图摆放：

datasets
	--coco
 		--annotations
 		--train2017
 		--val2017

之后我们就可以进行实例分割测试了。

（大家如果看detectron2/datasets/
readme.md）可以发现，后面还有语义分割PanopticFPN的数据集摆放和步骤，LVIS instance segmentation数据集的摆放和步骤，还有很多很多其他格式的数据集的摆放和步骤，但我按照步骤摆放PanopticFPN数据集的时候（这个数据集COCO官网可以下载），用train_net会提示没有 instances_{train,val}2017.json。大家要是能训练PanopticFPN的话跟我说一下。

接下来我们回到detectron2/GETTING_STARTED.md开始训练。

readme里提供了训练的方法，这里我们选择一个GPU进行训练：
在detectron2的根目录下输入：

python tools/train_net.py \
--num-gpus 1 \
--config-file configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml SOLVER.IMS_PER_BATCH 2 SOLVER.BASE_LR 0.0025

之后会在detectron2的根目录下生成一个output文件夹，里面存放着训练时生成的文件（如果完全训练完成，新生成的模型也会保存到output里）

当然我们也可以选择不同的模型进行训练，比如我们想训练retinanet,就可以输入以下内容

python tools/train_net.py \
--num-gpus 2 \
--config-file configs/COCO-Detection/retinanet_R_50_FPN_3x.yaml SOLVER.IMS_PER_BATCH 2 SOLVER.BASE_LR 0.0025

继续往下看detectron2/GETTING_STARTED.md，这里同样有对模型评估的内容。我们依此进行模型的评估：

这里我先在MODEL_ZOO.md里下载了一个RetinaNet训练好的模型

我模型存放的位置是：/home/fengsiyuan/detectron2/offimodel/model_final_5bd44e.pkl

运行评估代码：（这里第三行，MODEL.WEIGHTS 填自己模型保存的路径）
我们下载的retinanet的权重，这里也用retinanet进行评估

python tools/train_net.py \
--config-file configs/COCO-Detection/retinanet_R_50_FPN_3x.yaml \
--eval-only \
MODEL.WEIGHTS /home/fengsiyuan/detectron2/offimodel/model_final_5bd44e.pkl

就可以在COCO数据集上完成评估了。
当然，你也可以去评估自己训练好的模型。把路径改成你训练好的模型的路径就可以了

参考网址：
detectron2（目标检测框架）无死角玩转-00：目录

使用detectron2 api在一张图片上测试

参考网址：Object Detection with Detectron 2 - PyTorch
在我们安装好的detectron2的根目录下建立一个测试文件.ipynb。(注意，最好先学习一下jupyter notebook的使用方法)

jupyter的一些使用方法： 按M变marddown，按Y退出markdown
按A在上方添加一个代码块，按B在下方添加一个代码块，按两次D，删除一个代码块

接下来我们在根目录下建立的.ipynb文件里开始写代码

#检查一下我们有没有安装好pytorch
import torch, torchvision
print(torch.__version__, torch.cuda.is_available())

#Detectron2 Tutorial.ipynb内含google-lab内存放的一个cv2的函数
#其实可以学一下maskrcnn那样的定好根目录的环境的方法在import，否则像我们这样就只能在根目录下运行ipynb
#这里提示的要升级matplotlibrc不用管
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# 插入需要的公共头文件
import numpy as np
import os, json, cv2, random
import matplotlib.pyplot as plt
%matplotlib inline

# 插入detectron2的接口
#(这些都是保存在detectron2/detectron2文件夹下的一个个py文件)
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, DatasetCatalog
from detectron2.structures import BoxMode

#从这个网址上下载一个图片到咱们的根目录下进行测试
!wget http://images.cocodataset.org/val2017/000000439715.jpg -q -O input.jpg
im = cv2.imread("./input.jpg")
#Detectron2 Tutorial.ipynb是使用的他自己的cv2_imshow函数来展示图片
#咱们这里用plt来展示图片
plt.figure(figsize=(15,7.5))
plt.imshow(im[..., ::-1]) #bgr to rgb

cfg = get_cfg()
#我们这里使用的mask_rcnn_R_50_FPN_3x.yaml进行测试
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  # set threshold for this model

cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
outputs = predictor(im[..., ::-1])
print(outputs["instances"].pred_classes)
# print(outputs["instances"].pred_boxes)

在上面的输出中，我们看到一个数组，它向我们展示了模型所做的预测。 但是什么数字代表哪个等级？

我们使用一个元数据metadata与我们的数据集相关联。 它是一个键值映射，包含数据集的有关信息。 方便我们对数据集进行扩充、评估、可视化、日志记录等。

MetadataCatalog.get(cfg.DATASETS.TRAIN[0])

例如，您可以找到它可以识别的对象：For example you can find which objects it can recognize:

import pandas as pd
modelclasses = MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).thing_classes
df = pd.DataFrame(modelclasses,columns=['Model classes'])
df

现在我们可以通过元数据(metadata)知道每个标签。 让我们可视化来自 modelzoo 的预训练模型的结果。 为此，我们将使用 Visualizer 类。

v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
plt.figure(figsize=(20,10))
plt.imshow(out.get_image()[..., ::-1][..., ::-1])

我们还可以对使用keypoint_rcnn_R_50_FPN_3x.yaml这个模型对这张图片进行关键点检测:

im = cv2.imread("./input.jpg")
cfg = get_cfg()   # fresh config
cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7  # set threshold for this model
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
outputs = predictor(im)
v = Visualizer(im[:,:,::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
plt.figure(figsize=(15,7))
plt.imshow(out.get_image()[:, :, ::-1][..., ::-1])

当然也可以使用panoptic_fpn_R_101_3x.yaml进行语义分割

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml"))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml")
predictor = DefaultPredictor(cfg)
panoptic_seg, segments_info = predictor(im)["panoptic_seg"]
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_panoptic_seg_predictions(panoptic_seg.to("cpu"), segments_info)
plt.figure(figsize=(25,15))
plt.imshow(out.get_image()[:, :, ::-1][..., ::-1])

使用detectron2的api进行气球数据的训练

数据集下载

我们可以自己去https://github.com/matterport/Mask_RCNN/releases上下载balloon_dataset.zip，也可以服务器上用wget指令直接下载这个气球数据集

!wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip

我是在detectron2根目录下下载了这个气球数据集，然后解压到了根目录下

创建读入函数

因为这个气球数据集不是COCO格式的，所以我们要创建一个辅助函数用来读入气球数据，实际上，如果你的数据集是COCO格式的，直接用下面三行就可以读入

#如果你的数据集是 COCO 格式，下面这个单元格可以用以下三行替换：
# 咱们训练气球数据不要写这个！
from detectron2.data.datasets import register_coco_instances
#(注册的训练数据集名,{},训练数据的json文件,训练数据根目录)
register_coco_instances("my_dataset_train", {}, "json_annotation_train.json", "path/to/image/dir")
#(注册的测试数据集名,{},测试数据的json文件,测试数据根目录)
register_coco_instances("my_dataset_val", {}, "json_annotation_val.json", "path/to/image/dir")

读入气球数据的辅助函数:

# 读入气球数据的辅助函数
from detectron2.structures import BoxMode

def get_balloon_dicts(img_dir):
    json_file = os.path.join(img_dir, "via_region_data.json")
    with open(json_file) as f:
        imgs_anns = json.load(f)

    dataset_dicts = []
    for idx, v in enumerate(imgs_anns.values()):
        record = {}
        
        filename = os.path.join(img_dir, v["filename"])
        height, width = cv2.imread(filename).shape[:2]
        
        record["file_name"] = filename
        record["image_id"] = idx
        record["height"] = height
        record["width"] = width
      
        annos = v["regions"]
        objs = []
        for _, anno in annos.items():
            assert not anno["region_attributes"]
            anno = anno["shape_attributes"]
            px = anno["all_points_x"]
            py = anno["all_points_y"]
            poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
            poly = [p for x in poly for p in x]

            obj = {
                "bbox": [np.min(px), np.min(py), np.max(px), np.max(py)],
                "bbox_mode": BoxMode.XYXY_ABS,
                "segmentation": [poly],
                "category_id": 0,
            }
            objs.append(obj)
        record["annotations"] = objs
        dataset_dicts.append(record)
    return dataset_dicts

for d in ["train", "val"]:
    DatasetCatalog.register("balloon_" + d, lambda d=d: get_balloon_dicts("balloon/" + d))
    MetadataCatalog.get("balloon_" + d).set(thing_classes=["balloon"])
balloon_metadata = MetadataCatalog.get("balloon_train")

#用了上面的get_balloon_dicts
#估计传入路径函数的时候，直接就是我们解压在detectron2的根目录下的/ballon/train
dataset_dicts = get_balloon_dicts("balloon/train")
for d in random.sample(dataset_dicts, 3):
    img = cv2.imread(d["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=balloon_metadata, scale=0.5)
    out = visualizer.draw_dataset_dict(d)
    plt.figure(figsize=(15,7))
    plt.imshow(out.get_image()[:, :, ::-1][..., ::-1])

我们写入一下气球训练数据的config信息，然后设置一个trainer进行训练

训练数据

from detectron2.engine import DefaultTrainer

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("balloon_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
cfg.SOLVER.MAX_ITER = 300    # 300 iterations enough for this dataset; Train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # faster, enough for this dataset (default: 512)
#因为气球数据只有一个类别，所以这里class填1
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

我们训练好了以后就可以模型评估了.在OUTPUT_DIR里自动保存了我们的训练模型model_final.pth。我们首先建立我们的评估模型predictor

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")  # path to the model we trained
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set a testing threshold设置测试阈值
predictor = DefaultPredictor(cfg)

运行下方代码，随机抽取一些图片进行预测

from detectron2.utils.visualizer import ColorMode
dataset_dicts = get_balloon_dicts("balloon/val")
for d in random.sample(dataset_dicts, 3):    
    im = cv2.imread(d["file_name"])
    outputs = predictor(im) 
    v = Visualizer(im[:, :, ::-1],
                   metadata=balloon_metadata, 
                   scale=0.5, 
                   instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels. Only available for segmentation models
    )
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    plt.figure(figsize=(15,7))
    plt.imshow(out.get_image()[:, :, ::-1][..., ::-1])

模型评估:

Above we can see that our models performs pretty well! Let’s now evaluate our custom model with Evaluators. Two evaluators can be used:

• COCOEvaluator can evaluate AP (Average Precision) for box detection, instance segmentation and keypoint detection.

• SemSegEvaluator can evaluate semantic segmentation metrics.

Afterwards we’ll use the build_detection_test_loader which returns a torch DataLoader, that loads the given detection dataset.

At last we’ll use the model, evaluated and dataloader within the inference_on_dataset function. It runs the model on the dataloader and evaluates the metric with the evaluator.

# 模型评估代码
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader

evaluator = COCOEvaluator("balloon_val", ("bbox", "segm"), False, output_dir="./output/")
val_loader = build_detection_test_loader(cfg, "balloon_val")
print(inference_on_dataset(trainer.model, val_loader, evaluator))

可以看到评估结果还是很好的

使用detectron2训练自己的COCO格式数据集

有了上面的训练气球的例子，训练我们自己的数据就很容易了

我是使用的自己的coco格式数据集，大家可以在COCO官网上下载2017数据集跟着我一起测试

大家下载好了注意解压到一个文件夹里，摆放格式如下摆放:
coco/annotations/instances_train2017.json
coco/annotations/instances_val2017.json
coco/train2017
coco/val2017

我们还是像之前一样，在detectron2的根目录下再建立一个.ipynb文件，然后开始写代码

import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# 这里我把google colab的cv2_imshow函数删掉了
# 我们用plt来代替cv_imshow
import matplotlib.pyplot as plt
import numpy as np
import cv2

# common libraries
import json
import matplotlib.pyplot as plt
%matplotlib inline

# import some common detectron2 utilities
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.structures import BoxMode
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog

注册我们的数据

# 注意，如果大家是用coco2017来训练的，大家把路径改成对应的COCO2017的路径
# coco/annotations/instances_train2017.json
# coco/annotations/instances_val2017.json
# coco/train2017
# coco/val2017
from detectron2.data.datasets import register_coco_instances
register_coco_instances('self_coco_train', {},'/home/fengsiyuan/detectron2/datasets/coco/annotations/data_train.json','/home/fengsiyuan/detectron2/datasets/coco/train_real')               
register_coco_instances('self_coco_val', {},'/home/fengsiyuan/detectron2/datasets/coco/annotations/data_test.json','/home/fengsiyuan/detectron2/datasets/coco/test')       

查看一下train的元数据metadata

#测试一下注册的train数据的annotations
coco_train_metadata = MetadataCatalog.get("self_coco_train")
train_dataset_dicts = DatasetCatalog.get("self_coco_train")
coco_train_metadata

#测试一下注册的val数据的annotations
coco_val_metadata = MetadataCatalog.get("self_coco_val")
dataset_dicts = DatasetCatalog.get("self_coco_val")
coco_val_metadata

# 我们还可以这样测试一下读入的train数据里的第一个数据[0]
train_dataset_dicts[0]

展示我们读入的图片

import random

for d in random.sample(dataset_dicts, 3):
    img = cv2.imread(d["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=coco_val_metadata, scale=0.5)
    vis = visualizer.draw_dataset_dict(d)
    #cv2_imshow(vis.get_image()[:, :, ::-1])
    #我把cv2给删了，咱们用plt来展示
    #这个figsize是咱们输出的图片的大小，随便标，没关系的
    plt.figure(figsize=(15,7))
    plt.imshow(vis.get_image()[:, :, ::-1][..., ::-1])

这里展示的图片是我的数据集里的图片，大家应该展示的是COCO2017数据集里的图片



设计数据的congfig信息

from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg
from detectron2 import model_zoo

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("self_coco_train",)
cfg.DATASETS.TEST = ("self_coco_val", )
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
#这里是迭代次数，300次远远不够，我绝的至少要3000次
cfg.SOLVER.MAX_ITER = 3000    # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512   #(default: 512)
# 我的数据集是21个种类，COCO2017的话CLASSES改成80！
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 21

cfg

注意！我的数据集是21个种类，COCO2017的话上方的NUM_CLASSES改成80！

开始训练:

import os 

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)

trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
#设置此模型的测试阈值
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.8   # set the testing threshold for this model
cfg.DATASETS.TEST = ("self_coco_val", )
predictor = DefaultPredictor(cfg)

from detectron2.utils.visualizer import ColorMode

for d in random.sample(dataset_dicts, 3):    
    im = cv2.imread(d["file_name"])
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                   metadata=coco_val_metadata, 
                   scale=0.8, 
                   instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels
    )
    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    plt.figure(figsize=(15,7))
    plt.imshow(v.get_image()[:, :, ::-1][..., ::-1])

模型的评估部分:

from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
#evaluator = COCOEvaluator("self_coco_val", ("bbox", "segm"), False, output_dir="./output/")
evaluator = COCOEvaluator("self_coco_val",{}, False, output_dir="./output/")
val_loader = build_detection_test_loader(cfg, "self_coco_val")
print(inference_on_dataset(trainer.model, val_loader, evaluator))
# another equivalent way to evaluate the model is to use `trainer.test`

之后就可以展示出模型的评估数据了

一些进阶内容

tutoial里自己定义的imshow函数

Detectron2 Beginner’s Tutorial里自己定义的的cv2_imshow函数

def cv2_imshow(a):
    	"""A replacement for cv2.imshow() for use in Jupyter notebooks.
    	Args:
    	a : np.ndarray. shape (N, M) or (N, M, 1) is an NxM grayscale image. shape
      	(N, M, 3) is an NxM BGR color image. shape (N, M, 4) is an NxM BGRA color
      	image.
    	"""
    	a = a.clip(0, 255).astype('uint8')
    	# cv2 stores colors as BGR; convert to RGB
    	if a.ndim == 3:
            if a.shape[2] == 4:
           	a = cv2.cvtColor(a, cv2.COLOR_BGRA2RGBA)
            else:
            	a = cv2.cvtColor(a, cv2.COLOR_BGR2RGB)
    	display.display(PIL.Image.fromarray(a))

边训练边评估防止过拟合

#We are importing our own Trainer Module here to use the COCO validation evaluation during training.

from detectron2.engine import DefaultTrainer
from detectron2.evaluation import COCOEvaluator

class CocoTrainer(DefaultTrainer):

  @classmethod
  def build_evaluator(cls, cfg, dataset_name, output_folder=None):

    if output_folder is None:
        os.makedirs("coco_eval", exist_ok=True)
        output_folder = "coco_eval"

    return COCOEvaluator(dataset_name, cfg, False, output_folder)

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = CocoTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()