"""
Base64 图片抠图脚本
接收输入图片路径和输出图片路径，进行抠图处理
"""
import sys
import os
from rembg import remove
from PIL import Image

def main():
    if len(sys.argv) < 3:
        print("错误: 需要提供输入和输出路径")
        sys.exit(1)
    
    input_path = sys.argv[1]
    output_path = sys.argv[2]
    
    try:
        # 读取输入图像
        input_image = Image.open(input_path)
        original_size = input_image.size
        
        # 确保输入图像是RGB模式
        if input_image.mode != 'RGB':
            input_image = input_image.convert('RGB')
        
        # 使用 rembg 移除背景
        output_image = remove(input_image, post_process_mask=False)
        
        # 确保输出图像尺寸与输入图像一致
        if output_image.size != original_size:
            final_image = Image.new('RGBA', original_size, (0, 0, 0, 0))
            final_image.paste(output_image, (0, 0), output_image if output_image.mode == 'RGBA' else None)
            output_image = final_image
        else:
            if output_image.mode != 'RGBA':
                output_image = output_image.convert('RGBA')
        
        # 保存结果
        output_image.save(output_path, 'PNG')
        print("抠图完成")
        
    except Exception as error:
        print(f"错误: {error}", file=sys.stderr)
        sys.exit(1)

if __name__ == '__main__':
    main()