RMBG-1.4

此版本的 RMBG-1.4 已转换为在 Axera NPU 上运行，采用 w8a8 量化。

兼容 Pulsar2 版本：5.1

转换工具链接：

如果您对模型转换感兴趣，可以通过以下方式导出 axmodel：

• AXera 平台仓库，您可以获取详细的指南
• Pulsar2 链接，如何将 ONNX 转换为 axmodel

支持平台

• AX650

  • M4N-Dock(爱芯派Pro)

  • M.2 加速卡

如何使用

将此仓库中的所有文件下载到设备上

(base) root@ax650:/home/ax650/RMBG-1.4# tree
.
├── ax_inference.py
├── axmodel
│   ├── build_config.json
│   └── rmbg1_4_ax650.axmodel
├── ax_rmbg
├── img
│   ├── 3_1920x1080.jpg
│   └── example_input.jpg
├── mask.png
├── README.md
├── result.png
└── rmbg_1_4_sim.onnx

推理

如果您使用 Python 进行推理，需要安装 pyaxengine，可以通过以下方式安装：

wget https://github.com/AXERA-TECH/pyaxengine/releases/download/0.1.3.rc1/axengine-0.1.3-py3-none-any.whl
pip install axengine-0.1.3-py3-none-any.whl

在 AX650 主机上推理，例如 M4N-Dock(爱芯派Pro)

(base) root@ax650:/home/ax650/RMBG-1.4# python3 ax_inference.py -m ./axmodel/rmbg1_4_ax650.axmodel -i ./img/example_input.jpg
Command: ax_inference.py -m ./axmodel/rmbg1_4_ax650.axmodel -i ./img/example_input.jpg
Parameters:
  --model: ./axmodel/rmbg1_4_ax650.axmodel
  --img_path: ./img/example_input.jpg
  --save_path: ./result.png
[INFO] Available providers:  ['AxEngineExecutionProvider']
[INFO] Using provider: AxEngineExecutionProvider
[INFO] Chip type: ChipType.MC50
[INFO] VNPU type: VNPUType.DISABLED
[INFO] Engine version: 2.12.0s
[INFO] Model type: 0 (single core)
[INFO] Compiler version: 5.1 59c683ff
输入名称：input，输入尺寸：[1024, 1024]
推理时间：121.97 ms
推理完成，带透明通道的结果已保存至：./result.png

使用 C++ 版本推理：

(base) root@ax650:/home/ax650/RMBG-1.4# ./ax_rmbg -m ./axmodel/rmbg1_4_ax650.axmodel -i ./img/3_1920x1080.jpg
--------------------------------------
Model file: ./axmodel/rmbg1_4_ax650.axmodel
Input image: ./img/3_1920x1080.jpg
Output image: result.png
Model input size: 1024 x 1024
Repeat count: 1
--------------------------------------
Original image size: 1920 x 1080, channels: 3
Engine creating handle is done.
Engine creating context is done.
Engine get io info is done. 
Inputs:
  [0] name: input, shape: [1, 3, 1024, 1024]
Outputs:
  [0] name: output, shape: [1, 1, 1024, 1024]
Engine alloc io is done. 
Engine push input is done. 
--------------------------------------
post process cost time:23.89 ms 
--------------------------------------
Repeat 1 times, avg time 106.16 ms, max_time 106.16 ms, min_time 106.16 ms
--------------------------------------
Saved result image: result.png
Saved mask: mask.png