fitst commit

This commit is contained in:
2026-01-08 14:35:23 +08:00
commit b57e61f2f7
152 changed files with 7675 additions and 0 deletions

7
.flake8 Executable file
View File

@@ -0,0 +1,7 @@
[flake8]
ignore = W503, E203, E221, C901, C408, E741, C407, B017, F811, C101, EXE001, EXE002
max-line-length = 100
max-complexity = 18
select = B,C,E,F,W,T4,B9
per-file-ignores =
**/__init__.py:F401,F403,E402

42
.gitignore vendored Executable file
View File

@@ -0,0 +1,42 @@
.nfs*
# compilation and distribution
__pycache__
_ext
*.pyc
*.pyd
*.so
*.dll
*.egg-info/
build/
dist/
wheels/
# pytorch/python/numpy formats
*.pth
*.pkl
*.npy
*.ts
model_ts*.txt
# onnx models
*.onnx
# ipython/jupyter notebooks
**/.ipynb_checkpoints/
# Editor temporaries
*.swn
*.swo
*.swp
*~
# editor settings
.idea
.vscode
_darcs
# demo
**/node_modules
yarn.lock
package-lock.json

80
CODE_OF_CONDUCT.md Executable file
View File

@@ -0,0 +1,80 @@
# Code of Conduct
## Our Pledge
In the interest of fostering an open and welcoming environment, we as
contributors and maintainers pledge to make participation in our project and
our community a harassment-free experience for everyone, regardless of age, body
size, disability, ethnicity, sex characteristics, gender identity and expression,
level of experience, education, socio-economic status, nationality, personal
appearance, race, religion, or sexual identity and orientation.
## Our Standards
Examples of behavior that contributes to creating a positive environment
include:
* Using welcoming and inclusive language
* Being respectful of differing viewpoints and experiences
* Gracefully accepting constructive criticism
* Focusing on what is best for the community
* Showing empathy towards other community members
Examples of unacceptable behavior by participants include:
* The use of sexualized language or imagery and unwelcome sexual attention or
advances
* Trolling, insulting/derogatory comments, and personal or political attacks
* Public or private harassment
* Publishing others' private information, such as a physical or electronic
address, without explicit permission
* Other conduct which could reasonably be considered inappropriate in a
professional setting
## Our Responsibilities
Project maintainers are responsible for clarifying the standards of acceptable
behavior and are expected to take appropriate and fair corrective action in
response to any instances of unacceptable behavior.
Project maintainers have the right and responsibility to remove, edit, or
reject comments, commits, code, wiki edits, issues, and other contributions
that are not aligned to this Code of Conduct, or to ban temporarily or
permanently any contributor for other behaviors that they deem inappropriate,
threatening, offensive, or harmful.
## Scope
This Code of Conduct applies within all project spaces, and it also applies when
an individual is representing the project or its community in public spaces.
Examples of representing a project or community include using an official
project e-mail address, posting via an official social media account, or acting
as an appointed representative at an online or offline event. Representation of
a project may be further defined and clarified by project maintainers.
This Code of Conduct also applies outside the project spaces when there is a
reasonable belief that an individual's behavior may have a negative impact on
the project or its community.
## Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be
reported by contacting the project team at <opensource-conduct@fb.com>. All
complaints will be reviewed and investigated and will result in a response that
is deemed necessary and appropriate to the circumstances. The project team is
obligated to maintain confidentiality with regard to the reporter of an incident.
Further details of specific enforcement policies may be posted separately.
Project maintainers who do not follow or enforce the Code of Conduct in good
faith may face temporary or permanent repercussions as determined by other
members of the project's leadership.
## Attribution
This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
[homepage]: https://www.contributor-covenant.org
For answers to common questions about this code of conduct, see
https://www.contributor-covenant.org/faq

31
CONTRIBUTING.md Executable file
View File

@@ -0,0 +1,31 @@
# Contributing to segment-anything
We want to make contributing to this project as easy and transparent as
possible.
## Pull Requests
We actively welcome your pull requests.
1. Fork the repo and create your branch from `main`.
2. If you've added code that should be tested, add tests.
3. If you've changed APIs, update the documentation.
4. Ensure the test suite passes.
5. Make sure your code lints, using the `linter.sh` script in the project's root directory. Linting requires `black==23.*`, `isort==5.12.0`, `flake8`, and `mypy`.
6. If you haven't already, complete the Contributor License Agreement ("CLA").
## Contributor License Agreement ("CLA")
In order to accept your pull request, we need you to submit a CLA. You only need
to do this once to work on any of Facebook's open source projects.
Complete your CLA here: <https://code.facebook.com/cla>
## Issues
We use GitHub issues to track public bugs. Please ensure your description is
clear and has sufficient instructions to be able to reproduce the issue.
Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe
disclosure of security bugs. In those cases, please go through the process
outlined on that page and do not file a public issue.
## License
By contributing to segment-anything, you agree that your contributions will be licensed
under the LICENSE file in the root directory of this source tree.

201
LICENSE Executable file
View File

@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

183
README.md Executable file
View File

@@ -0,0 +1,183 @@
## Latest updates -- SAM 2: Segment Anything in Images and Videos
Please check out our new release on [**Segment Anything Model 2 (SAM 2)**](https://github.com/facebookresearch/segment-anything-2).
* SAM 2 code: https://github.com/facebookresearch/segment-anything-2
* SAM 2 demo: https://sam2.metademolab.com/
* SAM 2 paper: https://arxiv.org/abs/2408.00714
![SAM 2 architecture](https://github.com/facebookresearch/segment-anything-2/blob/main/assets/model_diagram.png?raw=true)
**Segment Anything Model 2 (SAM 2)** is a foundation model towards solving promptable visual segmentation in images and videos. We extend SAM to video by considering images as a video with a single frame. The model design is a simple transformer architecture with streaming memory for real-time video processing. We build a model-in-the-loop data engine, which improves model and data via user interaction, to collect [**our SA-V dataset**](https://ai.meta.com/datasets/segment-anything-video), the largest video segmentation dataset to date. SAM 2 trained on our data provides strong performance across a wide range of tasks and visual domains.
# Segment Anything
**[Meta AI Research, FAIR](https://ai.facebook.com/research/)**
[Alexander Kirillov](https://alexander-kirillov.github.io/), [Eric Mintun](https://ericmintun.github.io/), [Nikhila Ravi](https://nikhilaravi.com/), [Hanzi Mao](https://hanzimao.me/), Chloe Rolland, Laura Gustafson, [Tete Xiao](https://tetexiao.com), [Spencer Whitehead](https://www.spencerwhitehead.com/), Alex Berg, Wan-Yen Lo, [Piotr Dollar](https://pdollar.github.io/), [Ross Girshick](https://www.rossgirshick.info/)
[[`Paper`](https://ai.facebook.com/research/publications/segment-anything/)] [[`Project`](https://segment-anything.com/)] [[`Demo`](https://segment-anything.com/demo)] [[`Dataset`](https://segment-anything.com/dataset/index.html)] [[`Blog`](https://ai.facebook.com/blog/segment-anything-foundation-model-image-segmentation/)] [[`BibTeX`](#citing-segment-anything)]
![SAM design](assets/model_diagram.png?raw=true)
The **Segment Anything Model (SAM)** produces high quality object masks from input prompts such as points or boxes, and it can be used to generate masks for all objects in an image. It has been trained on a [dataset](https://segment-anything.com/dataset/index.html) of 11 million images and 1.1 billion masks, and has strong zero-shot performance on a variety of segmentation tasks.
<p float="left">
<img src="assets/masks1.png?raw=true" width="37.25%" />
<img src="assets/masks2.jpg?raw=true" width="61.5%" />
</p>
## Installation
The code requires `python>=3.8`, as well as `pytorch>=1.7` and `torchvision>=0.8`. Please follow the instructions [here](https://pytorch.org/get-started/locally/) to install both PyTorch and TorchVision dependencies. Installing both PyTorch and TorchVision with CUDA support is strongly recommended.
Install Segment Anything:
```
pip install git+https://github.com/facebookresearch/segment-anything.git
```
or clone the repository locally and install with
```
git clone git@github.com:facebookresearch/segment-anything.git
cd segment-anything; pip install -e .
```
The following optional dependencies are necessary for mask post-processing, saving masks in COCO format, the example notebooks, and exporting the model in ONNX format. `jupyter` is also required to run the example notebooks.
```
pip install opencv-python pycocotools matplotlib onnxruntime onnx
```
## <a name="GettingStarted"></a>Getting Started
First download a [model checkpoint](#model-checkpoints). Then the model can be used in just a few lines to get masks from a given prompt:
```
from segment_anything import SamPredictor, sam_model_registry
sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
predictor = SamPredictor(sam)
predictor.set_image(<your_image>)
masks, _, _ = predictor.predict(<input_prompts>)
```
or generate masks for an entire image:
```
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
mask_generator = SamAutomaticMaskGenerator(sam)
masks = mask_generator.generate(<your_image>)
```
Additionally, masks can be generated for images from the command line:
```
python scripts/amg.py --checkpoint <path/to/checkpoint> --model-type <model_type> --input <image_or_folder> --output <path/to/output>
```
See the examples notebooks on [using SAM with prompts](/notebooks/predictor_example.ipynb) and [automatically generating masks](/notebooks/automatic_mask_generator_example.ipynb) for more details.
<p float="left">
<img src="assets/notebook1.png?raw=true" width="49.1%" />
<img src="assets/notebook2.png?raw=true" width="48.9%" />
</p>
## ONNX Export
SAM's lightweight mask decoder can be exported to ONNX format so that it can be run in any environment that supports ONNX runtime, such as in-browser as showcased in the [demo](https://segment-anything.com/demo). Export the model with
```
python scripts/export_onnx_model.py --checkpoint <path/to/checkpoint> --model-type <model_type> --output <path/to/output>
```
See the [example notebook](https://github.com/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb) for details on how to combine image preprocessing via SAM's backbone with mask prediction using the ONNX model. It is recommended to use the latest stable version of PyTorch for ONNX export.
### Web demo
The `demo/` folder has a simple one page React app which shows how to run mask prediction with the exported ONNX model in a web browser with multithreading. Please see [`demo/README.md`](https://github.com/facebookresearch/segment-anything/blob/main/demo/README.md) for more details.
## <a name="Models"></a>Model Checkpoints
Three model versions of the model are available with different backbone sizes. These models can be instantiated by running
```
from segment_anything import sam_model_registry
sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
```
Click the links below to download the checkpoint for the corresponding model type.
- **`default` or `vit_h`: [ViT-H SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth)**
- `vit_l`: [ViT-L SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth)
- `vit_b`: [ViT-B SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth)
## Dataset
See [here](https://ai.facebook.com/datasets/segment-anything/) for an overview of the datastet. The dataset can be downloaded [here](https://ai.facebook.com/datasets/segment-anything-downloads/). By downloading the datasets you agree that you have read and accepted the terms of the SA-1B Dataset Research License.
We save masks per image as a json file. It can be loaded as a dictionary in python in the below format.
```python
{
"image" : image_info,
"annotations" : [annotation],
}
image_info {
"image_id" : int, # Image id
"width" : int, # Image width
"height" : int, # Image height
"file_name" : str, # Image filename
}
annotation {
"id" : int, # Annotation id
"segmentation" : dict, # Mask saved in COCO RLE format.
"bbox" : [x, y, w, h], # The box around the mask, in XYWH format
"area" : int, # The area in pixels of the mask
"predicted_iou" : float, # The model's own prediction of the mask's quality
"stability_score" : float, # A measure of the mask's quality
"crop_box" : [x, y, w, h], # The crop of the image used to generate the mask, in XYWH format
"point_coords" : [[x, y]], # The point coordinates input to the model to generate the mask
}
```
Image ids can be found in sa_images_ids.txt which can be downloaded using the above [link](https://ai.facebook.com/datasets/segment-anything-downloads/) as well.
To decode a mask in COCO RLE format into binary:
```
from pycocotools import mask as mask_utils
mask = mask_utils.decode(annotation["segmentation"])
```
See [here](https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/mask.py) for more instructions to manipulate masks stored in RLE format.
## License
The model is licensed under the [Apache 2.0 license](LICENSE).
## Contributing
See [contributing](CONTRIBUTING.md) and the [code of conduct](CODE_OF_CONDUCT.md).
## Contributors
The Segment Anything project was made possible with the help of many contributors (alphabetical):
Aaron Adcock, Vaibhav Aggarwal, Morteza Behrooz, Cheng-Yang Fu, Ashley Gabriel, Ahuva Goldstand, Allen Goodman, Sumanth Gurram, Jiabo Hu, Somya Jain, Devansh Kukreja, Robert Kuo, Joshua Lane, Yanghao Li, Lilian Luong, Jitendra Malik, Mallika Malhotra, William Ngan, Omkar Parkhi, Nikhil Raina, Dirk Rowe, Neil Sejoor, Vanessa Stark, Bala Varadarajan, Bram Wasti, Zachary Winstrom
## Citing Segment Anything
If you use SAM or SA-1B in your research, please use the following BibTeX entry.
```
@article{kirillov2023segany,
title={Segment Anything},
author={Kirillov, Alexander and Mintun, Eric and Ravi, Nikhila and Mao, Hanzi and Rolland, Chloe and Gustafson, Laura and Xiao, Tete and Whitehead, Spencer and Berg, Alexander C. and Lo, Wan-Yen and Doll{\'a}r, Piotr and Girshick, Ross},
journal={arXiv:2304.02643},
year={2023}
}
```

BIN
assets/masks1.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 3.5 MiB

BIN
assets/masks2.jpg Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 130 KiB

BIN
assets/minidemo.gif Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.9 MiB

BIN
assets/model_diagram.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 568 KiB

BIN
assets/notebook1.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 854 KiB

BIN
assets/notebook2.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.2 MiB

126
demo/README.md Executable file
View File

@@ -0,0 +1,126 @@
## 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

84
demo/configs/webpack/common.js Executable file
View File

@@ -0,0 +1,84 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
const { resolve } = require("path");
const HtmlWebpackPlugin = require("html-webpack-plugin");
const FriendlyErrorsWebpackPlugin = require("friendly-errors-webpack-plugin");
const CopyPlugin = require("copy-webpack-plugin");
const webpack = require("webpack");
module.exports = {
entry: "./src/index.tsx",
resolve: {
extensions: [".js", ".jsx", ".ts", ".tsx"],
},
output: {
path: resolve(__dirname, "dist"),
},
module: {
rules: [
{
test: /\.mjs$/,
include: /node_modules/,
type: "javascript/auto",
resolve: {
fullySpecified: false,
},
},
{
test: [/\.jsx?$/, /\.tsx?$/],
use: ["ts-loader"],
exclude: /node_modules/,
},
{
test: /\.css$/,
use: ["style-loader", "css-loader"],
},
{
test: /\.(scss|sass)$/,
use: ["style-loader", "css-loader", "postcss-loader"],
},
{
test: /\.(jpe?g|png|gif|svg)$/i,
use: [
"file-loader?hash=sha512&digest=hex&name=img/[contenthash].[ext]",
"image-webpack-loader?bypassOnDebug&optipng.optimizationLevel=7&gifsicle.interlaced=false",
],
},
{
test: /\.(woff|woff2|ttf)$/,
use: {
loader: "url-loader",
},
},
],
},
plugins: [
new CopyPlugin({
patterns: [
{
from: "node_modules/onnxruntime-web/dist/*.wasm",
to: "[name][ext]",
},
{
from: "model",
to: "model",
},
{
from: "src/assets",
to: "assets",
},
],
}),
new HtmlWebpackPlugin({
template: "./src/assets/index.html",
}),
new FriendlyErrorsWebpackPlugin(),
new webpack.ProvidePlugin({
process: "process/browser",
}),
],
};

25
demo/configs/webpack/dev.js Executable file
View File

@@ -0,0 +1,25 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
// development config
const { merge } = require("webpack-merge");
const commonConfig = require("./common");
module.exports = merge(commonConfig, {
mode: "development",
devServer: {
hot: true, // enable HMR on the server
open: true,
// These headers enable the cross origin isolation state
// needed to enable use of SharedArrayBuffer for ONNX
// multithreading.
headers: {
"Cross-Origin-Opener-Policy": "same-origin",
"Cross-Origin-Embedder-Policy": "credentialless",
},
},
devtool: "cheap-module-source-map",
});

22
demo/configs/webpack/prod.js Executable file
View File

@@ -0,0 +1,22 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
// production config
const { merge } = require("webpack-merge");
const { resolve } = require("path");
const Dotenv = require("dotenv-webpack");
const commonConfig = require("./common");
module.exports = merge(commonConfig, {
mode: "production",
output: {
filename: "js/bundle.[contenthash].min.js",
path: resolve(__dirname, "../../dist"),
publicPath: "/",
},
devtool: "source-map",
plugins: [new Dotenv()],
});

65
demo/package.json Executable file
View File

@@ -0,0 +1,65 @@
{
"engines": {
"node": ">=20.0.0"
},
"name": "segment-anything-mini-demo",
"version": "0.1.0",
"license": "MIT",
"scripts": {
"build": "yarn run clean-dist && webpack --config=configs/webpack/prod.js && mv dist/*.wasm dist/js",
"clean-dist": "rimraf dist/*",
"lint": "eslint './src/**/*.{js,ts,tsx}' --quiet",
"start": "yarn run start-dev",
"test": "yarn run start-model-test",
"start-dev": "webpack serve --config=configs/webpack/dev.js"
},
"devDependencies": {
"@babel/core": "^7.18.13",
"@babel/preset-env": "^7.18.10",
"@babel/preset-react": "^7.18.6",
"@babel/preset-typescript": "^7.18.6",
"@pmmmwh/react-refresh-webpack-plugin": "^0.5.7",
"@testing-library/react": "^13.3.0",
"@types/node": "^18.7.13",
"@types/react": "^18.0.17",
"@types/react-dom": "^18.0.6",
"@types/underscore": "^1.11.4",
"@typescript-eslint/eslint-plugin": "^5.35.1",
"@typescript-eslint/parser": "^5.35.1",
"babel-loader": "^8.2.5",
"copy-webpack-plugin": "^11.0.0",
"css-loader": "^6.7.1",
"dotenv": "^16.0.2",
"dotenv-webpack": "^8.0.1",
"eslint": "^8.22.0",
"eslint-plugin-react": "^7.31.0",
"file-loader": "^6.2.0",
"fork-ts-checker-webpack-plugin": "^7.2.13",
"friendly-errors-webpack-plugin": "^1.7.0",
"html-webpack-plugin": "^5.5.0",
"image-webpack-loader": "^8.1.0",
"postcss-loader": "^7.0.1",
"postcss-preset-env": "^7.8.0",
"process": "^0.11.10",
"rimraf": "^3.0.2",
"sass": "^1.54.5",
"sass-loader": "^13.0.2",
"style-loader": "^3.3.1",
"tailwindcss": "^3.1.8",
"ts-loader": "^9.3.1",
"typescript": "^4.8.2",
"webpack": "^5.74.0",
"webpack-cli": "^4.10.0",
"webpack-dev-server": "^4.10.0",
"webpack-dotenv-plugin": "^2.1.0",
"webpack-merge": "^5.8.0"
},
"dependencies": {
"npyjs": "^0.4.0",
"onnxruntime-web": "^1.14.0",
"react": "^18.2.0",
"react-dom": "^18.2.0",
"underscore": "^1.13.6",
"react-refresh": "^0.14.0"
}
}

10
demo/postcss.config.js Executable file
View File

@@ -0,0 +1,10 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
const tailwindcss = require("tailwindcss");
module.exports = {
plugins: ["postcss-preset-env", 'tailwindcss/nesting', tailwindcss],
};

130
demo/src/App.tsx Executable file
View File

@@ -0,0 +1,130 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
import { InferenceSession, Tensor } from "onnxruntime-web";
import React, { useContext, useEffect, useState } from "react";
import "./assets/scss/App.scss";
import { handleImageScale } from "./components/helpers/scaleHelper";
import { modelScaleProps } from "./components/helpers/Interfaces";
import { onnxMaskToImage } from "./components/helpers/maskUtils";
import { modelData } from "./components/helpers/onnxModelAPI";
import Stage from "./components/Stage";
import AppContext from "./components/hooks/createContext";
const ort = require("onnxruntime-web");
/* @ts-ignore */
import npyjs from "npyjs";
// Define image, embedding and model paths
const IMAGE_PATH = "/assets/data/dogs.jpg";
const IMAGE_EMBEDDING = "/assets/data/dogs_embedding.npy";
const MODEL_DIR = "/model/sam_onnx_quantized_example.onnx";
const App = () => {
const {
clicks: [clicks],
image: [, setImage],
maskImg: [, setMaskImg],
} = useContext(AppContext)!;
const [model, setModel] = useState<InferenceSession | null>(null); // ONNX model
const [tensor, setTensor] = useState<Tensor | null>(null); // Image embedding tensor
// The ONNX model expects the input to be rescaled to 1024.
// The modelScale state variable keeps track of the scale values.
const [modelScale, setModelScale] = useState<modelScaleProps | null>(null);
// Initialize the ONNX model. load the image, and load the SAM
// pre-computed image embedding
useEffect(() => {
// Initialize the ONNX model
const initModel = async () => {
try {
if (MODEL_DIR === undefined) return;
const URL: string = MODEL_DIR;
const model = await InferenceSession.create(URL);
setModel(model);
} catch (e) {
console.log(e);
}
};
initModel();
// Load the image
const url = new URL(IMAGE_PATH, location.origin);
loadImage(url);
// Load the Segment Anything pre-computed embedding
Promise.resolve(loadNpyTensor(IMAGE_EMBEDDING, "float32")).then(
(embedding) => setTensor(embedding)
);
}, []);
const loadImage = async (url: URL) => {
try {
const img = new Image();
img.src = url.href;
img.onload = () => {
const { height, width, samScale } = handleImageScale(img);
setModelScale({
height: height, // original image height
width: width, // original image width
samScale: samScale, // scaling factor for image which has been resized to longest side 1024
});
img.width = width;
img.height = height;
setImage(img);
};
} catch (error) {
console.log(error);
}
};
// Decode a Numpy file into a tensor.
const loadNpyTensor = async (tensorFile: string, dType: string) => {
let npLoader = new npyjs();
const npArray = await npLoader.load(tensorFile);
const tensor = new ort.Tensor(dType, npArray.data, npArray.shape);
return tensor;
};
// Run the ONNX model every time clicks has changed
useEffect(() => {
runONNX();
}, [clicks]);
const runONNX = async () => {
try {
if (
model === null ||
clicks === null ||
tensor === null ||
modelScale === null
)
return;
else {
// Preapre the model input in the correct format for SAM.
// The modelData function is from onnxModelAPI.tsx.
const feeds = modelData({
clicks,
tensor,
modelScale,
});
if (feeds === undefined) return;
// Run the SAM ONNX model with the feeds returned from modelData()
const results = await model.run(feeds);
const output = results[model.outputNames[0]];
// The predicted mask returned from the ONNX model is an array which is
// rendered as an HTML image using onnxMaskToImage() from maskUtils.tsx.
setMaskImg(onnxMaskToImage(output.data, output.dims[2], output.dims[3]));
}
} catch (e) {
console.log(e);
}
};
return <Stage />;
};
export default App;

BIN
demo/src/assets/data/dogs.jpg Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 438 KiB

18
demo/src/assets/index.html Executable file
View File

@@ -0,0 +1,18 @@
<!DOCTYPE html>
<html lang="en" dir="ltr" prefix="og: https://ogp.me/ns#" class="w-full h-full">
<head>
<meta charset="utf-8" />
<meta
name="viewport"
content="width=device-width, initial-scale=1, shrink-to-fit=no"
/>
<title>Segment Anything Demo</title>
<!-- Meta Tags -->
<meta property="og:type" content="website" />
<meta property="og:title" content="Segment Anything Demo" />
</head>
<body class="w-full h-full">
<div id="root" class="w-full h-full"></div>
</body>
</html>

3
demo/src/assets/scss/App.scss Executable file
View File

@@ -0,0 +1,3 @@
@tailwind base;
@tailwind components;
@tailwind utilities;

49
demo/src/components/Stage.tsx Executable file
View File

@@ -0,0 +1,49 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
import React, { useContext } from "react";
import * as _ from "underscore";
import Tool from "./Tool";
import { modelInputProps } from "./helpers/Interfaces";
import AppContext from "./hooks/createContext";
const Stage = () => {
const {
clicks: [, setClicks],
image: [image],
} = useContext(AppContext)!;
const getClick = (x: number, y: number): modelInputProps => {
const clickType = 1;
return { x, y, clickType };
};
// Get mouse position and scale the (x, y) coordinates back to the natural
// scale of the image. Update the state of clicks with setClicks to trigger
// the ONNX model to run and generate a new mask via a useEffect in App.tsx
const handleMouseMove = _.throttle((e: any) => {
let el = e.nativeEvent.target;
const rect = el.getBoundingClientRect();
let x = e.clientX - rect.left;
let y = e.clientY - rect.top;
const imageScale = image ? image.width / el.offsetWidth : 1;
x *= imageScale;
y *= imageScale;
const click = getClick(x, y);
if (click) setClicks([click]);
}, 15);
const flexCenterClasses = "flex items-center justify-center";
return (
<div className={`${flexCenterClasses} w-full h-full`}>
<div className={`${flexCenterClasses} relative w-[90%] h-[90%]`}>
<Tool handleMouseMove={handleMouseMove} />
</div>
</div>
);
};
export default Stage;

73
demo/src/components/Tool.tsx Executable file
View File

@@ -0,0 +1,73 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
import React, { useContext, useEffect, useState } from "react";
import AppContext from "./hooks/createContext";
import { ToolProps } from "./helpers/Interfaces";
import * as _ from "underscore";
const Tool = ({ handleMouseMove }: ToolProps) => {
const {
image: [image],
maskImg: [maskImg, setMaskImg],
} = useContext(AppContext)!;
// Determine if we should shrink or grow the images to match the
// width or the height of the page and setup a ResizeObserver to
// monitor changes in the size of the page
const [shouldFitToWidth, setShouldFitToWidth] = useState(true);
const bodyEl = document.body;
const fitToPage = () => {
if (!image) return;
const imageAspectRatio = image.width / image.height;
const screenAspectRatio = window.innerWidth / window.innerHeight;
setShouldFitToWidth(imageAspectRatio > screenAspectRatio);
};
const resizeObserver = new ResizeObserver((entries) => {
for (const entry of entries) {
if (entry.target === bodyEl) {
fitToPage();
}
}
});
useEffect(() => {
fitToPage();
resizeObserver.observe(bodyEl);
return () => {
resizeObserver.unobserve(bodyEl);
};
}, [image]);
const imageClasses = "";
const maskImageClasses = `absolute opacity-40 pointer-events-none`;
// Render the image and the predicted mask image on top
return (
<>
{image && (
<img
onMouseMove={handleMouseMove}
onMouseOut={() => _.defer(() => setMaskImg(null))}
onTouchStart={handleMouseMove}
src={image.src}
className={`${
shouldFitToWidth ? "w-full" : "h-full"
} ${imageClasses}`}
></img>
)}
{maskImg && (
<img
src={maskImg.src}
className={`${
shouldFitToWidth ? "w-full" : "h-full"
} ${maskImageClasses}`}
></img>
)}
</>
);
};
export default Tool;

View File

@@ -0,0 +1,29 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
import { Tensor } from "onnxruntime-web";
export interface modelScaleProps {
samScale: number;
height: number;
width: number;
}
export interface modelInputProps {
x: number;
y: number;
clickType: number;
}
export interface modeDataProps {
clicks?: Array<modelInputProps>;
tensor: Tensor;
modelScale: modelScaleProps;
}
export interface ToolProps {
handleMouseMove: (e: any) => void;
}

View File

@@ -0,0 +1,47 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
// Convert the onnx model mask prediction to ImageData
function arrayToImageData(input: any, width: number, height: number) {
const [r, g, b, a] = [0, 114, 189, 255]; // the masks's blue color
const arr = new Uint8ClampedArray(4 * width * height).fill(0);
for (let i = 0; i < input.length; i++) {
// Threshold the onnx model mask prediction at 0.0
// This is equivalent to thresholding the mask using predictor.model.mask_threshold
// in python
if (input[i] > 0.0) {
arr[4 * i + 0] = r;
arr[4 * i + 1] = g;
arr[4 * i + 2] = b;
arr[4 * i + 3] = a;
}
}
return new ImageData(arr, height, width);
}
// Use a Canvas element to produce an image from ImageData
function imageDataToImage(imageData: ImageData) {
const canvas = imageDataToCanvas(imageData);
const image = new Image();
image.src = canvas.toDataURL();
return image;
}
// Canvas elements can be created from ImageData
function imageDataToCanvas(imageData: ImageData) {
const canvas = document.createElement("canvas");
const ctx = canvas.getContext("2d");
canvas.width = imageData.width;
canvas.height = imageData.height;
ctx?.putImageData(imageData, 0, 0);
return canvas;
}
// Convert the onnx model mask output to an HTMLImageElement
export function onnxMaskToImage(input: any, width: number, height: number) {
return imageDataToImage(arrayToImageData(input, width, height));
}

View File

@@ -0,0 +1,71 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
import { Tensor } from "onnxruntime-web";
import { modeDataProps } from "./Interfaces";
const modelData = ({ clicks, tensor, modelScale }: modeDataProps) => {
const imageEmbedding = tensor;
let pointCoords;
let pointLabels;
let pointCoordsTensor;
let pointLabelsTensor;
// Check there are input click prompts
if (clicks) {
let n = clicks.length;
// If there is no box input, a single padding point with
// label -1 and coordinates (0.0, 0.0) should be concatenated
// so initialize the array to support (n + 1) points.
pointCoords = new Float32Array(2 * (n + 1));
pointLabels = new Float32Array(n + 1);
// Add clicks and scale to what SAM expects
for (let i = 0; i < n; i++) {
pointCoords[2 * i] = clicks[i].x * modelScale.samScale;
pointCoords[2 * i + 1] = clicks[i].y * modelScale.samScale;
pointLabels[i] = clicks[i].clickType;
}
// Add in the extra point/label when only clicks and no box
// The extra point is at (0, 0) with label -1
pointCoords[2 * n] = 0.0;
pointCoords[2 * n + 1] = 0.0;
pointLabels[n] = -1.0;
// Create the tensor
pointCoordsTensor = new Tensor("float32", pointCoords, [1, n + 1, 2]);
pointLabelsTensor = new Tensor("float32", pointLabels, [1, n + 1]);
}
const imageSizeTensor = new Tensor("float32", [
modelScale.height,
modelScale.width,
]);
if (pointCoordsTensor === undefined || pointLabelsTensor === undefined)
return;
// There is no previous mask, so default to an empty tensor
const maskInput = new Tensor(
"float32",
new Float32Array(256 * 256),
[1, 1, 256, 256]
);
// There is no previous mask, so default to 0
const hasMaskInput = new Tensor("float32", [0]);
return {
image_embeddings: imageEmbedding,
point_coords: pointCoordsTensor,
point_labels: pointLabelsTensor,
orig_im_size: imageSizeTensor,
mask_input: maskInput,
has_mask_input: hasMaskInput,
};
};
export { modelData };

View File

@@ -0,0 +1,18 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
// Helper function for handling image scaling needed for SAM
const handleImageScale = (image: HTMLImageElement) => {
// Input images to SAM must be resized so the longest side is 1024
const LONG_SIDE_LENGTH = 1024;
let w = image.naturalWidth;
let h = image.naturalHeight;
const samScale = LONG_SIDE_LENGTH / Math.max(h, w);
return { height: h, width: w, samScale };
};
export { handleImageScale };

View File

@@ -0,0 +1,31 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
import React, { useState } from "react";
import { modelInputProps } from "../helpers/Interfaces";
import AppContext from "./createContext";
const AppContextProvider = (props: {
children: React.ReactElement<any, string | React.JSXElementConstructor<any>>;
}) => {
const [clicks, setClicks] = useState<Array<modelInputProps> | null>(null);
const [image, setImage] = useState<HTMLImageElement | null>(null);
const [maskImg, setMaskImg] = useState<HTMLImageElement | null>(null);
return (
<AppContext.Provider
value={{
clicks: [clicks, setClicks],
image: [image, setImage],
maskImg: [maskImg, setMaskImg],
}}
>
{props.children}
</AppContext.Provider>
);
};
export default AppContextProvider;

View File

@@ -0,0 +1,27 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
import { createContext } from "react";
import { modelInputProps } from "../helpers/Interfaces";
interface contextProps {
clicks: [
clicks: modelInputProps[] | null,
setClicks: (e: modelInputProps[] | null) => void
];
image: [
image: HTMLImageElement | null,
setImage: (e: HTMLImageElement | null) => void
];
maskImg: [
maskImg: HTMLImageElement | null,
setMaskImg: (e: HTMLImageElement | null) => void
];
}
const AppContext = createContext<contextProps | null>(null);
export default AppContext;

17
demo/src/index.tsx Executable file
View File

@@ -0,0 +1,17 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
import * as React from "react";
import { createRoot } from "react-dom/client";
import AppContextProvider from "./components/hooks/context";
import App from "./App";
const container = document.getElementById("root");
const root = createRoot(container!);
root.render(
<AppContextProvider>
<App/>
</AppContextProvider>
);

12
demo/tailwind.config.js Executable file
View File

@@ -0,0 +1,12 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
// This source code is licensed under the license found in the
// LICENSE file in the root directory of this source tree.
/** @type {import('tailwindcss').Config} */
module.exports = {
content: ["./src/**/*.{html,js,tsx}"],
theme: {},
plugins: [],
};

24
demo/tsconfig.json Executable file
View File

@@ -0,0 +1,24 @@
{
"compilerOptions": {
"lib": ["dom", "dom.iterable", "esnext"],
"allowJs": true,
"skipLibCheck": true,
"strict": true,
"forceConsistentCasingInFileNames": true,
"noEmit": false,
"esModuleInterop": true,
"module": "esnext",
"moduleResolution": "node",
"resolveJsonModule": true,
"isolatedModules": true,
"jsx": "react",
"incremental": true,
"target": "ESNext",
"useDefineForClassFields": true,
"allowSyntheticDefaultImports": true,
"outDir": "./dist/",
"sourceMap": true
},
"include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", "src"],
"exclude": ["node_modules"]
}

32
linter.sh Executable file
View File

@@ -0,0 +1,32 @@
#!/bin/bash -e
# Copyright (c) Facebook, Inc. and its affiliates.
{
black --version | grep -E "23\." > /dev/null
} || {
echo "Linter requires 'black==23.*' !"
exit 1
}
ISORT_VERSION=$(isort --version-number)
if [[ "$ISORT_VERSION" != 5.12* ]]; then
echo "Linter requires isort==5.12.0 !"
exit 1
fi
echo "Running isort ..."
isort . --atomic
echo "Running black ..."
black -l 100 .
echo "Running flake8 ..."
if [ -x "$(command -v flake8)" ]; then
flake8 .
else
python3 -m flake8 .
fi
echo "Running mypy..."
mypy --exclude 'setup.py|notebooks' .

File diff suppressed because one or more lines are too long

BIN
notebooks/images/dog.jpg Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 98 KiB

BIN
notebooks/images/groceries.jpg Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 164 KiB

BIN
notebooks/images/truck.jpg Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 265 KiB

View File

@@ -0,0 +1,832 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "901c8ef3",
"metadata": {},
"outputs": [],
"source": [
"# Copyright (c) Meta Platforms, Inc. and affiliates."
]
},
{
"cell_type": "markdown",
"id": "1662bb7c",
"metadata": {},
"source": [
"# Produces masks from prompts using an ONNX model"
]
},
{
"cell_type": "markdown",
"id": "7fcc21a0",
"metadata": {},
"source": [
"SAM's prompt encoder and mask decoder are very lightweight, which allows for efficient computation of a mask given user input. This notebook shows an example of how to export and use this lightweight component of the model in ONNX format, allowing it to run on a variety of platforms that support an ONNX runtime."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "86daff77",
"metadata": {
"ExecuteTime": {
"end_time": "2025-08-06T03:16:23.533543Z",
"start_time": "2025-08-06T03:16:23.528904Z"
}
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb\">\n",
" <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
"</a>\n"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"from IPython.display import display, HTML\n",
"display(HTML(\n",
"\"\"\"\n",
"<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb\">\n",
" <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
"</a>\n",
"\"\"\"\n",
"))"
]
},
{
"cell_type": "markdown",
"id": "55ae4e00",
"metadata": {},
"source": [
"## Environment Set-up"
]
},
{
"cell_type": "markdown",
"id": "109a5cc2",
"metadata": {},
"source": [
"If running locally using jupyter, first install `segment_anything` in your environment using the [installation instructions](https://github.com/facebookresearch/segment-anything#installation) in the repository. The latest stable versions of PyTorch and ONNX are recommended for this notebook. If running from Google Colab, set `using_colab=True` below and run the cell. In Colab, be sure to select 'GPU' under 'Edit'->'Notebook Settings'->'Hardware accelerator'."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "39b99fc4",
"metadata": {
"ExecuteTime": {
"end_time": "2025-08-06T03:17:30.823461Z",
"start_time": "2025-08-06T03:17:30.818737Z"
}
},
"outputs": [],
"source": [
"using_colab = False"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "296a69be",
"metadata": {
"ExecuteTime": {
"end_time": "2025-08-06T03:18:04.726389Z",
"start_time": "2025-08-06T03:18:04.705650Z"
}
},
"outputs": [],
"source": [
"if using_colab:\n",
" import torch\n",
" import torchvision\n",
" print(\"PyTorch version:\", torch.__version__)\n",
" print(\"Torchvision version:\", torchvision.__version__)\n",
" print(\"CUDA is available:\", torch.cuda.is_available())\n",
" import sys\n",
" !{sys.executable} -m pip install opencv-python matplotlib onnx onnxruntime\n",
" !{sys.executable} -m pip install 'git+https://github.com/facebookresearch/segment-anything.git'\n",
" \n",
" !mkdir images\n",
" !wget -P images https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/truck.jpg\n",
" \n",
" !wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
]
},
{
"cell_type": "markdown",
"id": "dc4a58be",
"metadata": {},
"source": [
"## Set-up"
]
},
{
"cell_type": "markdown",
"id": "42396e8d",
"metadata": {},
"source": [
"Note that this notebook requires both the `onnx` and `onnxruntime` optional dependencies, in addition to `opencv-python` and `matplotlib` for visualization."
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "2c712610",
"metadata": {
"ExecuteTime": {
"end_time": "2025-08-06T03:18:15.221624Z",
"start_time": "2025-08-06T03:18:12.230502Z"
}
},
"outputs": [],
"source": [
"import torch\n",
"import numpy as np\n",
"import cv2\n",
"import matplotlib.pyplot as plt\n",
"from segment_anything import sam_model_registry, SamPredictor\n",
"from segment_anything.utils.onnx import SamOnnxModel\n",
"\n",
"import onnxruntime\n",
"from onnxruntime.quantization import QuantType\n",
"from onnxruntime.quantization.quantize import quantize_dynamic"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "f29441b9",
"metadata": {
"ExecuteTime": {
"end_time": "2025-08-06T03:18:18.375013Z",
"start_time": "2025-08-06T03:18:18.359691Z"
}
},
"outputs": [],
"source": [
"def show_mask(mask, ax):\n",
" color = np.array([30/255, 144/255, 255/255, 0.6])\n",
" h, w = mask.shape[-2:]\n",
" mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
" ax.imshow(mask_image)\n",
" \n",
"def show_points(coords, labels, ax, marker_size=375):\n",
" pos_points = coords[labels==1]\n",
" neg_points = coords[labels==0]\n",
" ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
" ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25) \n",
" \n",
"def show_box(box, ax):\n",
" x0, y0 = box[0], box[1]\n",
" w, h = box[2] - box[0], box[3] - box[1]\n",
" ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2)) "
]
},
{
"cell_type": "markdown",
"id": "bd0f6b2b",
"metadata": {},
"source": [
"## Export an ONNX model"
]
},
{
"cell_type": "markdown",
"id": "1540f719",
"metadata": {},
"source": [
"Set the path below to a SAM model checkpoint, then load the model. This will be needed to both export the model and to calculate embeddings for the model."
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "76fc53f4",
"metadata": {
"ExecuteTime": {
"end_time": "2025-08-06T03:18:52.842965Z",
"start_time": "2025-08-06T03:18:52.839365Z"
}
},
"outputs": [],
"source": [
"checkpoint = \"/home/alab/PycharmProjects/segment-anything/checkpoint/sam_vit_h_4b8939.pth\"\n",
"model_type = \"vit_h\""
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "11bfc8aa",
"metadata": {
"ExecuteTime": {
"end_time": "2025-08-06T03:19:02.574461Z",
"start_time": "2025-08-06T03:18:55.883913Z"
}
},
"outputs": [],
"source": [
"sam = sam_model_registry[model_type](checkpoint=checkpoint)"
]
},
{
"cell_type": "markdown",
"id": "450c089c",
"metadata": {},
"source": [
"The script `segment-anything/scripts/export_onnx_model.py` can be used to export the necessary portion of SAM. Alternatively, run the following code to export an ONNX model. If you have already exported a model, set the path below and skip to the next section. Assure that the exported ONNX model aligns with the checkpoint and model type set above. This notebook expects the model was exported with the parameter `return_single_mask=True`."
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "38a8add8",
"metadata": {
"ExecuteTime": {
"end_time": "2025-08-06T03:19:08.682003Z",
"start_time": "2025-08-06T03:19:08.677719Z"
}
},
"outputs": [],
"source": [
"onnx_model_path = None # Set to use an already exported model, then skip to the next section."
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "7da638ba",
"metadata": {
"ExecuteTime": {
"end_time": "2025-08-06T03:22:08.108690Z",
"start_time": "2025-08-06T03:22:07.697634Z"
}
},
"outputs": [],
"source": [
"import warnings\n",
"\n",
"onnx_model_path = \"/home/alab/PycharmProjects/segment-anything/notebooks/sam_onnx_example.onnx\"\n",
"\n",
"onnx_model = SamOnnxModel(sam, return_single_mask=True)\n",
"\n",
"dynamic_axes = {\n",
" \"point_coords\": {1: \"num_points\"},\n",
" \"point_labels\": {1: \"num_points\"},\n",
"}\n",
"\n",
"embed_dim = sam.prompt_encoder.embed_dim\n",
"embed_size = sam.prompt_encoder.image_embedding_size\n",
"mask_input_size = [4 * x for x in embed_size]\n",
"dummy_inputs = {\n",
" \"image_embeddings\": torch.randn(1, embed_dim, *embed_size, dtype=torch.float),\n",
" \"point_coords\": torch.randint(low=0, high=1024, size=(1, 5, 2), dtype=torch.float),\n",
" \"point_labels\": torch.randint(low=0, high=4, size=(1, 5), dtype=torch.float),\n",
" \"mask_input\": torch.randn(1, 1, *mask_input_size, dtype=torch.float),\n",
" \"has_mask_input\": torch.tensor([1], dtype=torch.float),\n",
" \"orig_im_size\": torch.tensor([1500, 2250], dtype=torch.float),\n",
"}\n",
"output_names = [\"masks\", \"iou_predictions\", \"low_res_masks\"]\n",
"\n",
"with warnings.catch_warnings():\n",
" warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)\n",
" warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
" with open(onnx_model_path, \"wb\") as f:\n",
" torch.onnx.export(\n",
" onnx_model,\n",
" tuple(dummy_inputs.values()),\n",
" f,\n",
" export_params=True,\n",
" verbose=False,\n",
" opset_version=17,\n",
" do_constant_folding=True,\n",
" input_names=list(dummy_inputs.keys()),\n",
" output_names=output_names,\n",
" dynamic_axes=dynamic_axes,\n",
" ) "
]
},
{
"cell_type": "markdown",
"id": "c450cf1a",
"metadata": {},
"source": [
"If desired, the model can additionally be quantized and optimized. We find this improves web runtime significantly for negligible change in qualitative performance. Run the next cell to quantize the model, or skip to the next section otherwise."
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "235d39fe",
"metadata": {
"ExecuteTime": {
"end_time": "2025-08-06T03:19:17.443097Z",
"start_time": "2025-08-06T03:19:17.408587Z"
}
},
"outputs": [
{
"ename": "TypeError",
"evalue": "quantize_dynamic() got an unexpected keyword argument 'optimize_model'",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[13], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m onnx_model_quantized_path \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msam_onnx_quantized_example.onnx\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m----> 2\u001b[0m \u001b[43mquantize_dynamic\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 3\u001b[0m \u001b[43m \u001b[49m\u001b[43mmodel_input\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43monnx_model_path\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 4\u001b[0m \u001b[43m \u001b[49m\u001b[43mmodel_output\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43monnx_model_quantized_path\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 5\u001b[0m \u001b[43m \u001b[49m\u001b[43moptimize_model\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m 6\u001b[0m \u001b[43m \u001b[49m\u001b[43mper_channel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m 7\u001b[0m \u001b[43m \u001b[49m\u001b[43mreduce_range\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m 8\u001b[0m \u001b[43m \u001b[49m\u001b[43mweight_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mQuantType\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mQUInt8\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 9\u001b[0m \u001b[43m)\u001b[49m\n\u001b[1;32m 10\u001b[0m onnx_model_path \u001b[38;5;241m=\u001b[39m onnx_model_quantized_path\n",
"\u001b[0;31mTypeError\u001b[0m: quantize_dynamic() got an unexpected keyword argument 'optimize_model'"
]
}
],
"source": [
"onnx_model_quantized_path = \"/home/alab/PycharmProjects/segment-anything/notebooks/sam_onnx_example.onnx\"\n",
"quantize_dynamic(\n",
" model_input=onnx_model_path,\n",
" model_output=onnx_model_quantized_path,\n",
" optimize_model=True,\n",
" per_channel=False,\n",
" reduce_range=False,\n",
" weight_type=QuantType.QUInt8,\n",
")\n",
"onnx_model_path = onnx_model_quantized_path"
]
},
{
"cell_type": "markdown",
"id": "927a928b",
"metadata": {},
"source": [
"## Example Image"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6be6eb55",
"metadata": {},
"outputs": [],
"source": [
"image = cv2.imread('images/truck.jpg')\n",
"image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b7e9a27a",
"metadata": {},
"outputs": [],
"source": [
"plt.figure(figsize=(10,10))\n",
"plt.imshow(image)\n",
"plt.axis('on')\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "027b177b",
"metadata": {},
"source": [
"## Using an ONNX model"
]
},
{
"cell_type": "markdown",
"id": "778d4593",
"metadata": {},
"source": [
"Here as an example, we use `onnxruntime` in python on CPU to execute the ONNX model. However, any platform that supports an ONNX runtime could be used in principle. Launch the runtime session below:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9689b1bf",
"metadata": {},
"outputs": [],
"source": [
"ort_session = onnxruntime.InferenceSession(onnx_model_path)"
]
},
{
"cell_type": "markdown",
"id": "7708ead6",
"metadata": {},
"source": [
"To use the ONNX model, the image must first be pre-processed using the SAM image encoder. This is a heavier weight process best performed on GPU. SamPredictor can be used as normal, then `.get_image_embedding()` will retreive the intermediate features."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "26e067b4",
"metadata": {},
"outputs": [],
"source": [
"sam.to(device='cuda')\n",
"predictor = SamPredictor(sam)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7ad3f0d6",
"metadata": {},
"outputs": [],
"source": [
"predictor.set_image(image)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8a6f0f07",
"metadata": {},
"outputs": [],
"source": [
"image_embedding = predictor.get_image_embedding().cpu().numpy()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5e112f33",
"metadata": {},
"outputs": [],
"source": [
"image_embedding.shape"
]
},
{
"cell_type": "markdown",
"id": "6337b654",
"metadata": {},
"source": [
"The ONNX model has a different input signature than `SamPredictor.predict`. The following inputs must all be supplied. Note the special cases for both point and mask inputs. All inputs are `np.float32`.\n",
"* `image_embeddings`: The image embedding from `predictor.get_image_embedding()`. Has a batch index of length 1.\n",
"* `point_coords`: Coordinates of sparse input prompts, corresponding to both point inputs and box inputs. Boxes are encoded using two points, one for the top-left corner and one for the bottom-right corner. *Coordinates must already be transformed to long-side 1024.* Has a batch index of length 1.\n",
"* `point_labels`: Labels for the sparse input prompts. 0 is a negative input point, 1 is a positive input point, 2 is a top-left box corner, 3 is a bottom-right box corner, and -1 is a padding point. *If there is no box input, a single padding point with label -1 and coordinates (0.0, 0.0) should be concatenated.*\n",
"* `mask_input`: A mask input to the model with shape 1x1x256x256. This must be supplied even if there is no mask input. In this case, it can just be zeros.\n",
"* `has_mask_input`: An indicator for the mask input. 1 indicates a mask input, 0 indicates no mask input.\n",
"* `orig_im_size`: The size of the input image in (H,W) format, before any transformation. \n",
"\n",
"Additionally, the ONNX model does not threshold the output mask logits. To obtain a binary mask, threshold at `sam.mask_threshold` (equal to 0.0)."
]
},
{
"cell_type": "markdown",
"id": "bf5a9f55",
"metadata": {},
"source": [
"### Example point input"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1c0deef0",
"metadata": {},
"outputs": [],
"source": [
"input_point = np.array([[500, 375]])\n",
"input_label = np.array([1])"
]
},
{
"cell_type": "markdown",
"id": "7256394c",
"metadata": {},
"source": [
"Add a batch index, concatenate a padding point, and transform."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4f69903e",
"metadata": {},
"outputs": [],
"source": [
"onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]\n",
"onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)\n",
"\n",
"onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)\n"
]
},
{
"cell_type": "markdown",
"id": "b188dc53",
"metadata": {},
"source": [
"Create an empty mask input and an indicator for no mask."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5cb52bcf",
"metadata": {},
"outputs": [],
"source": [
"onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)\n",
"onnx_has_mask_input = np.zeros(1, dtype=np.float32)"
]
},
{
"cell_type": "markdown",
"id": "a99c2cc5",
"metadata": {},
"source": [
"Package the inputs to run in the onnx model"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b1d7ea11",
"metadata": {},
"outputs": [],
"source": [
"ort_inputs = {\n",
" \"image_embeddings\": image_embedding,\n",
" \"point_coords\": onnx_coord,\n",
" \"point_labels\": onnx_label,\n",
" \"mask_input\": onnx_mask_input,\n",
" \"has_mask_input\": onnx_has_mask_input,\n",
" \"orig_im_size\": np.array(image.shape[:2], dtype=np.float32)\n",
"}"
]
},
{
"cell_type": "markdown",
"id": "4b6409c9",
"metadata": {},
"source": [
"Predict a mask and threshold it."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dc4cc082",
"metadata": {},
"outputs": [],
"source": [
"masks, _, low_res_logits = ort_session.run(None, ort_inputs)\n",
"masks = masks > predictor.model.mask_threshold"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d778a8fb",
"metadata": {},
"outputs": [],
"source": [
"masks.shape"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "badb1175",
"metadata": {},
"outputs": [],
"source": [
"plt.figure(figsize=(10,10))\n",
"plt.imshow(image)\n",
"show_mask(masks, plt.gca())\n",
"show_points(input_point, input_label, plt.gca())\n",
"plt.axis('off')\n",
"plt.show() "
]
},
{
"cell_type": "markdown",
"id": "1f1d4d15",
"metadata": {},
"source": [
"### Example mask input"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b319da82",
"metadata": {},
"outputs": [],
"source": [
"input_point = np.array([[500, 375], [1125, 625]])\n",
"input_label = np.array([1, 1])\n",
"\n",
"# Use the mask output from the previous run. It is already in the correct form for input to the ONNX model.\n",
"onnx_mask_input = low_res_logits"
]
},
{
"cell_type": "markdown",
"id": "b1823b37",
"metadata": {},
"source": [
"Transform the points as in the previous example."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8885130f",
"metadata": {},
"outputs": [],
"source": [
"onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]\n",
"onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)\n",
"\n",
"onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)"
]
},
{
"cell_type": "markdown",
"id": "28e47b69",
"metadata": {},
"source": [
"The `has_mask_input` indicator is now 1."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3ab4483a",
"metadata": {},
"outputs": [],
"source": [
"onnx_has_mask_input = np.ones(1, dtype=np.float32)"
]
},
{
"cell_type": "markdown",
"id": "d3781955",
"metadata": {},
"source": [
"Package inputs, then predict and threshold the mask."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0c1ec096",
"metadata": {},
"outputs": [],
"source": [
"ort_inputs = {\n",
" \"image_embeddings\": image_embedding,\n",
" \"point_coords\": onnx_coord,\n",
" \"point_labels\": onnx_label,\n",
" \"mask_input\": onnx_mask_input,\n",
" \"has_mask_input\": onnx_has_mask_input,\n",
" \"orig_im_size\": np.array(image.shape[:2], dtype=np.float32)\n",
"}\n",
"\n",
"masks, _, _ = ort_session.run(None, ort_inputs)\n",
"masks = masks > predictor.model.mask_threshold"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1e36554b",
"metadata": {},
"outputs": [],
"source": [
"plt.figure(figsize=(10,10))\n",
"plt.imshow(image)\n",
"show_mask(masks, plt.gca())\n",
"show_points(input_point, input_label, plt.gca())\n",
"plt.axis('off')\n",
"plt.show() "
]
},
{
"cell_type": "markdown",
"id": "2ef211d0",
"metadata": {},
"source": [
"### Example box and point input"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "51e58d2e",
"metadata": {},
"outputs": [],
"source": [
"input_box = np.array([425, 600, 700, 875])\n",
"input_point = np.array([[575, 750]])\n",
"input_label = np.array([0])"
]
},
{
"cell_type": "markdown",
"id": "6e119dcb",
"metadata": {},
"source": [
"Add a batch index, concatenate a box and point inputs, add the appropriate labels for the box corners, and transform. There is no padding point since the input includes a box input."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "bfbe4911",
"metadata": {},
"outputs": [],
"source": [
"onnx_box_coords = input_box.reshape(2, 2)\n",
"onnx_box_labels = np.array([2,3])\n",
"\n",
"onnx_coord = np.concatenate([input_point, onnx_box_coords], axis=0)[None, :, :]\n",
"onnx_label = np.concatenate([input_label, onnx_box_labels], axis=0)[None, :].astype(np.float32)\n",
"\n",
"onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)"
]
},
{
"cell_type": "markdown",
"id": "65edabd2",
"metadata": {},
"source": [
"Package inputs, then predict and threshold the mask."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2abfba56",
"metadata": {},
"outputs": [],
"source": [
"onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)\n",
"onnx_has_mask_input = np.zeros(1, dtype=np.float32)\n",
"\n",
"ort_inputs = {\n",
" \"image_embeddings\": image_embedding,\n",
" \"point_coords\": onnx_coord,\n",
" \"point_labels\": onnx_label,\n",
" \"mask_input\": onnx_mask_input,\n",
" \"has_mask_input\": onnx_has_mask_input,\n",
" \"orig_im_size\": np.array(image.shape[:2], dtype=np.float32)\n",
"}\n",
"\n",
"masks, _, _ = ort_session.run(None, ort_inputs)\n",
"masks = masks > predictor.model.mask_threshold"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8301bf33",
"metadata": {},
"outputs": [],
"source": [
"plt.figure(figsize=(10, 10))\n",
"plt.imshow(image)\n",
"show_mask(masks[0], plt.gca())\n",
"show_box(input_box, plt.gca())\n",
"show_points(input_point, input_label, plt.gca())\n",
"plt.axis('off')\n",
"plt.show()"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.23"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

11
notebooks/onnx_model_example.py Executable file
View File

@@ -0,0 +1,11 @@
# 导入必要的库
import numpy as np
import torch
import matplotlib.pyplot as plt
from PIL import Image
import requests
from io import BytesIO
# 从 SAM 2 库中导入图像预测器和模型构建工具
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

1016
notebooks/predictor_example.ipynb Executable file

File diff suppressed because one or more lines are too long

Binary file not shown.

After

Width:  |  Height:  |  Size: 7.5 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 6.8 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 3.1 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 3.3 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 3.3 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.9 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 3.5 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.9 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 3.4 KiB

View File

@@ -0,0 +1,10 @@
id,area,bbox_x0,bbox_y0,bbox_w,bbox_h,point_input_x,point_input_y,predicted_iou,stability_score,crop_box_x0,crop_box_y0,crop_box_w,crop_box_h
0,208104,346,119,353,794,496.0,624.0,1.0344942808151245,0.9846644997596741,0,0,1024,1024
1,826080,0,0,1023,1023,1008.0,912.0,1.0164952278137207,0.9937294721603394,0,0,1024,1024
2,3538,495,151,49,108,528.0,176.0,0.9655076861381531,0.960321843624115,0,0,1024,1024
3,2837,588,340,35,124,592.0,368.0,0.9651759266853333,0.9690757393836975,0,0,1024,1024
4,2722,420,347,34,131,432.0,432.0,0.9555162787437439,0.9638205766677856,0,0,1024,1024
5,980,590,447,27,49,592.0,464.0,0.9420101046562195,0.9660339951515198,0,0,1024,1024
6,10617,479,119,87,172,496.0,144.0,0.936710000038147,0.9556418061256409,0,0,1024,1024
7,609,681,528,15,55,688.0,560.0,0.9296139478683472,0.9707317352294922,0,0,1024,1024
8,4588,522,183,79,143,560.0,240.0,0.9280573725700378,0.9508405923843384,0,0,1024,1024
1 id area bbox_x0 bbox_y0 bbox_w bbox_h point_input_x point_input_y predicted_iou stability_score crop_box_x0 crop_box_y0 crop_box_w crop_box_h
2 0 208104 346 119 353 794 496.0 624.0 1.0344942808151245 0.9846644997596741 0 0 1024 1024
3 1 826080 0 0 1023 1023 1008.0 912.0 1.0164952278137207 0.9937294721603394 0 0 1024 1024
4 2 3538 495 151 49 108 528.0 176.0 0.9655076861381531 0.960321843624115 0 0 1024 1024
5 3 2837 588 340 35 124 592.0 368.0 0.9651759266853333 0.9690757393836975 0 0 1024 1024
6 4 2722 420 347 34 131 432.0 432.0 0.9555162787437439 0.9638205766677856 0 0 1024 1024
7 5 980 590 447 27 49 592.0 464.0 0.9420101046562195 0.9660339951515198 0 0 1024 1024
8 6 10617 479 119 87 172 496.0 144.0 0.936710000038147 0.9556418061256409 0 0 1024 1024
9 7 609 681 528 15 55 688.0 560.0 0.9296139478683472 0.9707317352294922 0 0 1024 1024
10 8 4588 522 183 79 143 560.0 240.0 0.9280573725700378 0.9508405923843384 0 0 1024 1024

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.3 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.3 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.3 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.9 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.9 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 3.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.2 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.2 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.1 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.1 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 3.4 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.3 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.9 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.9 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 KiB

View File

@@ -0,0 +1,28 @@
id,area,bbox_x0,bbox_y0,bbox_w,bbox_h,point_input_x,point_input_y,predicted_iou,stability_score,crop_box_x0,crop_box_y0,crop_box_w,crop_box_h
0,318264,4,0,533,899,354.90625,42.28125,1.025822401046753,0.9901363253593445,0,0,554,902
1,2908,194,8,67,91,233.71875,42.28125,0.9839344024658203,0.9722222089767456,0,0,554,902
2,10997,91,78,68,234,112.53125,267.78125,0.9781885743141174,0.9724088311195374,0,0,554,902
3,2910,280,7,67,92,320.28125,42.28125,0.977243185043335,0.9685065746307373,0,0,554,902
4,19292,391,302,146,207,441.46875,408.71875,0.9717334508895874,0.9758270978927612,0,0,554,902
5,204,262,572,17,14,268.34375,577.84375,0.9706144332885742,0.9950980544090271,0,0,554,902
6,209,261,456,18,15,268.34375,465.09375,0.9689342379570007,0.9763033390045166,0,0,554,902
7,10832,381,81,67,231,424.15625,267.78125,0.968802273273468,0.9669151306152344,0,0,554,902
8,190,263,398,15,15,268.34375,408.71875,0.967627227306366,0.9947643876075745,0,0,554,902
9,206,262,283,16,15,268.34375,295.96875,0.9655775427818298,0.9855769276618958,0,0,554,902
10,216,262,687,17,15,268.34375,690.59375,0.9645803570747375,0.9907407164573669,0,0,554,902
11,222,262,744,17,16,268.34375,746.96875,0.9638096690177917,0.9955157041549683,0,0,554,902
12,203,262,803,16,15,268.34375,803.34375,0.9622183442115784,0.9756097793579102,0,0,554,902
13,216,262,629,17,16,268.34375,634.21875,0.9615152478218079,0.9862385392189026,0,0,554,902
14,205,263,514,15,15,268.34375,521.46875,0.9610558152198792,0.9902439117431641,0,0,554,902
15,216,261,167,17,15,268.34375,183.21875,0.9606413245201111,0.9817351698875427,0,0,554,902
16,919,154,294,17,84,164.46875,352.34375,0.9570132493972778,0.9583777785301208,0,0,554,902
17,237,261,224,19,16,268.34375,239.59375,0.9552502632141113,0.9915611743927002,0,0,554,902
18,19664,4,302,147,207,129.84375,465.09375,0.9529095888137817,0.9699330925941467,0,0,554,902
19,19614,3,131,148,378,112.53125,436.90625,0.9519698619842529,0.9640650153160095,0,0,554,902
20,3890,224,4,93,61,302.96875,14.09375,0.9364364743232727,0.9555724859237671,0,0,554,902
21,224,262,109,17,16,268.34375,126.84375,0.9293429255485535,0.9647576808929443,0,0,554,902
22,203,263,69,16,15,268.34375,70.46875,0.9020635485649109,0.9708737730979919,0,0,554,902
23,338,310,252,11,94,320.28125,295.96875,0.8969216346740723,0.9504373073577881,0,0,554,902
24,520,261,345,21,30,268.34375,352.34375,0.8908350467681885,0.9606741666793823,0,0,554,902
25,149,333,300,9,43,337.59375,324.15625,0.8839569091796875,0.9599999785423279,0,0,554,902
26,39591,2,0,209,446,43.28125,42.28125,0.8828506469726562,0.9530083537101746,0,0,554,902
1 id area bbox_x0 bbox_y0 bbox_w bbox_h point_input_x point_input_y predicted_iou stability_score crop_box_x0 crop_box_y0 crop_box_w crop_box_h
2 0 318264 4 0 533 899 354.90625 42.28125 1.025822401046753 0.9901363253593445 0 0 554 902
3 1 2908 194 8 67 91 233.71875 42.28125 0.9839344024658203 0.9722222089767456 0 0 554 902
4 2 10997 91 78 68 234 112.53125 267.78125 0.9781885743141174 0.9724088311195374 0 0 554 902
5 3 2910 280 7 67 92 320.28125 42.28125 0.977243185043335 0.9685065746307373 0 0 554 902
6 4 19292 391 302 146 207 441.46875 408.71875 0.9717334508895874 0.9758270978927612 0 0 554 902
7 5 204 262 572 17 14 268.34375 577.84375 0.9706144332885742 0.9950980544090271 0 0 554 902
8 6 209 261 456 18 15 268.34375 465.09375 0.9689342379570007 0.9763033390045166 0 0 554 902
9 7 10832 381 81 67 231 424.15625 267.78125 0.968802273273468 0.9669151306152344 0 0 554 902
10 8 190 263 398 15 15 268.34375 408.71875 0.967627227306366 0.9947643876075745 0 0 554 902
11 9 206 262 283 16 15 268.34375 295.96875 0.9655775427818298 0.9855769276618958 0 0 554 902
12 10 216 262 687 17 15 268.34375 690.59375 0.9645803570747375 0.9907407164573669 0 0 554 902
13 11 222 262 744 17 16 268.34375 746.96875 0.9638096690177917 0.9955157041549683 0 0 554 902
14 12 203 262 803 16 15 268.34375 803.34375 0.9622183442115784 0.9756097793579102 0 0 554 902
15 13 216 262 629 17 16 268.34375 634.21875 0.9615152478218079 0.9862385392189026 0 0 554 902
16 14 205 263 514 15 15 268.34375 521.46875 0.9610558152198792 0.9902439117431641 0 0 554 902
17 15 216 261 167 17 15 268.34375 183.21875 0.9606413245201111 0.9817351698875427 0 0 554 902
18 16 919 154 294 17 84 164.46875 352.34375 0.9570132493972778 0.9583777785301208 0 0 554 902
19 17 237 261 224 19 16 268.34375 239.59375 0.9552502632141113 0.9915611743927002 0 0 554 902
20 18 19664 4 302 147 207 129.84375 465.09375 0.9529095888137817 0.9699330925941467 0 0 554 902
21 19 19614 3 131 148 378 112.53125 436.90625 0.9519698619842529 0.9640650153160095 0 0 554 902
22 20 3890 224 4 93 61 302.96875 14.09375 0.9364364743232727 0.9555724859237671 0 0 554 902
23 21 224 262 109 17 16 268.34375 126.84375 0.9293429255485535 0.9647576808929443 0 0 554 902
24 22 203 263 69 16 15 268.34375 70.46875 0.9020635485649109 0.9708737730979919 0 0 554 902
25 23 338 310 252 11 94 320.28125 295.96875 0.8969216346740723 0.9504373073577881 0 0 554 902
26 24 520 261 345 21 30 268.34375 352.34375 0.8908350467681885 0.9606741666793823 0 0 554 902
27 25 149 333 300 9 43 337.59375 324.15625 0.8839569091796875 0.9599999785423279 0 0 554 902
28 26 39591 2 0 209 446 43.28125 42.28125 0.8828506469726562 0.9530083537101746 0 0 554 902

BIN
output/dog_and_girl/0.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 9.8 KiB

BIN
output/dog_and_girl/1.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 6.9 KiB

BIN
output/dog_and_girl/10.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 6.4 KiB

BIN
output/dog_and_girl/11.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.2 KiB

BIN
output/dog_and_girl/12.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.1 KiB

BIN
output/dog_and_girl/13.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.0 KiB

BIN
output/dog_and_girl/14.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 6.2 KiB

BIN
output/dog_and_girl/15.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.2 KiB

BIN
output/dog_and_girl/16.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 7.0 KiB

BIN
output/dog_and_girl/17.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.1 KiB

BIN
output/dog_and_girl/18.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.0 KiB

BIN
output/dog_and_girl/19.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.1 KiB

BIN
output/dog_and_girl/2.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 6.0 KiB

BIN
output/dog_and_girl/20.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.1 KiB

BIN
output/dog_and_girl/21.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.5 KiB

BIN
output/dog_and_girl/22.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.3 KiB

BIN
output/dog_and_girl/23.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.1 KiB

BIN
output/dog_and_girl/24.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.1 KiB

BIN
output/dog_and_girl/25.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.3 KiB

BIN
output/dog_and_girl/26.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 6.6 KiB

BIN
output/dog_and_girl/27.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.0 KiB

Some files were not shown because too many files have changed in this diff Show More