williamzhangzhe / segment_anything_clone Goto Github PK

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.

Install Yarn

npm install --g yarn

Build and run:

yarn && yarn start

Navigate to http://localhost:8081/

Move your cursor around to see the mask prediction update in real time.

In the ONNX Model Example notebook upload the image of your choice and generate and save corresponding embedding.

Initialize the predictor:

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.

You also need to export the quantized ONNX model from the ONNX Model Example notebook.

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 following file paths at the top ofApp.tsx:

const IMAGE_PATH = "/assets/data/dogs.jpg";
const IMAGE_EMBEDDING = "/assets/data/dogs_embedding.npy";
const MODEL_DIR = "/model/sam_onnx_quantized_example.onnx";

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 for more details)

The headers below are set in configs/webpack/dev.js:

headers: {
    "Cross-Origin-Opener-Policy": "same-origin",
    "Cross-Origin-Embedder-Policy": "credentialless",
}

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