eardoi/RMBG-1.4-npu
模型介绍
文件和版本
Pull Requests
讨论
分析

RMBG-1.4（昇腾 NPU 适配）

RMBG-1.4 是 BRIA AI 发布的背景移除（background removal / matting）模型， 基于 Segformer 风格 U²-Net 架构（BriaRMBG，44M 参数），输入任意真实图片， 输出前景蒙版并生成透明背景 PNG，广泛用于电商图、证件照、素材处理等场景。

本项目完成其在**华为昇腾 NPU（Ascend 910B）**上的推理适配与验证。

模型规格

项目说明
模型架构BriaRMBG（Segformer + U²-Net，trust_remote_code）
参数量44M
输入尺寸1024×1024（内部缩放）
输出前景 alpha 蒙版 + 透明背景 PNG
权重格式safetensors（model.safetensors）
支持任务背景移除、图像抠图

环境要求

项目要求
硬件Ascend 910B 系列（≥1 卡）
CANN≥ 8.0（source /usr/local/Ascend/ascend-toolkit/set_env.sh）
Python3.10 – 3.11
torch / torch_npu2.9.0 / 2.9.0.post1
scikit-image预处理/后处理需要

环境准备

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple/
pip install torch_npu==2.9.0.post1 transformers pillow safetensors scikit-image \
    -i https://repo.huaweicloud.com/repository/pypi/simple/

模型权重

权重通过 HuggingFace 镜像获取：

git clone --depth 1 https://hf-mirror.com/briaai/RMBG-1.4 model_src

权重目录包含：model.safetensors / briarmbg.py / MyConfig.py / utilities.py / example_input.jpg（示例输入图）等。

推理

cd /workspace/RMBG-1.4
python3 inference.py --model_path ./model_src --image 图片.png --output no_bg.png --device npu:0

参数

参数说明默认值
--model_path模型权重目录./model_src
--image输入图片（真实图片，必选）-
--device推理设备（npu:0 / cpu）npu:0
--output输出透明背景 PNG 路径no_bg.png

设备解析：NPU 优先，自动回退 CPU。

验证结果（Ascend 910B, npu:0）

输入：模型仓库自带真实示例图（example_input.jpg，1405×933）

  • 背景移除完成，输出透明背景 PNG（933×1405）
  • 前景占比：8.4%（主体被正确抠出，背景置为透明）
  • NPU 推理链路完整跑通

适配说明

  • 加载方式：BriaRMBG 的 trust_remote_code 代码基于旧版 PreTrainedModel， 与 transformers 5.15 的 from_pretrained 不兼容（缺 all_tied_weights_keys 属性）。本脚本改为直接实例化模型类 + 手动加载 model.safetensors 权重 （load_state_dict(strict=False)），绕过兼容问题。
  • 预处理/后处理复用仓库自带 utilities.py（preprocess_image / postprocess_image）。
  • 推理使用 torch.no_grad()，输出 alpha 蒙版叠加到原图生成透明背景 PNG。
  • 设备解析 resolve_device()：NPU 可用性检测，失败自动回退 CPU。
  • 模型权重不推送到仓库：*.safetensors / *.bin / *.pth 已加入 .gitignore， 仓库仅包含推理脚本、模型卡片与截图素材。

目录结构

.
├── README.md          # 模型卡片（本文件）
├── inference.py       # NPU 推理脚本
├── utilities.py       # 预处理/后处理工具（来自模型仓库）
└── assets/            # 适配过程截图素材

引用

  • 模型：https://huggingface.co/briaai/RMBG-1.4