使用 Detectron2 中的 mask-RCNN 检测建筑物轮廓

文章目录[隐藏]

目标分割
目标分割
Mask-RCNN 算法
2. Detectron2
- 跑通 Demo
- 定义自己的数据

前言

兜兜转转试了好多框架与算法，从软件再到代码，中途有许多问题，从 ArcGIS Pro 再到 Tensorflow 的 Mask-RCNN，再到 Detectron2 。

目标检测与目标分割

目标分割

建筑物轮廓的提取属于目标分割（实例分割）的一种，与目标检测不同，不仅要框出目标所在的位置，还要以 mask 的形式框出物体的轮廓。

在这里插入图片描述

上图是 Detectron2 中采用 Mask-RCNN 算法提取建筑物的轮廓，包含了 mask 和 bounding box 以及 possible，当然，这仅仅是 RGB 三通道的数据训练后的结果，也可以加入更多的 channels，包括点云数据、高光谱的数据，数据量大后，也会带来一系列的问题。

目前比较常用的算法包括 U-Net（医学图像分割的鼻祖），Mask-RCNN…

目标分割

目标检测是检测物体所在的位置，并且框出其所在的位置。
在这里插入图片描述

下图是 Yolov5 （目标检测）算法检测出的物体，其并没有将 mask 给标出来。
常见的算法包括 Yolo、SSD…

算法部署

Mask-RCNN 算法

深度学习由于环境配置较为繁琐，除了常见的显卡驱动和 CUDA、CUDNN 的问题之外，还存在 Python 虚拟环境、各种 packages 的版本匹配的问题。

常见的深度学习框架通常用 Pytorch、Tensorflow 写，而 Tensorflow 的显存使用明显比 Pytorch 高，而且，Tensorflow 通常存在 tf 1.x 和 tf2.x 版本的问题，所以之前使用 tensorflow 版本的 mask-RCNN 存在许多问题，

在这里插入图片描述

关于 tensorflow 的 mask-RCNN 在之前的 demo 跑通已经有所记录（Colab 由于文件存于 Drive 中分享不方便，即不分享），而 tensorflow 版本的 mask-RCNN 对于 tensorflow-gpu 和 keras 的版本有较高的要求，配置起来较为麻烦，此外，其 mask 的轮廓提取相较于其他算法的提取较为不准确，所以采用其他框架（Pytorch）的 Mask-RCNN.

2. Detectron2

跑通 Demo

Detectron2 是 facebook 研究部门的一款深度学习的综合集成框架，其包括了目标检测和目标分割的算法，其包括不同格式（coco、PascalVOC）格式的目标检测和目标分割的算法，在其 configs 文件中存在许多文件。主要就是不同算法的 yaml 文件（许可权重）
在这里插入图片描述

关于 Detectron2 算法，官方有一个 Colab 教程文件，在 Github 中搜索 Detectron2 ，在 Getstarted 部分中的 Colab Notebook 部分打开即可跑通官方的代码，当然，官方代码采用的是 coco 训练完成后的权重，自定义数据集采用的是气球文件的数据集，当然训练自定义的数据有很多问题需要去解决，按照官方的文档需要慢慢耐心地解决。

定义自己的数据

本文在编写时采用了武汉大学提供的建筑数据集，由于想要尝试自定义数据集并进行训练以及检测，所以仅仅使用了 50 张 RGB 图像进行标注（labelme工具），原始的图像为 tif文件，
由于一开始尝试的时候使用了 Roboflow 深度学习数据集制作常用网站，而其只支持目标检测数据集的制作（初次尝试所以将 tif 图转成了 jpg 图像，可使用外部工具转化，也可以使用代码自己转化，通常深度学习的框架仅支持 jpg、png、jpeg 格式，所以为了保险也进行了保存）。

定义自己数据的前的准备

由于是使用自己的数据，通常需要自己进行标注，标注的时候需要使用外部的标注软件，对于目标检测可以使用 labelimg 工具（pyqt5 + labelimg 框架）其不支持 rotate box，安装与启动方法以前的文章中也有所介绍。对于目标分割，语义分割来说，需要标注 mask，即生成一个polygon 多边形文件，使用的工具为 labelme，

在这里插入图片描述

上图为 labelimg 工具标注的过程，其支持 PascalVOC、Yolo、CreateML 等格式的数据标注。

在这里插入图片描述

上图为 labelme 工具标注的过程，其生成自己的 labelme json 文件，后期训练后需要转化为 coco 数据的格式，由于建筑比较密集，且建筑物容易受到树的遮挡，在标注时需要自行抉择是否标注。

关于 labelme 和 labelimg 的安装与启动方式：

# labelimg 安装 
conda create -n labelimg python=3.7
conda activate labelme
pip install labeImg
labelImg

需要严格区别大小写，否则无法找到工具报错
在这里插入图片描述

在这里插入图片描述


#安装 labelme 并且启动
conda create -n labelme python=3.7     
conda activate  labelme
pip install labelme
labelme

在这里插入图片描述

都需要在终端（cmd皆可）输入命令，安装并且启动。

标注完成后生成各自的 json 文件，需要使用 labelme2coco.py 将 labelme 的 json 文件转为 coco 自己的 json 文件,可以使用 prepare_detection_dataset代码进行转化，该代码需要修改图像文件的路径以及存放的路径（如果没有定义存放的路径会自动生成，并且将数据分成训练集和验证集）

在这里插入图片描述

完成后，就会生成一个 coco_datacoco 的文件夹，在这里面会根据训练集合的数据分成 images 和 annotations 文件夹，annotation文件夹中会有 instances_train2017.json 文件，以及 instances_val2017.json 文件，images 文件夹下会有 train2017 和 val2017文件里面的 xxx.jpg 和 xxx.json 会放在一起，这样就生成了我们要的前期的数据准备，然后上传到 Google Drive 中，准备开始后期的训练。

训练自己的数据

由于自身的显卡显存有限，所以还是一样在 Google Colab 中进行训练，除了常见的将运行时改为 GPU 外，还有下载 Detectron2 的源码，挂载云盘后

%cd drive/MyDrive
!git clone https://github.com/facebookresearch/detectron2

将 annotations，train2017,val2017 放在根目录下（如果不放也行，只需要在代码的关键地方指过去即可）

按照官方教程安装 pytorch=1.5.

# 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

安装 detectron2

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

上述两步完成后，可能需要 Restart Runtime，然后运行：

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

切换一下当前路径，为后续的 coco 数据的注册切换下路径，也为后期的代码方便。

%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')

注册主要注册训练以及验证数据集，这里注册后会给两个类(OOP 思想)，self_coco_train
类，并且需要提供 coco 标注信息 json 文件所在的位置，以及训练图像的数据。

后面是注册验证集，验证集的 json 文件以及其图像所在的位置。

检查一下训练集合和验证集合是否被正确的读入:

#查看一下自己的元数据
##运行后可以看到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")
coco_val_metadata
coco_train_metadata

在这里插入图片描述

验证后应该会有下面的提示，主要检查下 thing_classes,thing_dataset_id_to_contigous_id,由于是单类别的训练，在 labelme 标注的时候就可以将 id 置为 0，

可视化一下 label 标签是否正确：

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

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)
    vis = visualizer.draw_dataset_dict(d)
    cv2_imshow(vis.get_image()[:, :, ::-1])

在这里插入图片描述

导入分类器并且导入 coco 的评估器

#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)

开始训练

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

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 = ()
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.MODEL.ROI_HEADS.NUM_CLASSES = 4

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

需要注意的是，这里获取 config 文件的时候建议不要使用 ./configs/xxx的路径，一开始尝试的时候遇到了错误，所以在之前将 yaml 文件以及 Base-RCNN-FPN.yaml放在了 detectron2 的根目录下。

在这里插入图片描述

由于数据量较小，所以训练只用了 8 分钟，训练完成后，有一个 pth 权重文件，后续需要加载这个权重文件，并且 inference…

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.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"))
    cv2_imshow(v.get_image()[:, :, ::-1])

然后是 inference，即使用训练完成的权重文件推理。

在这里插入图片描述

至此，小范围的 RGB 三通道的数据训练已经完成，后续，有以下的工作需要完成:

尝试将 Whu Building 的数据集训练一次，过程应该是将其 label 的 tif 文件转化为 labelme 再转化为 coco_json 文件。
引入更多的支撑数据，例如点云数据、高光谱数据，增加 channels。
引入 OSM 数据，学长的建议是尽量能用代码解决就尽量避免用可视化的图形界面解决，所以目前需要去学习 Python 中的 gdal,pyprj,pyshp 地理空间数据库的使用。
最后的目标是大范围的建筑的提取，在 inference 完成之后，需要修改代码文件，将推理完成的图像修改为二值图像，将 mask 置为 1，将其余部分转成 0，即完成栅格化的过程。
即最后的样子应该是这样的二值图像，并且带有地理坐标系和投影坐标系，并且能够适应 OSM 图层（后续可能还有额外的几何校正）。

在这里插入图片描述