如何用Python实现类似remove.bg的珠宝镂空图片抠图效果？

问题场景 #

使用Python去除戒指、手镯等珠宝产品背景时，外围背景可正常去除，但中间镂空区域的背景无法清除。尝试过rembg及其他工具包均存在此问题，但remove.bg网站能正确完成抠图。

尝试过的代码 #

基于PIL的实现 #

input_image  = Image.open( input_path )
input_array  = np.array( input_image )
output_array = remove( input_array )
output_image = Image.fromarray(output_array)
output_image.save( output_path )

基于OpenCV的实现 #

input_img = cv2.imread( input_path )
output = remove(input_img)
cv2.imwrite( output_path, output )

效果对比 #

• 待处理原图：
  

• 代码处理结果：
  

• remove.bg处理结果：
  

解决方案 #

1. 调用remove.bg官方API #

直接复用remove.bg的专用模型，完美适配珠宝镂空场景：
首先安装依赖：

pip install removebg

示例代码：

from removebg import RemoveBg

# 替换为你的API密钥，可在remove.bg官网获取
api_key = "YOUR_API_KEY"
rbg = RemoveBg(api_key, "error.log")
rbg.remove_background_from_img_file(input_path, output_path=output_path)

2. 基于SAM模型的自定义处理 #

依赖Meta的Segment Anything Model，通过轮廓检测修复镂空区域：

1. 安装依赖：

pip install segment-anything opencv-python numpy pillow

2. 下载SAM预训练模型（如sam_vit_h_4b8939.pth）
3. 示例代码：

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

# 初始化SAM模型
sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"
device = "cuda"  # 无GPU可改为"cpu"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

mask_generator = SamAutomaticMaskGenerator(sam)

# 读取图片
image = cv2.imread(input_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 生成掩码并筛选主体
masks = mask_generator.generate(image_rgb)
target_mask = None
for mask in masks:
    if mask["area"] > 10000 and mask["predicted_iou"] > 0.9:
        target_mask = mask["segmentation"]
        break

if target_mask is not None:
    # 轮廓检测填充镂空区域
    gray = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2GRAY)
    _, thresh = cv2.threshold(gray, 20, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    full_mask = np.zeros_like(target_mask)
    cv2.drawContours(full_mask, contours, -1, 255, thickness=cv2.FILLED)
    
    # 生成RGBA结果图
    result = np.dstack((image_rgb, full_mask.astype(np.uint8)*255))
    Image.fromarray(result).save(output_path)

3. rembg后处理优化 #

基于rembg的基础结果，补充镂空区域修复：

import cv2
import numpy as np
from PIL import Image
from rembg import remove

input_image = Image.open(input_path)
output_array = remove(np.array(input_image))

# 提取Alpha通道并检测轮廓
alpha_channel = output_array[:, :, 3]
gray = cv2.cvtColor(output_array[:, :, :3], cv2.COLOR_RGB2GRAY)
_, thresh = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

# 重新生成Alpha通道，填充镂空区域
new_alpha = np.zeros_like(alpha_channel)
cv2.drawContours(new_alpha, contours, -1, 255, thickness=cv2.FILLED)

output_array[:, :, 3] = new_alpha
Image.fromarray(output_array).save(output_path)

内容的提问来源于stack exchange，提问作者babak-maziar