适用于RKNN2的Segment Anything模型

Segment Anything Model for RKNN2 (English readme see below)

模型使用

• 安装RKNPU2 2.0.0b23版本运行库
• 在开发板上安装python-opencv, rknn-toolkit-lite2, onnxruntime等
• 从 https://huggingface.co/happyme531/segment-anything-rknn2 下载模型文件(sam_vit_b_01ec64.pth.encoder.patched.onnx.rknn,sam_vit_b_01ec64.pth.decoder.onnx)
• 执行run_sam_rknn.py即可

效果展示

输入:

提示：

{\"type\": \"point\", \"data\": [540, 512], \"label\": 1}

输出:

性能: RK3588，单NPU核心，耗时约22000ms
..性能瓶颈: Softmax太大，NPU无法执行

模型转换

• (使用RKNN-Toolkit2 2.0.0b23版本测试)
• 使用https://github.com/vietanhdev/samexporter 导出ONNX模型
• 编辑convert_encoder.py, 修改模型路径:

  ONNX_MODEL=\"sam_vit_b_01ec64.pth.encoder.onnx\"
  

• 执行convert_encoder.py
• 现在会输出一个rknn文件, 但它的执行速度非常慢(~120s), 因为模型结构需要调整
• 执行patch_graph.py, 会生成调整后的onnx文件
• 再次编辑convert_encoder.py, 修改模型路径, 执行转换即可
• decoder模型运行很快，因此无需转换，直接用onnxruntime cpu运行即可

英文说明

模型使用

• 安装RKNPU2 2.0.0b23版本运行库
• 在开发板上安装python-opencv, rknn-toolkit-lite2, onnxruntime等
• 从 https://huggingface.co/happyme531/segment-anything-rknn2 下载模型文件(sam_vit_b_01ec64.pth.encoder.patched.onnx.rknn, sam_vit_b_01ec64.pth.decoder.onnx)
• 执行run_sam_rknn.py

效果展示

输入:

提示:

{\"type\": \"point\", \"data\": [540, 512], \"label\": 1}

输出:

性能: RK3588，单NPU核心，耗时约22000ms
..性能瓶颈: Softmax太大，NPU无法执行

模型转换

• (使用RKNN-Toolkit2 2.0.0b23版本测试)
• 使用 https://github.com/vietanhdev/samexporter 导出ONNX模型
• 编辑 convert_encoder.py，修改模型路径:

  ONNX_MODEL=\"sam_vit_b_01ec64.pth.encoder.onnx\"
  

• 执行 convert_encoder.py
• 现在会输出一个rknn文件，但它的执行速度非常慢(~120s)，因为模型结构需要调整
• 执行 patch_graph.py，会生成调整后的onnx文件
• 再次编辑 convert_encoder.py，修改模型路径，并执行转换
• decoder模型运行很快，因此无需转换，可以直接使用onnxruntime CPU运行。