最强抠图模型 RMBG-2.0 ONNX模型本地部署

原创已于 2025-05-21 10:23:43 修改 · 838 阅读

CC 4.0 BY-SA版权

文章标签：

于 2025-05-20 19:46:49 首次发布

模型介绍
RMBG-2.0是BRIA AI推出的最新开源图像背景移除模型，基于先进的AI技术实现高精度的前景与背景分离，达到SOTA水平。RMBG-2.0在性能上超越前代版本，从1.4版本的73.26%准确率大幅提升至2.0版本的90.14%，支持各类图像高精度背景移除，抠图效果十分出色

官方仓库：https://huggingface.co/briaai/RMBG-2.0

如果上面的huggingface仓库打不开，可以试试下面的镜像仓库

https://hf-mirror.com/briaai/RMBG-2.0

不同模型图片处理效果对比

在这里插入图片描述

1.4 模型下载
X-AnyLabeling 项目仓库 RMBG模型

https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.4.3/bria-rmbg-1.4.onnx

2.0 模型下载
上面的官方仓库中可以下载，但是速度可能有点慢，可以在 ModelScope 中下载

ModelScope 社区中模型下载

https://www.modelscope.cn/models/AI-ModelScope/RMBG-2.0/files

在这里插入图片描述

推理脚本
这个脚本程序是由 X-AnyLabeling 项目 rmbg.py 程序修改而来的，修改之后方便本地部署和项目集成

import cv2
import numpy as np
from PIL import Image
import onnxruntime as ort
import time

def show_img(img):
    cv2.imshow("01",img)
    cv2.waitKey(0)

class RMBG2():
    def __init__(self, model_ptah) -> None:

        self.sess_opts = ort.SessionOptions()
        self.sess_opts.log_severity_level = 3

        provider = ["CPUExecutionProvider"]        # CPU 推理
        #provider = ["CUDAExecutionProvider","CPUExecutionProvider"]   # GPU 推理

        self.session = ort.InferenceSession(model_ptah, providers=provider,sess_options=self.sess_opts)

        self.input_name = self.session.get_inputs()[0].name
        self.input_shape = (1024,1024)   # [1024, 1024]
        #print(f"模型输入：{self.input_shape}")

        #self.output_name = self.session.get_outputs()[0].name
        #print(f"模型输出：{self.output_name}")
        #self.output_shape = self.session.get_outputs()[0].shape
        #print(f"模型输出：{self.output_shape}")

    def preprocess(self, image: np.ndarray) -> np.ndarray:
        """
        Preprocess the input image for the model.

        Args:
            image (np.ndarray): Input image as a numpy array.

        Returns:
            np.ndarray: Preprocessed image.
        """
        if len(image.shape) < 3:
            image = np.expand_dims(image, axis=2)
        image = cv2.resize(
            image, self.input_shape, interpolation=cv2.INTER_LINEAR
        )
        image = image.astype(np.float32) / 255.0
        image = (image - 0.5) / 1.0
        image = np.transpose(image, (2, 0, 1))
        return np.expand_dims(image, axis=0)

    def forward(self, blob):
        outs = self.session.run(None, {self.input_name: blob})
        outs = outs[0].squeeze(axis=0)
        return outs

    def postprocess(self, result: np.ndarray, original_size: tuple) -> np.ndarray:
        """
        Postprocess the model output.

        Args:
            result (np.ndarray): Model output.
            original_size (tuple): Original image size (height, width).

        Returns:
            np.ndarray: Postprocessed image as a numpy array.
        """
        result = cv2.resize(
            np.squeeze(result),
            original_size[::-1],
            interpolation=cv2.INTER_LINEAR,
        )
        max_val, min_val = np.max(result), np.min(result)
        result = (result - min_val) / (max_val - min_val)
        return (result * 255).astype(np.uint8)

    def infer(self, image):
        """
        Remove the background from an image and save the result.

        Args:
            image (np.ndarray): Input image as a numpy array.
        """

        blob = self.preprocess(image)
        output = self.forward(blob)
        # 输出为前景掩膜
        result_mask = self.postprocess(output, image.shape[:2])
        #print(result_mask.shape)
        #show_img(result_mask)

        pil_mask = Image.fromarray(result_mask)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 掩膜图像转为pil格式图像
        pil_image = Image.fromarray(image)              # 原图转为pil格式图像
        pil_image = pil_image.convert("RGBA")           # 增加透明图层通道
        pil_mask = pil_mask.convert("L")                # 转换为 L（灰度）格式

        output_image = Image.new("RGBA", pil_image.size, (0, 0, 0, 0))
        output_image.paste(pil_image, (0, 0), pil_mask)
        return output_image

if __name__ == '__main__':
    # bria-rmbg-1.4.onnx (1.4模型)或 model_quantized.onnx (2.0模型)
    model_path = r"bria-rmbg-1.4.onnx"          # 模型路径 
    tast = RMBG2(model_path )

    img_path = r"1.png"                         # 输入图片路径
    img =cv2.imread(img_path)

    start = time.time()

    output_image = tast.infer(img)

    end = time.time()
    total = round((end - start), 2)
    print(f"耗时：{total} s")
    print("完成抠图")

    output_image.save('11.png')                 # 保存结果图像
    # 显示去除背景的结果图像
    output_image.show()

说明：修改模型路径、输入图片路径和图片保存路径即可

1.4 模型 cpu推理不到 1s

2.0 模型 cpu推理约 9 s

原图

1.4 模型抠图效果：
1.4
2.0 模型抠图效果：
2.0