Segment Anything 2 in WebGPU

Segment Anything 2 with WebGPU

Segment Anything 2 with WebGPU delivers a fully browser-based implementation of the Segment Anything model, leveraged through WebGPU. This approach enhances performance by executing all processes client-side, ensuring no data transfer to external servers. Experience the capability online with the live demo.

Project Highlights:

• Efficiency: Utilizes WebGPU to run directly in the browser, minimizing latency and ensuring privacy.
• User Interaction: Enables users to segment images by interacting through points on the image canvas.
• Innovation: Inspiration drawn from extensive prior works such as the ONNX Runtime WebGPU support and ibaiGorodo's SAM 2 inference implementations in Python.

Application Architecture

Constructed using Svelte, this application maintains a clean and straightforward architecture. The project structure divides into decoder and encoder modules:

• Encoder: Responsible for generating image embeddings and features required for mask creation. It also deals with UI interactions such as file drag-and-drop.
• Decoder: Calculates and displays masks on the canvas using encodings and user-defined points.

Supports models up to 1024x1024 pixels for inputs and outputs masks of 256x256 pixels.

Model Management

To address model size limitations for web deployment, models are stored locally in the browser after initial download. Implementation leverages the origin private file system for efficient use of disk space.

Model Optimization for Web

Models are optimized for web use through the following steps:

1. Download the original Segment Anything 2 models.
2. Convert them using samexporter to ONNX format for compatibility with ONNX Runtime.
3. Use onnxruntime tools to further optimize models, reducing their size via .ORT format.

For a more detailed methodology and contributions, read the full write-up and explore collaborations that inspired this project.