PyCharm 跑通 SAM 全流程实战

PyCharm 跑通 SAM 全流程实战

目标：在 Windows 环境下，在 已正确安装显卡驱动和 Miniconda 的前提下，完整跑通 Segment Anything（SAM），并在 Notebook 中看到分割效果，成功导出 overlay 图片。

---

0. 前置条件

以下前置环境本文不再重复说明，直接参考我之前的记录即可：

• NVIDIA 显卡驱动已安装，nvidia-smi 能正常识别显卡
• Miniconda 已安装并可正常创建虚拟环境

参考：

• 显卡驱动安装
  https://blog.csdn.net/m0_58648890/article/details/154016219
• Miniconda 安装
  https://blog.csdn.net/m0_58648890/article/details/157805050

---

1. 创建并激活 Conda 虚拟环境（PyCharm 中创建 sam 解释器）

建议 不要使用 base 环境，单独为 SAM 创建一个虚拟环境，避免依赖冲突。

在 PyCharm 中创建并激活 Conda 虚拟环境后，先确认 Python 版本：

python --version

确保当前使用的是创建的虚拟环境。

如果不想安装 Miniconda，那么直接使用本地 Python 环境同样可以跑通 Segment Anything

---

2. 拉取 segment-anything 源码

进入你要存放代码的目录，然后克隆官方仓库：

git clone https://github.com/facebookresearch/segment-anything.git
cd segment-anything

---

3. 编译 / 安装（editable 模式）

在项目根目录执行以下命令：

pip install -e .

如果看到类似下面的输出：

Successfully installed segment-anything-1.0

说明源码已成功安装。

---

4. 安装运行所需的依赖包

补全运行和可视化所需的依赖：

pip install opencv-python matplotlib numpy pillow tqdm

如果需要在 PyCharm / VSCode 中运行 Notebook，再安装：

pip install jupyter ipykernel

---

5. 下载模型权重（checkpoints）

在项目根目录创建权重目录：

mkdir checkpoints
cd checkpoints

下载 ViT-H 模型权重（效果最好，但显存占用相对较高）：

wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

下载完成后，目录结构应为：

segment-anything/
  checkpoints/
    sam_vit_h_4b8939.pth

如果 wget 不可用，那么可以直接用浏览器打开链接下载，然后手动放到 checkpoints/ 目录

---

6. 在 PyCharm 中选择正确的虚拟环境

在 PyCharm 右下角选择 Python Interpreter：

1. 打开 segment-anything 项目
2. 选择 Python Interpreter / Conda Interpreter
3. 指定刚才创建的 sam 虚拟环境

这一点非常关键：
Notebook / Python 脚本必须运行在 sam 环境中，否则容易出现包缺失或 CUDA 不一致的问题。

---

7. 跑通 Notebook（最直观的验证方式）

项目自带多个 Notebook，这里重点验证两个：

• automatic_mask_generator_example.ipynb（自动生成整图 mask）
• test_sam.ipynb（自定义点提示测试）

---

7.1 运行 automatic_mask_generator_example.ipynb

打开：

notebooks/automatic_mask_generator_example.ipynb

逐个 Cell 运行后，如果能看到类似下图所示的效果（整张图片被多种颜色区域分割），说明：

• 模型权重加载正常
• 依赖安装正常
• 推理流程正常

---

8. 运行点提示分割测试（test_sam.ipynb）

这是我自己写的最小测试，用来验证 "真正跑通 SAM"。

8.1 路径设置（最容易踩坑）

由于 Notebook 位于 notebooks/ 目录，路径需要写成相对路径：

CHECKPOINT = "../checkpoints/sam_vit_h_4b8939.pth"
IMAGE_PATH = "images/dog.jpg"

---

8.2 点提示分割说明

核心参数说明：

• POINT = (500, 375)
  表示在图像中的一个前景点（x 向右，y 向下）
• label = 1
  表示前景点
• multimask_output = True
  会返回多个候选 mask，通常取 score 最大的一个

运行成功后，你将看到：

• 图像中叠加半透明 mask

• 标题中显示 SAM vit_h | score=...

• 控制台输出：

  Saved overlay to: ...\notebooks\images\out_overlay.png

这就表示 SAM 已完整跑通 ✅

---

8.3 Segment Anything 点提示测试代码

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

from segment_anything import sam_model_registry, SamPredictor

# =============================
# 配置区
# =============================
MODEL_TYPE = "vit_h"
CHECKPOINT = "checkpoints/sam_vit_h_4b8939.pth"
IMAGE_PATH = "notebooks/images/dog.jpg"

# 前景点 (x, y)，x 向右，y 向下
POINT = (500, 375)

# 保存结果（可选）
SAVE_PATH = "out_overlay.png"
ALPHA = 0.55

# =============================
# 环境检查
# =============================
print("Torch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())

if not os.path.exists(CHECKPOINT):
    raise FileNotFoundError(f"找不到权重文件: {CHECKPOINT}")

image_bgr = cv2.imread(IMAGE_PATH)
if image_bgr is None:
    raise FileNotFoundError(f"找不到图片: {IMAGE_PATH}")

image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

# =============================
# 加载模型
# =============================
device = "cuda" if torch.cuda.is_available() else "cpu"

sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT)
sam.to(device)

predictor = SamPredictor(sam)
predictor.set_image(image_rgb)

# =============================
# 点提示预测
# =============================
input_point = np.array([[POINT[0], POINT[1]]])
input_label = np.array([1])

masks, scores, _ = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True
)

best_idx = int(np.argmax(scores))
best_mask = masks[best_idx]
best_score = float(scores[best_idx])

print(f"Best mask score: {best_score:.3f}")

# =============================
# 显示结果
# =============================
plt.figure(figsize=(9, 7))
plt.imshow(image_rgb)
plt.imshow(best_mask, alpha=ALPHA)
plt.scatter(POINT[0], POINT[1], c="lime", s=200, marker="*")
plt.title(f"SAM {MODEL_TYPE} | score={best_score:.3f}")
plt.axis("off")
plt.show()

# =============================
# 保存 overlay
# =============================
if SAVE_PATH:
    overlay = image_rgb.copy()
    color = np.array([30, 144, 255], dtype=np.uint8)
    overlay[best_mask] = (
        overlay[best_mask] * (1 - ALPHA) + color * ALPHA
    ).astype(np.uint8)

    cv2.imwrite(SAVE_PATH, cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
    print("Saved overlay to:", SAVE_PATH)

---

9. 常见问题与排查

9.1 权重文件找不到（FileNotFoundError）

• 确认权重是否位于 segment-anything/checkpoints/
• Notebook 中路径是否写成 ../checkpoints/...

---

9.2 图片读取失败（cv2.imread 返回 None）

• 检查 IMAGE_PATH 是否正确
• 建议先使用项目自带的 dog.jpg 测试

---

9.3 显存不足（CUDA out of memory）

vit_h 模型显存占用较高，显存较小时建议换 vit_b：

MODEL_TYPE = "vit_b"
CHECKPOINT = "../checkpoints/sam_vit_b_01ec64.pth"

---

9.4 运行很慢 / 怀疑没用到 GPU

在 Notebook 中确认：

import torch
print(torch.cuda.is_available())

返回 True 才说明当前使用的是 GPU。

---

10. 总结

到这里，Segment Anything 的 完整跑通闭环已经完成：

1. Conda 虚拟环境（sam）

2. 拉取 segment-anything 源码

3. pip install -e . 编译安装

4. 安装必要依赖

5. 下载模型权重到 checkpoints

6. Notebook 成功运行

   • 自动分割
   • 点提示分割
   • 成功导出 overlay 结果

后续就可以在此基础上做进一步工程化封装，例如：

• 批量图片分割
• Box Prompt / 多点 Prompt
• 与目标检测模型联动（检测 → 精分割）