一、创建环境
bash
# 1. 创建名为 yolov26 的 Conda 环境,指定 Python 3.11
conda create --name yolov26 python=3.11
# 2. 激活该环境(后续安装的包都只在这个环境生效)
conda activate yolov26
# 3. 验证CUDA版本(确保后续装的PyTorch适配显卡)
nvcc --version
# 4. 我电脑装的是rtx4090装的CUDA 12.8的PyTorch(GPU加速必备)
# CUDA 12.8
pip install torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128
其他cuda斑斑自己从官网下载:https://pytorch.org/get-started/previous-versions/
# 5. 安装指定版本的Ultralytics
pip install ultralytics==8.4.0
# 6. 安装OpenCV(处理图像读写和掩码操作)
pip install opencv-python==4.10.0.84
二、测试环境是否安装成功
bash
# 1 验证 PyTorch 是否能正常使用 GPU
python -c"import torch; print(f'CUDA available: {torch.cuda.is_available()}'); print(f'CUDA version: {torch.version.cuda if torch.cuda.is_available() else None}')"
# 2. 验证 Ultralytics 是否安装成功
python -c"from ultralytics import YOLO, SAM; print('Ultralytics 安装成功!')"
三、核心代码解析
python
def main():
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
image_paths = find_images(IMAGE_DIR)
print(f"找到 {len(image_paths)} 张图片:")
for p in image_paths:
print(f" - {p.name}")
# ---- 加载模型 ----
print("\n加载 YOLO26 模型:", YOLO_WEIGHTS.name)
yolo = YOLO(str(YOLO_WEIGHTS))
print("加载 SAM3 模型:", SAM_WEIGHTS.name, "(首次加载较慢, 请耐心等待)")
sam = SAM(str(SAM_WEIGHTS))
# ---- 1. YOLO26 检测 ----
det_results, detections = detect_objects(yolo, image_paths)
# 保存 YOLO 检测结果图
for r, p in zip(det_results, image_paths):
if r.boxes is not None and len(r.boxes) > 0:
det_img = r.plot() # BGR numpy
cv2.imwrite(str(OUTPUT_DIR / f"{p.stem}_det.jpg"), det_img)
# ---- 2. SAM3 分割 + 3. 结果保存 ----
print(f"\n[2/3] SAM3 精准分割 (逐张处理)...")
sam_total_t = time.time()
for idx, (p, bboxes) in enumerate(zip(image_paths, detections), 1):
img = cv2.imread(str(p))
if img is None:
print(f" [{idx}/{len(image_paths)}] {p.name}: 读取失败, 跳过")
continue
if len(bboxes) == 0:
print(f" [{idx}/{len(image_paths)}] {p.name}: YOLO 未检测到目标, 跳过分割")
continue
t0 = time.time()
sam_results = segment_with_sam(sam, p, bboxes)
r = sam_results[0]
if r.masks is None or len(r.masks) == 0:
print(f" [{idx}/{len(image_paths)}] {p.name}: SAM3 未生成掩码")
continue
masks = r.masks.data.cpu().numpy() # (N, H, W), bool/0-1
h, w = img.shape[:2]
# 合并二值掩码 (+ 可选形态学优化)
binary = merge_binary_mask(masks, (h, w))
if REFINE_MASK:
binary = refine_mask(binary, REFINE_KERNEL)
# 彩色叠加图 (SAM 结果里的 boxes 转 numpy)
boxes_np = r.boxes.data.cpu().numpy() if r.boxes is not None else None
overlay = colorize_overlay(img, masks, boxes_np, r.names)
# 保存: 二值掩码 / 叠加图 / 对比图
cv2.imwrite(str(OUTPUT_DIR / f"{p.stem}_mask.png"), binary)
cv2.imwrite(str(OUTPUT_DIR / f"{p.stem}_overlay.jpg"), overlay)
# YOLO 检测图如果没存过, 这里用原图占位
det_img = cv2.imread(str(OUTPUT_DIR / f"{p.stem}_det.jpg"))
if det_img is None:
det_img = img.copy()
cv2.imwrite(str(OUTPUT_DIR / f"{p.stem}_compare.jpg"),
make_comparison(img, det_img, overlay, binary))
print(f" [{idx}/{len(image_paths)}] {p.name}: {len(masks)} 个掩码, "
f"{time.time() - t0:.1f}s -> {p.stem}_mask.png / _overlay.jpg / _compare.jpg")
print(f"\n[3/3] 全部完成! SAM3 总耗时 {time.time() - sam_total_t:.1f}s, 结果保存在: {OUTPUT_DIR}")
四、操作步骤
环境要求
-
Python 3.9 ~ 3.12
-
NVIDIA GPU(建议显存 ≥ 8GB,SAM3 权重 3.4GB)
-
CUDA 版 PyTorch:
> 无需克隆 ultralytics 仓库 ------ 项目根目录已自带源码,脚本会自动优先使用本地版本。
运行
bash
cd E:\python\YOLO26
python YOLO26_SAM3.py
把要分割的图片放进 `pic/` 目录(支持 jpg / jpeg / png / bmp / webp),运行后结果保存在 `output/`:
| 文件 | 内容 |
|---|---|
| `*_det.jpg` | YOLO26 检测结果(边界框 + 类别) |
| `*_mask.png` | 合并后的二值掩码(白=目标,黑=背景) |
| `*_overlay.jpg` | 原图 + 彩色半透明掩码 + 轮廓 + 标签 |
| `*_compare.jpg` | 原图 / 检测 / 分割 / 掩码 四联对比图 |
常用配置(脚本顶部 `配置区`)
-
`CONF_THRESHOLD` --- YOLO 置信度阈值,默认 0.25
-
`DEVICE` --- `""` 自动(优先 GPU),`"0"` 指定显卡,`"cpu"` 纯 CPU
-
`REFINE_MASK` / `REFINE_KERNEL` --- 掩码形态学闭运算去孔洞
-
模型路径、`pic` 目录路径均可修改
流程
-
YOLO26n 检测出目标的边界框(XYXY)
-
每个边界框作为提示(prompt)送入 SAM3,生成像素级掩码
-
所有掩码按位或合并为一张二值掩码,可选形态学优化
五、效果图
六、源码
全部源码已上传,带yolo权重文件
sam3的模型文件大概3g,请自己下载