关于MOBILESAM的部署

本文记录我部署MOBILESAM的历程，也供他人参考
首先获取代码
win+cmd
获取git仓库

git clone https://github.com/ChaoningZhang/MobileSAM.git cd MobileSAM

这个代码就会出现在你的文件夹下面来

关于环境安装就略过了

文件里面是由mobile_sam.pt的在那个/weights/的目录下面
创建一个outputs的文件夹，用来保存图片结果

核心代码是这个
 

import warnings
warnings.filterwarnings("ignore", category=UserWarning)
import torch
import cv2
import numpy as np
import matplotlib.pyplot as plt
from mobile_sam import SamPredictor, sam_model_registry
from segment_anything import SamAutomaticMaskGenerator

# 加载模型，使用 'vit_t'
checkpoint_path = "./weights/mobile_sam.pt"  # 指定预训练模型的路径
sam = sam_model_registry["vit_t"](checkpoint=checkpoint_path)  # 加载 "vit_t" 模型
predictor = SamPredictor(sam)  # 初始化 SamPredictor，用于执行推理任务

# 读取输入图像
image_path = r"" # 图像路径  只需要修改这个就可以了
image = cv2.imread(image_path)  # 使用 OpenCV 读取图像

# 检查图像是否成功读取
if image is None:
    print(f"Error: Could not load image at {image_path}")
    exit()

# 模型期望输入图像为 RGB 格式，因此将 BGR 转为 RGB
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 调整图像大小为更高分辨率，如 1500x1500
image = cv2.resize(image, (1500, 1500))

# 自动生成分割掩码
device = "cpu"
sam.to(device=device)
sam.eval()  # 将模型切换为评估模式
mask_generator = SamAutomaticMaskGenerator(sam)  # 初始化自动分割生成器

# 生成分割掩码
masks = mask_generator.generate(image)
print(f"Number of masks generated: {len(masks)}")
print("Keys for each mask:", masks[0].keys())  # 输出掩码的键值，检查包含哪些信息

# 可视化函数，显示分割区域
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.35]])  # 生成随机颜色
        img[m] = color_mask
    ax.imshow(img)

# 显示原始图像和分割结果
plt.figure(figsize=(20, 20))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show()

# 保存生成的最佳分割结果
best_mask = masks[0]['segmentation']  # 取出第一个分割结果
best_mask_uint8 = (best_mask * 255).astype(np.uint8)  # 将分割掩码转换为 uint8 格式

# 使用形态学操作（闭运算填补空隙，开运算去除噪点）
kernel = np.ones((5, 5), np.uint8)  # 定义一个 5x5 的核
refined_mask = cv2.morphologyEx(best_mask_uint8, cv2.MORPH_CLOSE, kernel)  # 闭运算填补小空隙
refined_mask = cv2.morphologyEx(refined_mask, cv2.MORPH_OPEN, kernel)  # 开运算去除小噪点

# 创建一个彩色图层，并将分割区域填充为蓝色（或其他颜色）
colored_mask = np.zeros_like(image)  # 创建与原图相同大小的空彩色图层
colored_mask[:, :, 2] = refined_mask  # 将分割区域设置为蓝色 (BGR: R=0, G=0, B=255)

# 将遮罩作为透明图层叠加到原始图像上（使用 alpha 混合）
alpha = 0.5  # 透明度系数
overlay_image = cv2.addWeighted(image, 1 - alpha, colored_mask, alpha, 0)
# 注意：此时 overlay_image 仍然是 RGB 颜色空间

# 这里注意一下，因为要看到原来论文的效果
# 但是Matplotlib 使用的是 RGB 颜色空间来显示图像。
# OpenCV 默认使用的是 BGR 颜色空间
# 将结果保存为彩色图像
overlay_image_bgr = cv2.cvtColor(overlay_image, cv2.COLOR_RGB2BGR)  # 将 RGB 转回 BGR 以便保存
cv2.imwrite('./outputs/segmented_output_best_refined.png', overlay_image_bgr)

print("Segmentation with color overlay and refined mask completed.")

时间有点长吖，等待一下就行
这个就是出来的结果

 

原图长这样：
 

保存了是这样，查了资料，这个是由于这个是Matplotlib 和 OpenCV 在处理图像时的颜色空间差异
可以解决来着，不过我没在代码里面写了，有需要的可以自己试试