gartajackhats1985's picture
Upload 1830 files
07f408f verified
## 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