Segment Anything (SAM) 抠图后，如何批量处理并保存结果？一个脚本搞定-平芜编程栈

Segment Anything (SAM) 批量处理与自动化保存实战指南

当我们需要处理数百张图片时，单张操作显然效率低下。本文将分享一套完整的Python自动化流程，从批量加载图片、智能分割到结构化保存，帮你把SAM模型的强大能力转化为生产力工具。

1. 环境配置与核心思路

在开始编写批量处理脚本前，需要确保环境配置正确。不同于单张图片处理，批量操作需要特别注意内存管理和文件组织。

# 基础环境配置 import os import cv2 import numpy as np from segment_anything import sam_model_registry, SamPredictor

关键组件选择建议：

模型版本：根据硬件条件选择，vit_b适合CPU环境，vit_l/vit_h适合GPU
内存管理：处理大图时建议分块加载
文件组织：采用/originals,/masks,/cropped三级目录结构

提示：首次运行时SAM会自动下载预训练模型，建议提前下载好放入./models目录

2. 批量处理框架设计

一个健壮的批量处理系统需要考虑以下要素：

输入输出流设计
- 支持常见图片格式（jpg/png/tiff）
- 自动跳过损坏文件
- 保留原始目录结构
处理流程优化
- 预加载模型减少重复初始化
- 智能批处理大小控制
- 多线程/进程加速

def process_folder(input_dir, output_dir): # 创建输出目录 os.makedirs(os.path.join(output_dir, 'masks'), exist_ok=True) os.makedirs(os.path.join(output_dir, 'cropped'), exist_ok=True) # 初始化SAM模型 sam_checkpoint = "sam_vit_b_01ec64.pth" model_type = "vit_b" device = "cuda" if torch.cuda.is_available() else "cpu" sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) sam.to(device=device) predictor = SamPredictor(sam) # 遍历处理每张图片 for img_name in os.listdir(input_dir): if not img_name.lower().endswith(('.png', '.jpg', '.jpeg')): continue img_path = os.path.join(input_dir, img_name) process_single_image(predictor, img_path, output_dir)

3. 核心处理函数实现

单张图片的处理需要封装为独立函数，便于维护和复用。以下是关键步骤的实现：

def process_single_image(predictor, img_path, output_dir): # 加载并预处理图片 image = cv2.imread(img_path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 设置图像并生成自动提示点 predictor.set_image(image) input_point = np.array([[image.shape[1]//2, image.shape[0]//2]]) input_label = np.array([1]) # 生成多mask预测 masks, scores, _ = predictor.predict( point_coords=input_point, point_labels=input_label, multimask_output=True, ) # 保存结果 base_name = os.path.splitext(os.path.basename(img_path))[0] save_results(image, masks, scores, output_dir, base_name)

性能优化技巧：

使用predictor.set_image的image_format参数控制内存占用
对相似图片复用提示点坐标
批量处理时适当降低multimask_output的选项

4. 结果保存与后处理

合理的文件组织能极大提升后续使用效率。我们采用以下结构：

output_dir/ ├── originals/ # 原始图片副本 ├── masks/ # 所有mask图像 │ ├── image1_mask1.png │ ├── image1_mask2.png │ └── ... └── cropped/ # 裁剪后的目标区域 ├── image1_crop1.png ├── image1_crop2.png └── ...

保存函数的实现细节：

def save_results(image, masks, scores, output_dir, base_name): # 保存原始图片副本 cv2.imwrite(f"{output_dir}/originals/{base_name}.jpg", cv2.cvtColor(image, cv2.COLOR_RGB2BGR)) # 保存所有mask和对应裁剪区域 for i, (mask, score) in enumerate(zip(masks, scores)): # 保存mask mask_img = (mask * 255).astype(np.uint8) cv2.imwrite(f"{output_dir}/masks/{base_name}_mask{i+1}.png", mask_img) # 保存裁剪区域 masked_image = cv2.bitwise_and(image, image, mask=mask.astype(np.uint8)) cv2.imwrite(f"{output_dir}/cropped/{base_name}_crop{i+1}.png", cv2.cvtColor(masked_image, cv2.COLOR_RGB2BGR)) # 可选：保存元数据 with open(f"{output_dir}/masks/{base_name}_mask{i+1}.txt", 'w') as f: f.write(f"score: {score:.4f}\n")

5. 高级功能扩展

基础功能实现后，可以考虑添加以下增强特性：

5.1 智能提示点生成

def generate_smart_points(image, num_points=3): # 使用边缘检测或显著性分析生成提示点 gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) edges = cv2.Canny(gray, 100, 200) y, x = np.where(edges > 0) if len(x) > num_points: indices = np.random.choice(len(x), num_points, replace=False) return np.column_stack([x[indices], y[indices]]) return np.array([[image.shape[1]//2, image.shape[0]//2]])

5.2 结果过滤与优化

# 在process_single_image中添加过滤 valid_masks = [m for m, s in zip(masks, scores) if s > 0.7] if not valid_masks: valid_masks = [masks[np.argmax(scores)]]

5.3 并行处理加速

from concurrent.futures import ThreadPoolExecutor def batch_process(input_dir, output_dir, workers=4): with ThreadPoolExecutor(max_workers=workers) as executor: for img_name in os.listdir(input_dir): img_path = os.path.join(input_dir, img_name) executor.submit(process_single_image, predictor, img_path, output_dir)

6. 异常处理与日志记录

健壮的生产环境脚本需要完善的错误处理机制：

def safe_process_image(predictor, img_path, output_dir): try: process_single_image(predictor, img_path, output_dir) logging.info(f"Successfully processed {img_path}") except Exception as e: logging.error(f"Failed to process {img_path}: {str(e)}") # 可选：将失败文件移动到单独目录 os.makedirs(f"{output_dir}/failed", exist_ok=True) os.rename(img_path, f"{output_dir}/failed/{os.path.basename(img_path)}")

常见问题处理清单：