flybroken/sam2-tensorrt

sam2-tensorrt is a C++ project that takes Meta's SAM2 video object segmentation model and rebuilds the inference pipeline in NVIDIA TensorRT so it can run fast on real video instead of just single images. SAM2 is a model that can find and outline a specific object in a video frame, given a starting box or point, and keep tracking that same object as the video plays. This project focuses purely on making that tracking run in real time on an NVIDIA GPU. On an NVIDIA A10 GPU, the project reports processing a single tracked target at around 30 milliseconds per frame, close to 33 frames per second, using 16 bit floating point precision across all six of the TensorRT engines that make up the pipeline. Optimizations described in the README include reusing GPU memory buffers instead of allocating new ones each frame, running image resizing and color conversion on the GPU, doing bounding box extraction on the GPU as well, and using a Kalman filter to help pick the most likely correct mask when tracking gets difficult, such as when an object is partly hidden or changes shape. The pipeline is split into six separate TensorRT engines covering image encoding, memory attention, mask decoding, and memory encoding, with a memory bank that stores past frames to help the model stay locked onto its target. It supports tracking a single object or multiple objects at once, each with its own memory bank. Using it requires a fairly involved native build: specific versions of CUDA, TensorRT, cuDNN, OpenCV built with CUDA support, ONNX Runtime, Intel TBB, and CMake, plus manually converting provided ONNX model files into TensorRT engine files before running the compiled binary on an image or video file. It is aimed at engineers building real time computer vision or video tracking systems who are comfortable with GPU inference tooling.