【Python&语义分割】Segment Anything（SAM）模型交互式分割+掩膜保存（三）

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)


def show_points(coords, labels, ax, marker_size=25):
    pos_points = coords[labels == 1]
    neg_points = coords[labels == 0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='o', s=marker_size, edgecolor='white',
               linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='o', s=marker_size, edgecolor='white',
               linewidth=1.25)

def select_point(event):
    # 鼠标点击事件
    global predictor, img_raster, image  # 声明全局变量
    if event.button == 1:
        x, y = event.xdata, event.ydata
        input_point = np.array([[int(x),int(y)]])
        # 单点 prompt  输入格式为(x, y)和并表示出点所带有的标签1(前景点)或0(背景点)。
        input_label = np.array([1])  # 点所对应的标签
        masks, scores, logit = predictor.predict(
            point_coords=input_point,
            point_labels=input_label,
            multimask_output=False,  # 为False时，它将返回一个掩码
        )
        # print(masks.shape)  # (3, 2160, 3840)波段，高，宽
        for i, (mask, score) in enumerate(zip(masks, scores)):
            # 三个置信度不同的图
            show_mask(mask, plt.gca())
            img_raster[mask] = 1
            show_points(input_point, input_label, plt.gca())
            print("[%s]坐标：%s，%s已分割完成......" % (datetime.now().strftime('%Y-%m-%d %H:%M:%S'), int(x), int(y)))
    if event.button == 3:  # 如果点击了鼠标右键
        print("[%s]正在清除内容......" % datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
        img_raster = np.zeros((image.shape[0], image.shape[1]))  # 覆盖原来的掩膜，全新的0数组
        plt.clf()  # 清楚所有图形
        plt.imshow(image)  # 显示原始图片
        plt.axis('off')  # 不显示坐标系
    plt.draw()  # 更新窗口中的图片

global predictor, img_raster, image  # 声明全局变量
print("【程序准备阶段】")
print("[%s]正在读取图片......" % datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
image = cv2.imread(image_path)  # 读取的图像以NumPy数组的形式存储在变量image中
print("[%s]正在转换图片格式......" % datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 将图像从BGR颜色空间转换为RGB颜色空间，还原图片色彩（图像处理库所认同的格式）
print("[%s]正在初始化模型参数......" % datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
img_raster = np.zeros((image.shape[0], image.shape[1]))
# 创建一个二维数组，用于保存掩膜做栅格转面
fit = plt.figure(figsize=(15, 15))
plt.imshow(image)
sys.path.append("..")  # 将当前路径上一级目录添加到sys.path列表，这里模型使用绝对路径所以这行没啥用
sam_checkpoint = model_path  # 定义模型路径
model_type = model_type  # 定义模型类型
device = device  # "cpu"  or  "cuda"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)  # 定义模型参数
predictor = SamPredictor(sam)  # 调用预测模型
predictor.set_image(image)
# 调用`SamPredictor.set_image`来处理图像以产生一个图像嵌入。`SamPredictor`会记住这个嵌入，并将其用于随后的掩码预测
print("【单点分割阶段】")
print("[%s]正在分割图片......" % datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
fit.canvas.mpl_connect('button_press_event', select_point)
# 窗口点击事件
plt.axis('off')
plt.show()
print("【结果保存阶段】")
print("[%s]正在保存分割数据......" % datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
driver = gdal.GetDriverByName('GTiff')  # 载入数据驱动，用于存储内存中的数组
ds_result_raster = driver.Create(out_path, image.shape[1], image.shape[0], bands=1, eType=gdal.GDT_Float64)
ds_result_raster.GetRasterBand(1).SetNoDataValue(0)  # 将无效值设为0
ds_result_raster.GetRasterBand(1).WriteArray(img_raster)  # 将结果写入数组