import cv2
import numpy as np

def simple_background_subtraction(background_path, foreground_path, output_path, threshold=30):
    """
    使用前后两张图做差分，分离前景
    """
    bg = cv2.imread(background_path)
    fg_img = cv2.imread(foreground_path)

    if bg is None or fg_img is None:
        raise FileNotFoundError("图片没找到！检查路径")

    # 确保两图大小一致
    if bg.shape != fg_img.shape:
        print("⚠️ 图片大小不一致，正在调整...")
        fg_img = cv2.resize(fg_img, (bg.shape[1], bg.shape[0]))

    # 转灰度 + 差分
    bg_gray = cv2.cvtColor(bg, cv2.COLOR_BGR2GRAY)
    fg_gray = cv2.cvtColor(fg_img, cv2.COLOR_BGR2GRAY)
    diff = cv2.absdiff(fg_gray, bg_gray)

    # 二值化
    _, mask = cv2.threshold(diff, threshold, 255, cv2.THRESH_BINARY)

    # 形态学去噪
    kernel = np.ones((5,5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

    # 提取前景：背景变白
    result_white = fg_img.copy()
    result_white[mask == 0] = [255, 255, 255]  # 背景白

    # 提取前景：背景透明（PNG）
    b, g, r = cv2.split(fg_img)
    alpha = np.where(mask > 0, 255, 0).astype(np.uint8)
    result_transparent = cv2.merge([b, g, r, alpha])

    # 保存
    cv2.imwrite(output_path + "_white.jpg", result_white)
    cv2.imwrite(output_path + "_transparent.png", result_transparent)

    print(f"✅ 保存成功：\n  {output_path}_white.jpg\n  {output_path}_transparent.png")

    # 显示（可选）
    cv2.imshow('Original', fg_img)
    cv2.imshow('Diff', diff)
    cv2.imshow('Mask', mask)
    cv2.imshow('Foreground', result_white)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


# === 使用 ===
if __name__ == '__main__':
    simple_background_subtraction(
        background_path='f.jpg',
        foreground_path='b.jpg',
        output_path='output/foreground',
        threshold=30
    )