Linaom1214 / TensorRT-FastSAM

TensorRT-FastSAM(https://github.com/CASIA-IVA-Lab/FastSAM)

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

TensorRT-FastSAM

Requirements

pip install -r requirements.txt

Export ONNX model

python export.py --weights < path to FastSAM-x.pt > --output <onnx_model_path>
##--------- example ---------- 
# python export.py --weights ../models/FastSAM-x.pt --output ../models/FastSAM-x2.onnx

Export trt engine

/usr/src/tensorrt/bin/trtexec \ 
--onnx=<onnx_model_path> \
--saveEngine=<engine_path> \
--explicitBatch \ 
--minShapes=images:1x3x1024x1024 \
--optShapes=images:1x3x1024x1024 \
--maxShapes=images:4x3x1024x1024 \ 
--verbose \ 
--device=0
##--------- example ---------- 
# /usr/src/tensorrt/bin/trtexec \
# --onnx="../models/FastSAM-x2.onnx" \
# --saveEngine="../trt_engines/FastSAM-x.engine" \
# --explicitBatch \
# --minShapes=images:1x3x1024x1024 \
# --optShapes=images:1x3x1024x1024 \
# --maxShapes=images:4x3x1024x1024 \
# --verbose \
# --device=0

Do inference

python trt_infer.py <engine_path> <image_path> --output=<output_path>
##--------- example ---------- 
# python trt_infer.py "../trt_engines/FastSAM-x.engine" --output="../unnamed.jpg"

References

About

TensorRT-FastSAM(https://github.com/CASIA-IVA-Lab/FastSAM)

License:Apache License 2.0


Languages

Language:Python 100.0%