sam2_cpp

About This Project

该项目是SAM2算法的c++实现，包括TensorRT、OnnxRuntime三种硬件平台(推理引擎)，使用SAM2Export导出onnx模型，并对MemoryAttention模块进行了优化。

Demo

image	masks	track_result

以下带有opt标志的代表在模型导出工程SAM2Export基础上，优化MemoryAttention模型结构后导出的onnx模型，具体请查看pr_link.

jetson-orin-nx-16GB	qps	cpu
sam2_track(fp16) - origin	0.98	4%
sam2_track(fp16) - opt	4.25	17%

Usage

Download Project

下载git项目

git clone [email protected]:zz990099/sam2_cpp.git
cd sam2_cpp
git submodule init && git submodule update

Build Enviroment

使用docker构建工作环境

cd sam2_cpp
bash easy_deploy_tool/docker/easy_deploy_startup.sh
# Select `jetson` -> `trt10_u2204`/`trt8_u2204`
bash easy_deploy_tool/docker/into_docker.sh

Compile Codes

在docker容器内，编译工程. 使用 -DENABLE_*宏来启用某种推理框架，可用的有: -DENABLE_TENSORRT=ON、-DENABLE_ORT=ON，可以兼容。

cd /workspace
mdkir build && cd build
cmake .. -DBUILD_TESTING=ON -DENABLE_TENSORRT=ON
make -j

Convert Model

1. 从google driver中下载模型，放到/workspace/models/下

2. 在docker容器内，运行模型转换脚本

cd /workspace
bash tools/cvt_onnx2trt.sh

Run Test Cases

1. 下载测试数据link，放到/workspace/test_data下，文件夹名为golf

2. 运行测试用例，具体测试用例请参考代码。

cd /workspace/build
./bin/test_sam_sam2 --gtest_filter=*correctness
# 限制GLOG输出
GLOG_minloglevel=1 ./bin/test_sam_sam2 --gtest_filter=*speed
# 使用ctest查看可用测试用例
ctest -N
# 直接使用ctest测试全部用例
ctest

References

• sam2
• SAM2Export
• EasyDeployTool