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

??鸽了好久，好久没写文章了，主要是一个技术难题，刚入手确实比较困难，所以一直在调试（打游戏），终于在 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')

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

【学习笔记|使用 Detectron2 中的 Mask-RCNN model 并使用预训练好的模型进行 inference 并提取建筑物轮廓】还是获取数据

#visualize training data my_dataset_train_metadata = https://www.it610.com/article/MetadataCatalog.get("self_coco_train") dataset_dicts = DatasetCatalog.get("self_coco_train") my_dataset_val_medata = https://www.it610.com/article/MetadataCatalog.get("self_coco_val") dataset_dicts1 = DatasetCatalog.get("self_coco_val")import random from detectron2.utils.visualizer import Visualizerfor d in random.sample(dataset_dicts, 3): img = cv2.imread(d["file_name"]) visualizer = Visualizer(img[:, :, ::-1], metadata=https://www.it610.com/article/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 COCOEvaluatorclass 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 = 1000cfg.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=https://www.it610.com/article/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 里面的，然后自己测试了一下，就可以了的。