使用Segment Anything实现图像自动处理结果不符的求助

需求与环境 #

• 自动化操作目标：
  1. 识别图像中的目标对象
  2. 围绕对象裁剪图像，保留少量间隙
  3. 转换为1:1比例，最终导出为800x800px的JPG格式，要求对象居中、背景为白色
• 运行环境：Win11 64位，已完成前置配置：
  1. Python虚拟环境搭建
  2. 安装opencv-python-headless、pillow、numpy及适配CUDA 11.8的PyTorch
  3. 克隆并安装segment-anything仓库
  4. 下载sam_vit_b_01ec64.pth模型文件

现有实现代码 #

import os
import cv2
import numpy as np
from PIL import Image
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator

def load_image(image_path):
    return cv2.imread(image_path)

def save_image(image, path):
    cv2.imwrite(path + '.jpg', image)

def select_object(image):
    sam = sam_model_registry["vit_b"](checkpoint="sam_vit_b_01ec64.pth")
    mask_generator = SamAutomaticMaskGenerator(sam)
    masks = mask_generator.generate(image)
    largest_mask = max(masks, key=lambda x: x['area'])
    return largest_mask['segmentation']

def crop_to_object(image, mask):
    x, y, w, h = cv2.boundingRect(mask.astype(np.uint8))
    padding = 5
    x = max(0, x - padding)
    y = max(0, y - padding)
    w = min(image.shape[1] - x, w + 2 * padding)
    h = min(image.shape[0] - y, h + 2 * padding)
    
    cropped_image = image[y:y+h, x:x+w]
    return cropped_image

def resize_to_square(image, size=800):
    h, w = image.shape[:2]
    scale = size / max(h, w)
    new_h, new_w = int(h * scale), int(w * scale)
    resized_image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)

    new_image = np.ones((size, size, 3), dtype=np.uint8) * 255

    top = (size - new_h) // 2
    left = (size - new_w) // 2
    bottom = top + new_h
    right = left + new_w

    new_image[top:top+new_h, left:left+new_w] = resized_image

    return new_image

def process_image(image_path, output_path):

    image = load_image(image_path)
    mask = select_object(image)
    cropped_image = crop_to_object(image, mask)
    final_image = resize_to_square(cropped_image, 800)
    save_image(final_image, output_path + '.jpg')

def process_folder(input_folder, output_folder):

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    for root, _, files in os.walk(input_folder):
        for filename in files:
            if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
                input_path = os.path.join(root, filename)

                relative_path = os.path.relpath(input_path, input_folder)
                output_path = os.path.join(output_folder, relative_path)
                
                output_dir = os.path.dirname(output_path)
                if not os.path.exists(output_dir):
                    os.makedirs(output_dir)
                
                try:
                    process_image(input_path, output_path)
                    print(f"Processed {input_path}")
                except Exception as e:
                    print(f"Failed to process {input_path}: {e}")

if __name__ == "__main__":
    input_folder = ""
    output_folder = ""
    process_folder(input_folder, output_folder)

问题现象 #

实际处理结果与预期不符，多组输入输出存在明显差异。

排查方向与修复方案 #

1. 对象识别逻辑缺陷 #

当前仅选取面积最大的掩码，但实际场景中最大掩码可能是背景或干扰物，而非目标对象；且未启用CUDA加速，CPU推理会降低掩码质量。

• 修复：
  • 结合stability_score（SAM掩码自带参数，值越高掩码越稳定）和面积筛选最优掩码；
  • 初始化SAM时指定CUDA设备：

    # 全局初始化SAM，避免重复加载
    sam = sam_model_registry["vit_b"](checkpoint="sam_vit_b_01ec64.pth").to("cuda")
    mask_generator = SamAutomaticMaskGenerator(sam)
    
    def select_object(image):
        masks = mask_generator.generate(image)
        # 按面积×稳定性得分排序，取最优掩码
        masks_sorted = sorted(masks, key=lambda x: x['area'] * x['stability_score'], reverse=True)
        return masks_sorted[0]['segmentation']
    

2. 裁剪间隙与比例控制问题 #

固定padding=5导致不同尺寸图像的间隙比例不一致；裁剪后未强制1:1比例，后续补白会导致对象占比不稳定。

• 修复：
  • 改用相对间隙（如取边界框最大边的5%）；
  • 裁剪时直接生成1:1区域，以对象中心为基准扩展：

    def crop_to_square_object(image, mask):
        x, y, w, h = cv2.boundingRect(mask.astype(np.uint8))
        # 计算对象中心
        center_x = x + w // 2
        center_y = y + h // 2
        # 基于最大边设置间隙和边长
        gap = int(max(w, h) * 0.05)
        side_length = max(w, h) + 2 * gap
        # 计算裁剪区域
        x1 = center_x - side_length // 2
        y1 = center_y - side_length // 2
        x2 = x1 + side_length
        y2 = y1 + side_length
        # 处理超出图像边界的情况，提前补白
        cropped = np.ones((side_length, side_length, 3), dtype=np.uint8)*255
        img_x1, img_y1 = max(0, x1), max(0, y1)
        img_x2, img_y2 = min(image.shape[1], x2), min(image.shape[0], y2)
        # 映射到裁剪图的位置
        crop_x1, crop_y1 = img_x1 - x1, img_y1 - y1
        cropped[crop_y1:crop_y1+(img_y2-img_y1), crop_x1:crop_x1+(img_x2-img_x1)] = image[img_y1:img_y2, img_x1:img_x2]
        return cropped
    

3. OpenCV颜色空间不匹配 #

OpenCV默认以BGR格式读取图像，而SAM预期输入为RGB，会导致掩码识别错误；保存时同样以BGR写入，最终图像颜色偏差。

• 修复：

  def load_image(image_path):
      img = cv2.imread(image_path)
      return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  
  def save_image(image, path):
      img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
      cv2.imwrite(path, img_bgr)
  

4. 保存路径重复添加后缀 #

process_image中调用save_image时重复添加.jpg，会生成如xxx.jpg.jpg的错误路径。

• 修复：

  def process_image(image_path, output_path):
      # ... 其他代码 ...
      # 移除原后缀，统一添加.jpg
      output_path = os.path.splitext(output_path)[0] + '.jpg'
      save_image(final_image, output_path)
  

内容的提问来源于stack exchange，提问作者Cyril F.