segment anything进行高分辨率遥感影像分割及效果对比_segment anything 边界分类

一、SAM模型介绍

1. Segement Anything Model，简称SAM，是4月初Meta 发布的史上首个图像分割基础模型。它是通过将三个相互关联的元素：Task、Model和Data结合，来构建形成的一个大模型。Task的构成如下图所示，通过输入分割提示和图片，经模型运算生成掩膜
   
   SAM的输入提示可以是标记点、规则/不规则框选边界、或者输入词语。如输入“Cat”，模型就会识别出来猫，并生成一个掩膜。

Model中主要运算的是词语和图像，架构如下图所示：

最后是Data，它是作为输入用于模型训练，训练过程中对数据进行注释，实现模型优化。SAM对于Data的解释如下：

二、模型使用

1. set-up

import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2

# 用来显示掩膜
def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)

    img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))
    img[:,:,3] = 0
    for ann in sorted_anns:
        m = ann['segmentation']
        color_mask = np.concatenate([np.random.random(3), [0.65]])
        img[m] = color_mask
    ax.imshow(img)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

2. example image导入图片

image = cv2.imread('images/lzu.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
1
2

3. automatic mask generation自动掩膜生成器

import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

sam_checkpoint = "sam_vit_b_01ec64.pth" # 模型
model_type = "vit_b"

device = "cuda" # 使用GPU

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

mask_generator = SamAutomaticMaskGenerator(sam)
1
2
3
4
5
6
7
8
9
10
11
12
13

注：SamAutomaticMaskGenerator生成器有几个可调整的参数，可用来控制采样密度或移除低质量和重复掩膜，以及设置生成器运行在裁剪后更小的对象上以提高性能、后处理移除零散像素和生成的洞等。SamAutomaticMaskGenerator参数设置如下：

# SamAutomaticMaskGenerator 参数输入
mask_generator_2 = SamAutomaticMaskGenerator(
    model=sam,
    points_per_side=32,
    pred_iou_thresh=0.86,
    stability_score_thresh=0.92,
    crop_n_layers=1,
    crop_n_points_downscale_factor=2,
    min_mask_region_area=100,  # Requires open-cv to run post-processing
)
1
2
3
4
5
6
7
8
9
10

4. 调用generate方法，生成掩膜。这里如果使用自定义参数的生成器mask_generator_2 ，则替换为masks = mask_generator_2.generate(image)来生成掩膜即可

masks = mask_generator.generate(image)
1

5. show image图像显示，包括分割图片和掩膜

plt.figure(figsize=(20,20))
plt.imshow(image)
show_anns(masks2)
plt.axis('off')
plt.show() 
1
2
3
4
5

三、结果对比

1. 城市复杂图像分割：目标是建筑物提取

Fig1. 默认参数分割结果（mask_generator生成器）
Fig2. 自定义参数分割结果（mask_generator_2生成器）

2. 类别简单的高分辨率遥感影像分割：目标是进行水域提取

Fig1. 默认参数分割结果（mask_generator生成器）

Fig2. 自定义参数分割结果（mask_generator_2生成器）

3. 结论

可以看出：
（1）对于复杂的城市影像：默认参数存在建筑物漏识别的情况，自定义参数识别的完整性相对较高；
（2）对于内容简单的图像：使用默认参数或自定义参数的结果没有太大差异；
（3）方案一和方案二对比，可见模型复杂图像（如城市）分割的精确度不如简单图像高；
（4）总之，可以使用该模型进行自动分割，再人为对结果进行处理。一方面可以提高图像分割和识别的效率，另一方面可以提高相对于手动数字化的精度；
（5）从SAM官方示例来看，官方示例所使用的数据集也有比较复杂的，但分割也效果却很好。我们本次实验与官方示例的区别是：官方的都是一些拍照的图片，本示例我们使用的高分辨率遥影像。再结合示例一和示例二的结果，我们猜想，这可能与这两种图片的像素问题等原因有关吧，所以分割效果不如官方示例的效果好。后续可以尝试将遥感影像的分割效果与图片的效果做对照。