实时Image segmentation background subtraction代码转静态图片处理需求

修改说明 #

针对原实时摄像头背景移除代码，做了以下适配调整以支持静态图片处理：

• 移除实时摄像头捕获循环，改为加载指定静态前景图片
• 保留背景图片切换功能，切换后自动重新生成处理结果
• 新增结果保存功能，可将处理后的图片导出到本地
• 调整窗口交互逻辑，适配静态图片的查看与操作需求

适配后代码 #

# 静态图片背景移除实现
import os
import cv2
import numpy as np
import mediapipe as mp

# 初始化MediaPipe人像分割模型
mp_selfie_segmentation = mp.solutions.selfie_segmentation
selfie_segmentation = mp_selfie_segmentation.SelfieSegmentation(model_selection=1)

# 配置路径
foreground_img_path = "input_photo.jpg"  # 替换为你的目标前景图片路径
bg_images_dir = "images"                 # 背景图片存放目录

# 加载背景图片列表
bg_images_list = os.listdir(bg_images_dir)
if not bg_images_list:
    raise ValueError("背景图片目录中没有图片，请检查路径是否正确")

current_bg_idx = 0
bg_image = cv2.imread(os.path.join(bg_images_dir, bg_images_list[current_bg_idx]))

# 加载前景图片
foreground_img = cv2.imread(foreground_img_path)
if foreground_img is None:
    raise ValueError("无法加载前景图片，请检查路径是否正确")
foreground_img = cv2.flip(foreground_img, 1)  # 可选：翻转图片，和原实时逻辑保持一致
height, width, channel = foreground_img.shape

def generate_output_image(foreground, background):
    """根据前景和背景生成移除背景后的图片"""
    # 转换为RGB格式（MediaPipe要求输入为RGB）
    rgb_frame = cv2.cvtColor(foreground, cv2.COLOR_BGR2RGB)
    # 执行人像分割
    results = selfie_segmentation.process(rgb_frame)
    # 生成掩码条件：分割置信度大于0.6的区域保留前景
    condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.6
    # 调整背景图片尺寸匹配前景
    resized_bg = cv2.resize(background, (width, height))
    # 合成结果图片
    output_img = np.where(condition, foreground, resized_bg)
    return output_img

# 首次生成处理结果
output_image = generate_output_image(foreground_img, bg_image)

# 窗口交互逻辑
while True:
    cv2.imshow("Original Foreground", foreground_img)
    cv2.imshow("Processed Result", output_image)
    
    key = cv2.waitKey(0)
    if key == ord('q'):
        # 退出程序
        break
    elif key == ord('d'):
        # 切换下一张背景图片
        current_bg_idx = (current_bg_idx + 1) % len(bg_images_list)
        bg_image = cv2.imread(os.path.join(bg_images_dir, bg_images_list[current_bg_idx]))
        # 重新生成处理结果
        output_image = generate_output_image(foreground_img, bg_image)
    elif key == ord('s'):
        # 保存当前处理结果
        save_path = "output_photo.jpg"
        cv2.imwrite(save_path, output_image)
        print(f"处理结果已保存至: {save_path}")

# 清理所有OpenCV窗口
cv2.destroyAllWindows()

使用说明 #

1. 替换代码中的foreground_img_path为你的目标静态图片路径
2. 确保bg_images_dir目录下存放了至少一张背景图片
3. 运行代码后，通过按键操作：
   • q：关闭窗口并退出程序
   • d：切换到下一张背景图片，自动更新处理结果
   • s：将当前处理后的图片保存到本地

内容的提问来源于stack exchange，提问作者rabbit b