使用 Rembg 和双重处理消除视频边缘的白色边框

from rembg import remove, new_session
from PIL import Image

# 初始化 Rembg 会话
rembg_session_u2net = new_session("u2net")
rembg_session_u2net_human_seg = new_session("u2net_human_seg")

def process_image(input_image_path, output_path):
    """
    使用双重 Rembg 处理流程移除图片背景并消除边缘白边。

    Args:
        input_image_path (str): 输入图片路径。
        output_path (str): 输出图片路径。
    """
    try:
        input_image = Image.open(input_image_path)
    except FileNotFoundError:
        print(f"Error: Input image not found at {input_image_path}")
        return

    # 第一次处理：使用 u2net_human_seg 模型
    first_pass_output_image = remove(
        input_image,
        session=rembg_session_u2net_human_seg
    )

    # 第二次处理：使用 u2net 模型，进行 alpha matting
    second_pass_output_image = remove(first_pass_output_image,
                                        post_process_mask=True,
                                        alpha_matting=True,
                                        alpha_matting_foreground_threshold=240,
                                        alpha_matting_background_threshold=10,
                                        alpha_matting_erode_size=15,
                                        session=rembg_session_u2net)

    try:
        second_pass_output_image.save(output_path)
        print(f"Successfully processed image and saved to {output_path}")
    except Exception as e:
        print(f"Error saving the output image: {e}")

# 示例用法
input_image_path = "input.png"  # 替换为你的输入图片路径
output_image_path = "output.png"  # 替换为你的输出图片路径
process_image(input_image_path, output_image_path)