RAFT-Stereo on Axera NPU
This repository contains the RAFT-Stereo model converted for high-performance inference on the Axera NPU.
This version is optimized with w8a16 quantization and is compatible with Pulsar2 version 4.2.
Convert Tools Links
For model conversion and deployment guidance:
- AXera Platform GitHub Repo: Sample code and optimization guides for Axera NPU.
- Pulsar2 Documentation: Guide for converting ONNX models to
.axmodel.
Support Platforms
- AX650N/AX8850
- AX630C
- AX637
| Chips | Model Variant | Latency (ms) |
|---|---|---|
| AX650 | raft_steoro256x640_r1.axmodel | 20.9 |
| AX650 | raft_steoro256x640_r4.axmodel | 111.8 |
| Chips | Model Variant | Latency (ms) |
|---|---|---|
| AX630C | raft_steoro256x640_r1_npu2.axmodel | 317.765 |
| AX630C | raft_steoro256x640_r4_npu2.axmodel | 825.793 |
| Chips | Model Variant | Latency (ms) |
|---|---|---|
| AX637 | raft_steoro256x640_r1_npu1.axmodel | 61.033 |
| AX637 | raft_steoro256x640_r4_npu1.axmodel | 99.094 |
How to Use
Download the repository and ensure the directory structure is organized as follows:
├── ax630c
│ ├── raft_steoro256x640_r1.onnx
│ ├── raft_steoro256x640_r1_npu2.axmodel
│ ├── raft_steoro256x640_r4.onnx
│ └── raft_steoro256x640_r4_npu2.axmodel
├── ax650
│ ├── raft_steoro256x640_r1.axmodel
│ ├── raft_steoro256x640_r1.onnx
│ ├── raft_steoro384x1280_r4.axmodel
│ └── raft_steoro384x1280_r4.onnx
├── ax637
│ ├── raft_steoro256x640_r1_npu1.axmodel
│ └── raft_steoro256x640_r4_npu1.axmodel
├── examples
│ ├── left
│ │ ├── 000051_11.png
│ │ ├── 000058_11.png
│ │ ├── 000059_10.png
│ │ ├── 000121_10.png
│ │ ├── 000164_10.png
│ │ ├── 000167_11.png
│ │ ├── 000172_11.png
│ │ ├── 000179_10.png
│ │ ├── 000193_10.png
│ │ └── 000195_10.png
│ └── right
│ ├── 000051_11.png
│ ├── 000058_11.png
│ ├── 000059_10.png
│ ├── 000121_10.png
│ ├── 000164_10.png
│ ├── 000167_11.png
│ ├── 000172_11.png
│ ├── 000179_10.png
│ ├── 000193_10.png
│ └── 000195_10.png
├── infer.py
└── infer_onnx.py
Prerequisites
- Environment: Ensure you have the required Python environment activated with the following core packages installed:
- NPU Runtime:
axengine(PyAXEngine) - Core Libraries:
numpy,opencv-python, andmatplotlib.
- Model/Data: Ensure the compiled
.axmodel, and input stereo pairs are available on the host.
Inference Command
Run the inference script by providing the compiled model, and input data.
python3 infer.py --left examples/left/000051_11.png --right examples/right/000051_11.png --model raft_steoro256x640_r1_npu2.axmodel --width 640 --height 256
Inference with AX630C Host
Results
The model generates a dense disparity map representing the depth of the scene.
(base) root@ax630c:~/data# /python3 infer.py --left examples/left/000051_11.png --right examples/right/000051_11.png --model raft_steoro256x640_r4_npu2.axmodel --width 640 --height 256
[INFO] Available providers: ['AxEngineExecutionProvider']
[INFO] Using provider: AxEngineExecutionProvider
[INFO] Chip type: ChipType.MC20E
[INFO] VNPU type: VNPUType.DISABLED
[INFO] Engine version: 2.7.2a
[INFO] Model type: 1 (full core)
[INFO] Compiler version: 5.1-patch1-dirty e0fdc66b-dirty
Saved: output-ax.png