import cv2
import numpy as np
import onnxruntime as ort
from pathlib import Path
import os
import time

class BiRefNetMatting:
    def __init__(self, model_path=None, input_size=(1024, 1024)):
        """
        初始化BiRefNet模型
        """
        # 如果未指定模型路径，尝试自动查找
        if model_path is None:
            model_path = self.find_model_file()
        
        self.model_path = model_path
        self.input_size = input_size
        
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"模型文件不存在: {model_path}")
        
        # 获取可用的执行提供程序
        available_providers = ort.get_available_providers()
        print(f"可用的执行提供程序: {available_providers}")
        
        # 创建ONNX Runtime会话（使用可用的提供程序）
        self.session = ort.InferenceSession(
            model_path,
            providers=available_providers  # 自动使用可用的提供程序
        )
        
        # 获取输入输出信息
        self.input_info = self.session.get_inputs()[0]
        self.output_info = self.session.get_outputs()[0]
        self.input_name = self.input_info.name
        self.output_name = self.output_info.name
        
        # 获取模型期望的输入数据类型
        self.input_type = self.input_info.type
        print(f"模型期望的输入类型: {self.input_type}")
        
        print(f"模型加载成功: {model_path}")
        print(f"输入名称: {self.input_name}, 输出名称: {self.output_name}")
        print(f"输入尺寸: {self.input_size}")
    
    def find_model_file(self):
        """尝试自动查找模型文件"""
        possible_names = [
            'birefnet_1024x1024.onnx',
            'mvanet_1024x1024.onnx',
            'birefnet.onnx', 
            'BiRefNet.onnx',
            'birefnet_512x512.onnx',
            '*.onnx'  # 任何onnx文件
        ]
        
        for name in possible_names:
            if name == '*.onnx':
                # 查找所有onnx文件
                onnx_files = list(Path('.').glob('*.onnx'))
                if onnx_files:
                    return str(onnx_files[0])
            elif Path(name).exists():
                return name
        
        # 如果都没找到，返回默认名称
        return 'birefnet_1024x1024.onnx'
    
    def preprocess(self, image):
        """
        图像预处理 - 确保使用float32类型
        """
        # 保存原始尺寸
        self.original_size = image.shape[:2]
        
        # 调整图像尺寸
        resized_image = cv2.resize(image, self.input_size, interpolation=cv2.INTER_LINEAR)
        
        # 确保图像是float32类型
        if resized_image.dtype != np.float32:
            resized_image = resized_image.astype(np.float32)
        
        # 归一化 (假设模型使用ImageNet归一化)
        # 使用float32常数确保结果保持float32
        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
        
        # 除以255.0并转换为float32
        normalized_image = resized_image / 255.0
        normalized_image = (normalized_image - mean) / std
        
        # 调整通道顺序 HWC -> CHW
        input_tensor = normalized_image.transpose(2, 0, 1)
        
        # 添加batch维度
        input_tensor = np.expand_dims(input_tensor, axis=0)
        
        # 确保最终类型是float32
        if input_tensor.dtype != np.float32:
            input_tensor = input_tensor.astype(np.float32)
            
        print(f"输入张量形状: {input_tensor.shape}, 类型: {input_tensor.dtype}")
        
        return input_tensor
    
    def postprocess(self, output, apply_threshold=True):
        """
        后处理输出
        """
        # 移除batch维度
        mask = output[0, 0]
        
        # 调整回原始尺寸
        mask = cv2.resize(mask, (self.original_size[1], self.original_size[0]), 
                         interpolation=cv2.INTER_LINEAR)
        
        # 应用阈值（可选）
        if apply_threshold:
            # 使用自适应阈值或固定阈值
            mask = (mask * 255).astype(np.uint8)
            # 使用OTSU阈值或固定阈值
            if np.max(mask) > np.min(mask):  # 确保有变化
                _, binary_mask = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            else:
                _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
            return binary_mask
        else:
            # 返回软掩码（0-1范围）
            return np.clip(mask, 0, 1)
    
    def predict(self, image, apply_threshold=True):
        """
        预测图像
        """
        # 预处理
        input_tensor = self.preprocess(image)
        
        # 推理
        start_time = time.time()
        outputs = self.session.run([self.output_name], {self.input_name: input_tensor})
        inference_time = time.time() - start_time
        
        print(f"推理时间: {inference_time:.3f}秒")
        
        # 后处理
        mask = self.postprocess(outputs[0], apply_threshold)
        
        return mask

def create_composite_result(image, mask, background_color=(255, 255, 255)):
    """
    创建合成结果
    """
    # 创建带透明通道的图像
    b, g, r = cv2.split(image)
    
    if mask.dtype == np.float32 or mask.dtype == np.float64:
        alpha = mask
    else:
        alpha = mask.astype(np.float32) / 255.0
    
    # 应用alpha通道
    b = (b * alpha + background_color[0] * (1 - alpha)).astype(np.uint8)
    g = (g * alpha + background_color[1] * (1 - alpha)).astype(np.uint8)
    r = (r * alpha + background_color[2] * (1 - alpha)).astype(np.uint8)
    
    # 创建RGBA图像
    rgba = cv2.merge([b, g, r, (alpha * 255).astype(np.uint8)])
    
    # 创建白色背景的合成图像
    white_bg = np.ones_like(image) * 255
    white_bg = (image * alpha[..., None] + white_bg * (1 - alpha[..., None])).astype(np.uint8)
    
    return rgba, white_bg

def save_results(output_dir, base_name, image, mask, rgba, white_bg):
    """
    保存所有结果
    """
    # 保存掩码
    cv2.imwrite(str(output_dir / f"{base_name}_mask.png"), mask)
    
    # 保存透明背景图像
    cv2.imwrite(str(output_dir / f"{base_name}_transparent.png"), rgba)
    
    # 保存白色背景图像
    cv2.imwrite(str(output_dir / f"{base_name}_white_bg.png"), white_bg)
    
    # 保存掩码可视化（增强对比度以便观察）
    mask_visual = cv2.normalize(mask, None, 0, 255, cv2.NORM_MINMAX)
    cv2.imwrite(str(output_dir / f"{base_name}_mask_visual.png"), mask_visual)
    
    print(f"结果已保存到: {output_dir}/{base_name}_*")

def main():
    # 配置参数
    image_path = "888.jpg"
    output_dir = Path("mvanet_results")
    output_dir.mkdir(exist_ok=True)
    
    try:
        # 首先检查模型文件
        print("正在查找模型文件...")
        model_files = list(Path('.').glob('*.onnx'))
        if not model_files:
            print("未找到任何.onnx模型文件！")
            print("请将模型文件放在当前目录")
            return
        
        model_path = str(model_files[0])
        print(f"使用模型: {model_path}")
        
        # 初始化模型
        print("正在加载模型...")
        matting_model = BiRefNetMatting(model_path, input_size=(1024, 1024))
        
        # 加载图像
        print(f"正在加载图像: {image_path}")
        image = cv2.imread(image_path)
        if image is None:
            raise FileNotFoundError(f"无法读取图像: {image_path}")
        
        print(f"图像尺寸: {image.shape}")
        
        # 尝试不同参数进行预测
        print("正在进行图像分割...")
        
        # 方法1: 使用软掩码（不二值化）
        try:
            mask_soft = matting_model.predict(image, apply_threshold=False)
            
            # 保存软掩码结果
            rgba_soft, white_bg_soft = create_composite_result(image, mask_soft)
            save_results(output_dir, "soft", image, 
                        (mask_soft * 255).astype(np.uint8), rgba_soft, white_bg_soft)
            
            print("软掩码处理完成")
        except Exception as e:
            print(f"软掩码处理失败: {e}")
        
        # 方法2: 使用硬掩码（二值化）
        try:
            mask_hard = matting_model.predict(image, apply_threshold=True)
            
            # 保存硬掩码结果
            rgba_hard, white_bg_hard = create_composite_result(image, 
                                                              mask_hard.astype(np.float32) / 255.0)
            save_results(output_dir, "hard", image, mask_hard, rgba_hard, white_bg_hard)
            
            print("硬掩码处理完成")
        except Exception as e:
            print(f"硬掩码处理失败: {e}")
        
        print("\n所有处理完成！")
        print(f"结果已保存到: {output_dir}/")
        print("\n生成的文件:")
        print("  soft_*: 软掩码（有透明度渐变）")
        print("  hard_*: 硬掩码（二值化）")
        
    except FileNotFoundError as e:
        print(f"文件错误: {e}")
        print("请确保:")
        print("1. 模型文件(.onnx)在当前目录")
        print("2. 888.jpg 文件存在")
    except Exception as e:
        print(f"处理过程中出现错误: {e}")
        import traceback
        traceback.print_exc()

def test_with_different_input_sizes():
    """测试不同输入尺寸的预处理"""
    image = cv2.imread("888.jpg")
    
    # 测试不同尺寸
    test_sizes = [(512, 512), (1024, 1024), (640, 480)]
    
    for size in test_sizes:
        print(f"\n测试尺寸: {size}")
        
        # 调整尺寸
        resized = cv2.resize(image, size, interpolation=cv2.INTER_LINEAR)
        
        # 转换为float32
        resized = resized.astype(np.float32)
        
        # 归一化
        normalized = resized / 255.0
        normalized = (normalized - np.array([0.485, 0.456, 0.406], dtype=np.float32)) / \
                     np.array([0.229, 0.224, 0.225], dtype=np.float32)
        
        # 调整通道顺序
        tensor = normalized.transpose(2, 0, 1)
        tensor = np.expand_dims(tensor, axis=0)
        
        print(f"  最终形状: {tensor.shape}, 类型: {tensor.dtype}, 范围: [{tensor.min():.3f}, {tensor.max():.3f}]")

if __name__ == "__main__":
    # 如果需要测试不同尺寸的预处理，取消下面的注释
    # test_with_different_input_sizes()
    
    # 运行主程序
    main()

需要安装环境

pip install opencv-python onnxruntime numpy
pip install torch torchvision
pip install rembg


其实不用装那么多，反正我觉得都要测大模型了，当然装的越多越好了

onnxruntime 要是安装不了得用国内境像了

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple onnxruntime

效果如下