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
  • M4N-Dock(爱芯派Pro)
  • M.2 Accelerator card
• AX630C
  • 爱芯派2
  • Module-LLM
  • LLM630 Compute Kit
• 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

1. Environment: Ensure you have the required Python environment activated with the following core packages installed:

• NPU Runtime: axengine (PyAXEngine)
• Core Libraries: numpy, opencv-python, and matplotlib.

2. 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

Example Visualization