import os
from PIL import Image
import numpy as np

def process_image(image_path, output_path):
    """
    处理单个PNG图像：去除白色背景并根据最外侧的黑边进行裁剪。

    参数:
        image_path (str): 输入PNG图像的路径。
        output_path (str): 保存处理后PNG图像的路径。
    """
    try:
        # 打开图像并确保其为RGBA格式
        img = Image.open(image_path).convert('RGBA')
        
        # 将图像转换为NumPy数组以便高效处理
        data = np.array(img)
        
        # 定义黑色的阈值（例如，RGB值均小于50）
        # 这将创建一个布尔掩码，其中True表示黑色像素
        black_pixels = (data[:, :, 0] < 50) & (data[:, :, 1] < 50) & (data[:, :, 2] < 50)
        
        # 查找所有黑色像素的坐标
        black_coords = np.where(black_pixels)
        
        # 如果没有找到黑色像素，我们将无法裁剪。
        # 此时仅去除白色背景并保存。
        if len(black_coords[0]) == 0:
            print(f"警告: 在 {os.path.basename(image_path)} 中未找到黑色边框。仅去除白色背景。")
            # 为不裁剪设置左上角和右下角坐标
            top, left = 0, 0
            bottom, right = data.shape[0] - 1, data.shape[1] - 1
        else:
            # 确定黑色边框的边界框
            top = np.min(black_coords[0])
            bottom = np.max(black_coords[0])
            left = np.min(black_coords[1])
            right = np.max(black_coords[1])

        # 使用边界框裁剪原始图像。
        # 裁剪区域指定为 (left, upper, right, lower)。
        # right 和 lower 是排他性的，所以我们加1。
        cropped_img = img.crop((left, top, right + 1, bottom + 1))
        
        # 将裁剪后的图像转换为NumPy数组
        cropped_data = np.array(cropped_img)
        
        # 定义白色的阈值（例如，RGB值均大于220）
        white_pixels_mask = (cropped_data[:, :, 0] > 220) & (cropped_data[:, :, 1] > 220) & (cropped_data[:, :, 2] > 220)
        
        # 将白色像素的alpha通道设置为0（完全透明）
        cropped_data[white_pixels_mask, 3] = 0
        
        # 从修改后的数组创建新图像
        final_img = Image.fromarray(cropped_data)
        
        # 保存最终图像
        final_img.save(output_path)
        print(f"成功处理 {os.path.basename(image_path)} 并保存到 {os.path.basename(output_path)}")

    except Exception as e:
        print(f"处理 {os.path.basename(image_path)} 时出错: {e}")

def main():
    """
    主函数，用于查找并处理指定目录中的所有PNG图像。
    """
    # 用户指定的输入目录
    input_dir = r"C:\Users\83500\Desktop\示例文件夹"
    
    # 在输入目录内创建一个'processed'子目录用于存放输出文件
    output_dir = os.path.join(input_dir, "processed")
    
    if not os.path.isdir(input_dir):
        print(f"错误: 输入目录 '{input_dir}' 未找到。")
        return
        
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"已创建输出目录: {output_dir}")
        
    # 遍历输入目录中的所有文件
    file_found = False
    for filename in os.listdir(input_dir):
        # 检查文件是否为PNG图像
        if filename.lower().endswith(".png"):
            file_found = True
            image_path = os.path.join(input_dir, filename)
            output_path = os.path.join(output_dir, filename)
            process_image(image_path, output_path)
    
    if not file_found:
        print(f"在目录 '{input_dir}' 中没有找到PNG文件。")

if __name__ == "__main__":
    main()