# -*- coding: utf-8 -*-
"""
产品图片抠图工具
使用rembg库自动去除图片背景
"""
import os
import sys

try:
    from rembg import remove
    from PIL import Image
except ImportError:
    print("Missing dependencies, installing...")
    os.system("pip install rembg pillow")
    from rembg import remove
    from PIL import Image

def remove_background(input_path, output_path):
    """
    去除图片背景
    """
    try:
        print(f"Processing: {input_path}")
        
        # 打开原始图片
        with open(input_path, 'rb') as input_file:
            input_data = input_file.read()
        
        # 去除背景
        output_data = remove(input_data)
        
        # 保存结果
        with open(output_path, 'wb') as output_file:
            output_file.write(output_data)
        
        print(f"Saved: {output_path}")
        return True
    except Exception as e:
        print(f"Error: {e}")
        return False

if __name__ == "__main__":
    # 产品图片路径
    image_dir = r"g:\stonelan-workspace\membank\产品\飞牛 NAS\fyd-亚克力机箱 - 单\image"
    
    # 要处理的图片
    input_image = os.path.join(image_dir, "正面.jpg")
    output_image = os.path.join(image_dir, "正面_抠图.png")
    
    # 检查文件是否存在
    if not os.path.exists(input_image):
        print(f"File not found: {input_image}")
        sys.exit(1)
    
    # 执行抠图
    success = remove_background(input_image, output_image)
    
    if success:
        print(f"\nBackground removal completed!")
        print(f"Original: {input_image}")
        print(f"Result: {output_image}")
    else:
        print("\nBackground removal failed!")
        sys.exit(1)