使用 Detectron2 中的 Mask-RCNN model 并使用预训练好的模型 进行 inference 并提取建筑物轮廓

  鸽了好久，好久没写文章了，主要是一个技术难题，刚入手确实比较困难，所以一直在调试（打游戏），终于在 Colab 上把 RGB 三通道的图像的建筑物轮廓给跑出来了。先看一组效果图吧！

原图

推理图

轮廓图

卡脖子的难点

  在使用的时候遇到了好多问题，而最关键的还是自己没有仔细看官方的 tutorial 中有两行很小的注释，推理的时候 cfg 需要使用和训练时候一样的参数，并添加额外的参数。所以一直报错。

# Inference should use the config with parameters that are used in training
# cfg now already contains everything we've set previously. We changed it a little bit for inference:

情况是什么呢？就是如果电脑重新启动了或者变量消失了，如果直接用官方 tutorial 中的 inference 代码，发现加载权重还是会错误，要么出现随意乱框的 box ，要么找不到某一个字段，错误是 Cannot find field "pred_masks" in the given Instances!。因为没有加载原来训练时候的 cfg 预设的参数！

流程

  本来不想再贴一次代码的，但是为了完整性就慢慢贴吧，也可以看后面录的视频！

# install dependencies: (use cu101 because colab has CUDA 10.1)
!pip install -U torch==1.5 torchvision==0.6 -f https://download.pytorch.org/whl/cu101/torch_stable.html 
!pip install cython pyyaml==5.1
!pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
import torch, torchvision
print(torch.__version__, torch.cuda.is_available())
!gcc --version
# opencv is pre-installed on colab

首先安装需要的库，pytorch…

!pip install detectron2==0.1.3 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.5/index.html

安装 detectron2，由于是 Colab 环境，所以要加感叹号。

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

# import some common libraries
import numpy as np
import cv2
import random
from google.colab.patches import cv2_imshow

# import some common detectron2 utilities
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
from detectron2.data.catalog import DatasetCatalog

导入一些 packages.

%cd drive/MyDrive/detectron2

切换当前路径

from detectron2.data import MetadataCatalog
from detectron2.data.datasets import register_coco_instances
##第一个参数为自己注册的数据名称，第二个参数无需管，第三个参数为数据的地址
#训练集
register_coco_instances('self_coco_train', {}, 
                        './annotations/instances_train2017.json', 
                       './train2017')
#验证集
register_coco_instances('self_coco_val', {}, 
                        './annotations/instances_val2017.json', 
                       './val2017')

注册 coco 的数据集合，训练集，验证集。后续可以用 MetadataCatalog.get('name') 去设置集合和对应的键值.

#查看一下自己的元数据
##运行后可以看到thing_classes=["自己的数据类别"]
coco_val_metadata = MetadataCatalog.get("self_coco_val")
dataset_dicts = DatasetCatalog.get("self_coco_val")
coco_train_metadata = MetadataCatalog.get("self_coco_train")
dataset_dicts1 = DatasetCatalog.get("self_coco_train")

还是获取数据

#visualize training data
my_dataset_train_metadata = MetadataCatalog.get("self_coco_train")
dataset_dicts = DatasetCatalog.get("self_coco_train")
my_dataset_val_medata = MetadataCatalog.get("self_coco_val")
dataset_dicts1 = DatasetCatalog.get("self_coco_val")

import random
from detectron2.utils.visualizer import Visualizer

for d in random.sample(dataset_dicts, 3):
    img = cv2.imread(d["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=my_dataset_train_metadata, scale=0.5)     #cv2 RGB 需要倒序
    vis = visualizer.draw_dataset_dict(d)                                #字典
    cv2_imshow(vis.get_image()[:, :, ::-1])      

可视化一下数据，看看有没报错而已，可能不需要，但是不想尝试了！

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

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)

定义一个训练器的类。

#需要加载原来训练时候的 cfg 
import os 

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.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 1000

cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR,"model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5   # set the testing threshold for this model
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4
cfg.DATASETS.TEST = ("self_coco_val", )
predictor = DefaultPredictor(cfg)

这一步，最为关键，如果重启了代码执行程序，而你又没有重新执行之前的训练代码，那么 cfg 文件里的一些键是没有这些关键值的，会出现之前所说的两个错误（胡乱的box，找不到instances）。

from detectron2.utils.visualizer import ColorMode
im = cv2.imread('./test_images/144_20.jpg')
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"))

mask_array = outputs['instances'].to("cpu").pred_masks.numpy()
num_instances = mask_array.shape[0]     #有几个目标
mask_array = np.moveaxis(mask_array, 0, -1)  #移动 shape 的尺寸
mask_array_instance = []
output = np.zeros_like(im)          # 创建一个0矩阵
num,h,width = mask_array.shape

#num_instances 为检测到的 instance 数量
for i in range(num_instances):
    mask_array_instance.append(mask_array[:, :, i:(i+1)])
    output = np.where(mask_array_instance[i] == True, 255, output)

cv2.imwriter('your_path',img_array)

然后就可以执行 inference 的代码了。这里其实就是读取图像，然后获取里面的 mask 然后再创建一个矩阵，矩阵里面一些关键位置赋值为 1，即 mask 的位置,代码也是在 issue 里面的，然后自己测试了一下，就可以了的。

原图

推理图

轮廓图

原图

推理图

轮廓图

  后期也需要平滑轮廓，有什么问题可以慢慢看视频（听歌）！

Detectron2 识别建筑物轮廓