README.md

---
license: apache-2.0

language:
- en

base_model:
- princeton-vl/RAFT-Stereo

pipeline_tag: depth-estimation

tags:
- Axera
- RAFT-Stereo
- NPU
- Stereo-Matching
- Depth-Estimation
- Computer-Vision
---

# RAFT-Stereo on Axera NPU

This repository contains the [RAFT-Stereo](https://arxiv.org/abs/2109.07547) 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](https://github.com/AXERA-TECH/RAFT-Stereo.axera): Sample code and optimization guides for Axera NPU.
* [Pulsar2 Documentation](https://pulsar2-docs.readthedocs.io/en/latest/pulsar2/introduction.html): Guide for converting ONNX models to `.axmodel`.

## Support Platforms

- **AX650N/AX8850**
  - [M4N-Dock(爱芯派Pro)](https://wiki.sipeed.com/hardware/zh/maixIV/m4ndock/m4ndock.html)
  - [M.2 Accelerator card](https://docs.m5stack.com/en/ai_hardware/LLM-8850_Card)
- **AX630C**
  - [爱芯派2](https://axera-pi-2-docs-cn.readthedocs.io/zh-cn/latest/index.html)
  - [Module-LLM](https://docs.m5stack.com/zh_CN/module/Module-LLM)
  - [LLM630 Compute Kit](https://docs.m5stack.com/zh_CN/core/LLM630%20Compute%20Kit)
- **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:

```text
├── 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)](https://github.com/AXERA-TECH/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.

```bash
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

![](output-ax.png)