Spaces:
Running
Running
## Segment Anything Simple Web demo | |
This **front-end only** React based web demo shows how to load a fixed image and corresponding `.npy` file of the SAM image embedding, and run the SAM ONNX model in the browser using Web Assembly with mulithreading enabled by `SharedArrayBuffer`, Web Worker, and SIMD128. | |
<img src="https://github.com/facebookresearch/segment-anything/raw/main/assets/minidemo.gif" width="500"/> | |
## Run the app | |
Install Yarn | |
``` | |
npm install --g yarn | |
``` | |
Build and run: | |
``` | |
yarn && yarn start | |
``` | |
Navigate to [`http://localhost:8081/`](http://localhost:8081/) | |
Move your cursor around to see the mask prediction update in real time. | |
## Export the image embedding | |
In the [ONNX Model Example notebook](https://github.com/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb) upload the image of your choice and generate and save corresponding embedding. | |
Initialize the predictor: | |
```python | |
checkpoint = "sam_vit_h_4b8939.pth" | |
model_type = "vit_h" | |
sam = sam_model_registry[model_type](checkpoint=checkpoint) | |
sam.to(device='cuda') | |
predictor = SamPredictor(sam) | |
``` | |
Set the new image and export the embedding: | |
``` | |
image = cv2.imread('src/assets/dogs.jpg') | |
predictor.set_image(image) | |
image_embedding = predictor.get_image_embedding().cpu().numpy() | |
np.save("dogs_embedding.npy", image_embedding) | |
``` | |
Save the new image and embedding in `src/assets/data`. | |
## Export the ONNX model | |
You also need to export the quantized ONNX model from the [ONNX Model Example notebook](https://github.com/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb). | |
Run the cell in the notebook which saves the `sam_onnx_quantized_example.onnx` file, download it and copy it to the path `/model/sam_onnx_quantized_example.onnx`. | |
Here is a snippet of the export/quantization code: | |
``` | |
onnx_model_path = "sam_onnx_example.onnx" | |
onnx_model_quantized_path = "sam_onnx_quantized_example.onnx" | |
quantize_dynamic( | |
model_input=onnx_model_path, | |
model_output=onnx_model_quantized_path, | |
optimize_model=True, | |
per_channel=False, | |
reduce_range=False, | |
weight_type=QuantType.QUInt8, | |
) | |
``` | |
**NOTE: if you change the ONNX model by using a new checkpoint you need to also re-export the embedding.** | |
## Update the image, embedding, model in the app | |
Update the following file paths at the top of`App.tsx`: | |
```py | |
const IMAGE_PATH = "/assets/data/dogs.jpg"; | |
const IMAGE_EMBEDDING = "/assets/data/dogs_embedding.npy"; | |
const MODEL_DIR = "/model/sam_onnx_quantized_example.onnx"; | |
``` | |
## ONNX multithreading with SharedArrayBuffer | |
To use multithreading, the appropriate headers need to be set to create a cross origin isolation state which will enable use of `SharedArrayBuffer` (see this [blog post](https://cloudblogs.microsoft.com/opensource/2021/09/02/onnx-runtime-web-running-your-machine-learning-model-in-browser/) for more details) | |
The headers below are set in `configs/webpack/dev.js`: | |
```js | |
headers: { | |
"Cross-Origin-Opener-Policy": "same-origin", | |
"Cross-Origin-Embedder-Policy": "credentialless", | |
} | |
``` | |
## Structure of the app | |
**`App.tsx`** | |
- Initializes ONNX model | |
- Loads image embedding and image | |
- Runs the ONNX model based on input prompts | |
**`Stage.tsx`** | |
- Handles mouse move interaction to update the ONNX model prompt | |
**`Tool.tsx`** | |
- Renders the image and the mask prediction | |
**`helpers/maskUtils.tsx`** | |
- Conversion of ONNX model output from array to an HTMLImageElement | |
**`helpers/onnxModelAPI.tsx`** | |
- Formats the inputs for the ONNX model | |
**`helpers/scaleHelper.tsx`** | |
- Handles image scaling logic for SAM (longest size 1024) | |
**`hooks/`** | |
- Handle shared state for the app | |