liuyizhang commited on
Commit
b9ad21d
1 Parent(s): 53720d1

delete segment_anything

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. requirements.txt +1 -1
  2. segment_anything/.flake8 +0 -7
  3. segment_anything/CODE_OF_CONDUCT.md +0 -80
  4. segment_anything/CONTRIBUTING.md +0 -31
  5. segment_anything/LICENSE +0 -201
  6. segment_anything/README.md +0 -107
  7. segment_anything/linter.sh +0 -32
  8. segment_anything/notebooks/automatic_mask_generator_example.ipynb +0 -0
  9. segment_anything/notebooks/images/dog.jpg +0 -0
  10. segment_anything/notebooks/images/groceries.jpg +0 -0
  11. segment_anything/notebooks/images/truck.jpg +0 -0
  12. segment_anything/notebooks/onnx_model_example.ipynb +0 -774
  13. segment_anything/notebooks/predictor_example.ipynb +0 -0
  14. segment_anything/scripts/amg.py +0 -238
  15. segment_anything/scripts/export_onnx_model.py +0 -204
  16. segment_anything/segment_anything.egg-info/PKG-INFO +0 -13
  17. segment_anything/segment_anything.egg-info/SOURCES.txt +0 -24
  18. segment_anything/segment_anything.egg-info/dependency_links.txt +0 -1
  19. segment_anything/segment_anything.egg-info/requires.txt +0 -13
  20. segment_anything/segment_anything.egg-info/top_level.txt +0 -1
  21. segment_anything/segment_anything/__init__.py +0 -15
  22. segment_anything/segment_anything/__pycache__/__init__.cpython-39.pyc +0 -0
  23. segment_anything/segment_anything/__pycache__/automatic_mask_generator.cpython-39.pyc +0 -0
  24. segment_anything/segment_anything/__pycache__/build_sam.cpython-39.pyc +0 -0
  25. segment_anything/segment_anything/__pycache__/predictor.cpython-39.pyc +0 -0
  26. segment_anything/segment_anything/automatic_mask_generator.py +0 -372
  27. segment_anything/segment_anything/build_sam.py +0 -107
  28. segment_anything/segment_anything/modeling/__init__.py +0 -11
  29. segment_anything/segment_anything/modeling/__pycache__/__init__.cpython-39.pyc +0 -0
  30. segment_anything/segment_anything/modeling/__pycache__/common.cpython-39.pyc +0 -0
  31. segment_anything/segment_anything/modeling/__pycache__/image_encoder.cpython-39.pyc +0 -0
  32. segment_anything/segment_anything/modeling/__pycache__/mask_decoder.cpython-39.pyc +0 -0
  33. segment_anything/segment_anything/modeling/__pycache__/prompt_encoder.cpython-39.pyc +0 -0
  34. segment_anything/segment_anything/modeling/__pycache__/sam.cpython-39.pyc +0 -0
  35. segment_anything/segment_anything/modeling/__pycache__/transformer.cpython-39.pyc +0 -0
  36. segment_anything/segment_anything/modeling/common.py +0 -43
  37. segment_anything/segment_anything/modeling/image_encoder.py +0 -395
  38. segment_anything/segment_anything/modeling/mask_decoder.py +0 -176
  39. segment_anything/segment_anything/modeling/prompt_encoder.py +0 -214
  40. segment_anything/segment_anything/modeling/sam.py +0 -174
  41. segment_anything/segment_anything/modeling/transformer.py +0 -240
  42. segment_anything/segment_anything/predictor.py +0 -269
  43. segment_anything/segment_anything/utils/__init__.py +0 -5
  44. segment_anything/segment_anything/utils/__pycache__/__init__.cpython-39.pyc +0 -0
  45. segment_anything/segment_anything/utils/__pycache__/amg.cpython-39.pyc +0 -0
  46. segment_anything/segment_anything/utils/__pycache__/transforms.cpython-39.pyc +0 -0
  47. segment_anything/segment_anything/utils/amg.py +0 -346
  48. segment_anything/segment_anything/utils/onnx.py +0 -144
  49. segment_anything/segment_anything/utils/transforms.py +0 -102
  50. segment_anything/setup.cfg +0 -11
requirements.txt CHANGED
@@ -20,7 +20,7 @@ torchvision
20
  transformers
21
  yapf
22
  numba
23
- segment_anything
24
 
25
  lama-cleaner==0.25.0
26
  openmim==0.1.5
 
20
  transformers
21
  yapf
22
  numba
23
+ # segment_anything
24
 
25
  lama-cleaner==0.25.0
26
  openmim==0.1.5
segment_anything/.flake8 DELETED
@@ -1,7 +0,0 @@
1
- [flake8]
2
- ignore = W503, E203, E221, C901, C408, E741, C407, B017, F811, C101, EXE001, EXE002
3
- max-line-length = 100
4
- max-complexity = 18
5
- select = B,C,E,F,W,T4,B9
6
- per-file-ignores =
7
- **/__init__.py:F401,F403,E402
 
 
 
 
 
 
 
 
segment_anything/CODE_OF_CONDUCT.md DELETED
@@ -1,80 +0,0 @@
1
- # Code of Conduct
2
-
3
- ## Our Pledge
4
-
5
- In the interest of fostering an open and welcoming environment, we as
6
- contributors and maintainers pledge to make participation in our project and
7
- our community a harassment-free experience for everyone, regardless of age, body
8
- size, disability, ethnicity, sex characteristics, gender identity and expression,
9
- level of experience, education, socio-economic status, nationality, personal
10
- appearance, race, religion, or sexual identity and orientation.
11
-
12
- ## Our Standards
13
-
14
- Examples of behavior that contributes to creating a positive environment
15
- include:
16
-
17
- * Using welcoming and inclusive language
18
- * Being respectful of differing viewpoints and experiences
19
- * Gracefully accepting constructive criticism
20
- * Focusing on what is best for the community
21
- * Showing empathy towards other community members
22
-
23
- Examples of unacceptable behavior by participants include:
24
-
25
- * The use of sexualized language or imagery and unwelcome sexual attention or
26
- advances
27
- * Trolling, insulting/derogatory comments, and personal or political attacks
28
- * Public or private harassment
29
- * Publishing others' private information, such as a physical or electronic
30
- address, without explicit permission
31
- * Other conduct which could reasonably be considered inappropriate in a
32
- professional setting
33
-
34
- ## Our Responsibilities
35
-
36
- Project maintainers are responsible for clarifying the standards of acceptable
37
- behavior and are expected to take appropriate and fair corrective action in
38
- response to any instances of unacceptable behavior.
39
-
40
- Project maintainers have the right and responsibility to remove, edit, or
41
- reject comments, commits, code, wiki edits, issues, and other contributions
42
- that are not aligned to this Code of Conduct, or to ban temporarily or
43
- permanently any contributor for other behaviors that they deem inappropriate,
44
- threatening, offensive, or harmful.
45
-
46
- ## Scope
47
-
48
- This Code of Conduct applies within all project spaces, and it also applies when
49
- an individual is representing the project or its community in public spaces.
50
- Examples of representing a project or community include using an official
51
- project e-mail address, posting via an official social media account, or acting
52
- as an appointed representative at an online or offline event. Representation of
53
- a project may be further defined and clarified by project maintainers.
54
-
55
- This Code of Conduct also applies outside the project spaces when there is a
56
- reasonable belief that an individual's behavior may have a negative impact on
57
- the project or its community.
58
-
59
- ## Enforcement
60
-
61
- Instances of abusive, harassing, or otherwise unacceptable behavior may be
62
- reported by contacting the project team at <[email protected]>. All
63
- complaints will be reviewed and investigated and will result in a response that
64
- is deemed necessary and appropriate to the circumstances. The project team is
65
- obligated to maintain confidentiality with regard to the reporter of an incident.
66
- Further details of specific enforcement policies may be posted separately.
67
-
68
- Project maintainers who do not follow or enforce the Code of Conduct in good
69
- faith may face temporary or permanent repercussions as determined by other
70
- members of the project's leadership.
71
-
72
- ## Attribution
73
-
74
- This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
75
- available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
76
-
77
- [homepage]: https://www.contributor-covenant.org
78
-
79
- For answers to common questions about this code of conduct, see
80
- https://www.contributor-covenant.org/faq
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/CONTRIBUTING.md DELETED
@@ -1,31 +0,0 @@
1
- # Contributing to segment-anything
2
- We want to make contributing to this project as easy and transparent as
3
- possible.
4
-
5
- ## Pull Requests
6
- We actively welcome your pull requests.
7
-
8
- 1. Fork the repo and create your branch from `main`.
9
- 2. If you've added code that should be tested, add tests.
10
- 3. If you've changed APIs, update the documentation.
11
- 4. Ensure the test suite passes.
12
- 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`.
13
- 6. If you haven't already, complete the Contributor License Agreement ("CLA").
14
-
15
- ## Contributor License Agreement ("CLA")
16
- In order to accept your pull request, we need you to submit a CLA. You only need
17
- to do this once to work on any of Facebook's open source projects.
18
-
19
- Complete your CLA here: <https://code.facebook.com/cla>
20
-
21
- ## Issues
22
- We use GitHub issues to track public bugs. Please ensure your description is
23
- clear and has sufficient instructions to be able to reproduce the issue.
24
-
25
- Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe
26
- disclosure of security bugs. In those cases, please go through the process
27
- outlined on that page and do not file a public issue.
28
-
29
- ## License
30
- By contributing to segment-anything, you agree that your contributions will be licensed
31
- under the LICENSE file in the root directory of this source tree.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/LICENSE DELETED
@@ -1,201 +0,0 @@
1
- Apache License
2
- Version 2.0, January 2004
3
- http://www.apache.org/licenses/
4
-
5
- TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
-
7
- 1. Definitions.
8
-
9
- "License" shall mean the terms and conditions for use, reproduction,
10
- and distribution as defined by Sections 1 through 9 of this document.
11
-
12
- "Licensor" shall mean the copyright owner or entity authorized by
13
- the copyright owner that is granting the License.
14
-
15
- "Legal Entity" shall mean the union of the acting entity and all
16
- other entities that control, are controlled by, or are under common
17
- control with that entity. For the purposes of this definition,
18
- "control" means (i) the power, direct or indirect, to cause the
19
- direction or management of such entity, whether by contract or
20
- otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
- outstanding shares, or (iii) beneficial ownership of such entity.
22
-
23
- "You" (or "Your") shall mean an individual or Legal Entity
24
- exercising permissions granted by this License.
25
-
26
- "Source" form shall mean the preferred form for making modifications,
27
- including but not limited to software source code, documentation
28
- source, and configuration files.
29
-
30
- "Object" form shall mean any form resulting from mechanical
31
- transformation or translation of a Source form, including but
32
- not limited to compiled object code, generated documentation,
33
- and conversions to other media types.
34
-
35
- "Work" shall mean the work of authorship, whether in Source or
36
- Object form, made available under the License, as indicated by a
37
- copyright notice that is included in or attached to the work
38
- (an example is provided in the Appendix below).
39
-
40
- "Derivative Works" shall mean any work, whether in Source or Object
41
- form, that is based on (or derived from) the Work and for which the
42
- editorial revisions, annotations, elaborations, or other modifications
43
- represent, as a whole, an original work of authorship. For the purposes
44
- of this License, Derivative Works shall not include works that remain
45
- separable from, or merely link (or bind by name) to the interfaces of,
46
- the Work and Derivative Works thereof.
47
-
48
- "Contribution" shall mean any work of authorship, including
49
- the original version of the Work and any modifications or additions
50
- to that Work or Derivative Works thereof, that is intentionally
51
- submitted to Licensor for inclusion in the Work by the copyright owner
52
- or by an individual or Legal Entity authorized to submit on behalf of
53
- the copyright owner. For the purposes of this definition, "submitted"
54
- means any form of electronic, verbal, or written communication sent
55
- to the Licensor or its representatives, including but not limited to
56
- communication on electronic mailing lists, source code control systems,
57
- and issue tracking systems that are managed by, or on behalf of, the
58
- Licensor for the purpose of discussing and improving the Work, but
59
- excluding communication that is conspicuously marked or otherwise
60
- designated in writing by the copyright owner as "Not a Contribution."
61
-
62
- "Contributor" shall mean Licensor and any individual or Legal Entity
63
- on behalf of whom a Contribution has been received by Licensor and
64
- subsequently incorporated within the Work.
65
-
66
- 2. Grant of Copyright License. Subject to the terms and conditions of
67
- this License, each Contributor hereby grants to You a perpetual,
68
- worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69
- copyright license to reproduce, prepare Derivative Works of,
70
- publicly display, publicly perform, sublicense, and distribute the
71
- Work and such Derivative Works in Source or Object form.
72
-
73
- 3. Grant of Patent License. Subject to the terms and conditions of
74
- this License, each Contributor hereby grants to You a perpetual,
75
- worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76
- (except as stated in this section) patent license to make, have made,
77
- use, offer to sell, sell, import, and otherwise transfer the Work,
78
- where such license applies only to those patent claims licensable
79
- by such Contributor that are necessarily infringed by their
80
- Contribution(s) alone or by combination of their Contribution(s)
81
- with the Work to which such Contribution(s) was submitted. If You
82
- institute patent litigation against any entity (including a
83
- cross-claim or counterclaim in a lawsuit) alleging that the Work
84
- or a Contribution incorporated within the Work constitutes direct
85
- or contributory patent infringement, then any patent licenses
86
- granted to You under this License for that Work shall terminate
87
- as of the date such litigation is filed.
88
-
89
- 4. Redistribution. You may reproduce and distribute copies of the
90
- Work or Derivative Works thereof in any medium, with or without
91
- modifications, and in Source or Object form, provided that You
92
- meet the following conditions:
93
-
94
- (a) You must give any other recipients of the Work or
95
- Derivative Works a copy of this License; and
96
-
97
- (b) You must cause any modified files to carry prominent notices
98
- stating that You changed the files; and
99
-
100
- (c) You must retain, in the Source form of any Derivative Works
101
- that You distribute, all copyright, patent, trademark, and
102
- attribution notices from the Source form of the Work,
103
- excluding those notices that do not pertain to any part of
104
- the Derivative Works; and
105
-
106
- (d) If the Work includes a "NOTICE" text file as part of its
107
- distribution, then any Derivative Works that You distribute must
108
- include a readable copy of the attribution notices contained
109
- within such NOTICE file, excluding those notices that do not
110
- pertain to any part of the Derivative Works, in at least one
111
- of the following places: within a NOTICE text file distributed
112
- as part of the Derivative Works; within the Source form or
113
- documentation, if provided along with the Derivative Works; or,
114
- within a display generated by the Derivative Works, if and
115
- wherever such third-party notices normally appear. The contents
116
- of the NOTICE file are for informational purposes only and
117
- do not modify the License. You may add Your own attribution
118
- notices within Derivative Works that You distribute, alongside
119
- or as an addendum to the NOTICE text from the Work, provided
120
- that such additional attribution notices cannot be construed
121
- as modifying the License.
122
-
123
- You may add Your own copyright statement to Your modifications and
124
- may provide additional or different license terms and conditions
125
- for use, reproduction, or distribution of Your modifications, or
126
- for any such Derivative Works as a whole, provided Your use,
127
- reproduction, and distribution of the Work otherwise complies with
128
- the conditions stated in this License.
129
-
130
- 5. Submission of Contributions. Unless You explicitly state otherwise,
131
- any Contribution intentionally submitted for inclusion in the Work
132
- by You to the Licensor shall be under the terms and conditions of
133
- this License, without any additional terms or conditions.
134
- Notwithstanding the above, nothing herein shall supersede or modify
135
- the terms of any separate license agreement you may have executed
136
- with Licensor regarding such Contributions.
137
-
138
- 6. Trademarks. This License does not grant permission to use the trade
139
- names, trademarks, service marks, or product names of the Licensor,
140
- except as required for reasonable and customary use in describing the
141
- origin of the Work and reproducing the content of the NOTICE file.
142
-
143
- 7. Disclaimer of Warranty. Unless required by applicable law or
144
- agreed to in writing, Licensor provides the Work (and each
145
- Contributor provides its Contributions) on an "AS IS" BASIS,
146
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147
- implied, including, without limitation, any warranties or conditions
148
- of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149
- PARTICULAR PURPOSE. You are solely responsible for determining the
150
- appropriateness of using or redistributing the Work and assume any
151
- risks associated with Your exercise of permissions under this License.
152
-
153
- 8. Limitation of Liability. In no event and under no legal theory,
154
- whether in tort (including negligence), contract, or otherwise,
155
- unless required by applicable law (such as deliberate and grossly
156
- negligent acts) or agreed to in writing, shall any Contributor be
157
- liable to You for damages, including any direct, indirect, special,
158
- incidental, or consequential damages of any character arising as a
159
- result of this License or out of the use or inability to use the
160
- Work (including but not limited to damages for loss of goodwill,
161
- work stoppage, computer failure or malfunction, or any and all
162
- other commercial damages or losses), even if such Contributor
163
- has been advised of the possibility of such damages.
164
-
165
- 9. Accepting Warranty or Additional Liability. While redistributing
166
- the Work or Derivative Works thereof, You may choose to offer,
167
- and charge a fee for, acceptance of support, warranty, indemnity,
168
- or other liability obligations and/or rights consistent with this
169
- License. However, in accepting such obligations, You may act only
170
- on Your own behalf and on Your sole responsibility, not on behalf
171
- of any other Contributor, and only if You agree to indemnify,
172
- defend, and hold each Contributor harmless for any liability
173
- incurred by, or claims asserted against, such Contributor by reason
174
- of your accepting any such warranty or additional liability.
175
-
176
- END OF TERMS AND CONDITIONS
177
-
178
- APPENDIX: How to apply the Apache License to your work.
179
-
180
- To apply the Apache License to your work, attach the following
181
- boilerplate notice, with the fields enclosed by brackets "[]"
182
- replaced with your own identifying information. (Don't include
183
- the brackets!) The text should be enclosed in the appropriate
184
- comment syntax for the file format. We also recommend that a
185
- file or class name and description of purpose be included on the
186
- same "printed page" as the copyright notice for easier
187
- identification within third-party archives.
188
-
189
- Copyright [yyyy] [name of copyright owner]
190
-
191
- Licensed under the Apache License, Version 2.0 (the "License");
192
- you may not use this file except in compliance with the License.
193
- You may obtain a copy of the License at
194
-
195
- http://www.apache.org/licenses/LICENSE-2.0
196
-
197
- Unless required by applicable law or agreed to in writing, software
198
- distributed under the License is distributed on an "AS IS" BASIS,
199
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200
- See the License for the specific language governing permissions and
201
- limitations under the License.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/README.md DELETED
@@ -1,107 +0,0 @@
1
- # Segment Anything
2
-
3
- **[Meta AI Research, FAIR](https://ai.facebook.com/research/)**
4
-
5
- [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/)
6
-
7
- [[`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/)]
8
-
9
- ![SAM design](assets/model_diagram.png?raw=true)
10
-
11
- 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.
12
-
13
- <p float="left">
14
- <img src="assets/masks1.png?raw=true" width="37.25%" />
15
- <img src="assets/masks2.jpg?raw=true" width="61.5%" />
16
- </p>
17
-
18
- ## Installation
19
-
20
- 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.
21
-
22
- Install Segment Anything:
23
-
24
- ```
25
- pip install git+https://github.com/facebookresearch/segment-anything.git
26
- ```
27
-
28
- or clone the repository locally and install with
29
-
30
- ```
31
- git clone [email protected]:facebookresearch/segment-anything.git
32
- cd segment-anything; pip install -e .
33
- ```
34
-
35
- 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.
36
- ```
37
- pip install opencv-python pycocotools matplotlib onnxruntime onnx
38
- ```
39
-
40
-
41
- ## <a name="GettingStarted"></a>Getting Started
42
-
43
- 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:
44
-
45
- ```
46
- from segment_anything import build_sam, SamPredictor
47
- predictor = SamPredictor(build_sam(checkpoint="</path/to/model.pth>"))
48
- predictor.set_image(<your_image>)
49
- masks, _, _ = predictor.predict(<input_prompts>)
50
- ```
51
-
52
- or generate masks for an entire image:
53
-
54
- ```
55
- from segment_anything import build_sam, SamAutomaticMaskGenerator
56
- mask_generator = SamAutomaticMaskGenerator(build_sam(checkpoint="</path/to/model.pth>"))
57
- masks = mask_generator_generate(<your_image>)
58
- ```
59
-
60
- Additionally, masks can be generated for images from the command line:
61
-
62
- ```
63
- python scripts/amg.py --checkpoint <path/to/sam/checkpoint> --input <image_or_folder> --output <output_directory>
64
- ```
65
-
66
- 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.
67
-
68
- <p float="left">
69
- <img src="assets/notebook1.png?raw=true" width="49.1%" />
70
- <img src="assets/notebook2.png?raw=true" width="48.9%" />
71
- </p>
72
-
73
- ## ONNX Export
74
-
75
- 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
76
-
77
- ```
78
- python scripts/export_onnx_model.py --checkpoint <path/to/checkpoint> --output <path/to/output>
79
- ```
80
-
81
- 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.
82
-
83
- ## <a name="Models"></a>Model Checkpoints
84
-
85
- Three model versions of the model are available with different backbone sizes. These models can be instantiated by running
86
- ```
87
- from segment_anything import sam_model_registry
88
- sam = sam_model_registry["<name>"](checkpoint="<path/to/checkpoint>")
89
- ```
90
- Click the links below to download the checkpoint for the corresponding model name. The default model in bold can also be instantiated with `build_sam`, as in the examples in [Getting Started](#getting-started).
91
-
92
- * **`default` or `vit_h`: [ViT-H SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth)**
93
- * `vit_l`: [ViT-L SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth)
94
- * `vit_b`: [ViT-B SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth)
95
-
96
- ## License
97
- The model is licensed under the [Apache 2.0 license](LICENSE).
98
-
99
- ## Contributing
100
-
101
- See [contributing](CONTRIBUTING.md) and the [code of conduct](CODE_OF_CONDUCT.md).
102
-
103
- ## Contributors
104
-
105
- The Segment Anything project was made possible with the help of many contributors (alphabetical):
106
-
107
- 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
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/linter.sh DELETED
@@ -1,32 +0,0 @@
1
- #!/bin/bash -e
2
- # Copyright (c) Facebook, Inc. and its affiliates.
3
-
4
- {
5
- black --version | grep -E "23\." > /dev/null
6
- } || {
7
- echo "Linter requires 'black==23.*' !"
8
- exit 1
9
- }
10
-
11
- ISORT_VERSION=$(isort --version-number)
12
- if [[ "$ISORT_VERSION" != 5.12* ]]; then
13
- echo "Linter requires isort==5.12.0 !"
14
- exit 1
15
- fi
16
-
17
- echo "Running isort ..."
18
- isort . --atomic
19
-
20
- echo "Running black ..."
21
- black -l 100 .
22
-
23
- echo "Running flake8 ..."
24
- if [ -x "$(command -v flake8)" ]; then
25
- flake8 .
26
- else
27
- python3 -m flake8 .
28
- fi
29
-
30
- echo "Running mypy..."
31
-
32
- mypy --exclude 'setup.py|notebooks' .
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/notebooks/automatic_mask_generator_example.ipynb DELETED
The diff for this file is too large to render. See raw diff
 
segment_anything/notebooks/images/dog.jpg DELETED
Binary file (99.8 kB)
 
segment_anything/notebooks/images/groceries.jpg DELETED
Binary file (168 kB)
 
segment_anything/notebooks/images/truck.jpg DELETED
Binary file (271 kB)
 
segment_anything/notebooks/onnx_model_example.ipynb DELETED
@@ -1,774 +0,0 @@
1
- {
2
- "cells": [
3
- {
4
- "cell_type": "code",
5
- "execution_count": null,
6
- "id": "901c8ef3",
7
- "metadata": {},
8
- "outputs": [],
9
- "source": [
10
- "# Copyright (c) Meta Platforms, Inc. and affiliates."
11
- ]
12
- },
13
- {
14
- "cell_type": "markdown",
15
- "id": "1662bb7c",
16
- "metadata": {},
17
- "source": [
18
- "# Produces masks from prompts using an ONNX model"
19
- ]
20
- },
21
- {
22
- "cell_type": "markdown",
23
- "id": "7fcc21a0",
24
- "metadata": {},
25
- "source": [
26
- "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."
27
- ]
28
- },
29
- {
30
- "cell_type": "code",
31
- "execution_count": 4,
32
- "id": "86daff77",
33
- "metadata": {},
34
- "outputs": [
35
- {
36
- "data": {
37
- "text/html": [
38
- "\n",
39
- "<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb\">\n",
40
- " <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
41
- "</a>\n"
42
- ],
43
- "text/plain": [
44
- "<IPython.core.display.HTML object>"
45
- ]
46
- },
47
- "metadata": {},
48
- "output_type": "display_data"
49
- }
50
- ],
51
- "source": [
52
- "from IPython.display import display, HTML\n",
53
- "display(HTML(\n",
54
- "\"\"\"\n",
55
- "<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb\">\n",
56
- " <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
57
- "</a>\n",
58
- "\"\"\"\n",
59
- "))"
60
- ]
61
- },
62
- {
63
- "cell_type": "markdown",
64
- "id": "55ae4e00",
65
- "metadata": {},
66
- "source": [
67
- "## Environment Set-up"
68
- ]
69
- },
70
- {
71
- "cell_type": "markdown",
72
- "id": "109a5cc2",
73
- "metadata": {},
74
- "source": [
75
- "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_collab=True` below and run the cell. In Colab, be sure to select 'GPU' under 'Edit'->'Notebook Settings'->'Hardware accelerator'."
76
- ]
77
- },
78
- {
79
- "cell_type": "code",
80
- "execution_count": 5,
81
- "id": "39b99fc4",
82
- "metadata": {},
83
- "outputs": [],
84
- "source": [
85
- "using_colab = False"
86
- ]
87
- },
88
- {
89
- "cell_type": "code",
90
- "execution_count": 6,
91
- "id": "296a69be",
92
- "metadata": {},
93
- "outputs": [],
94
- "source": [
95
- "if using_colab:\n",
96
- " import torch\n",
97
- " import torchvision\n",
98
- " print(\"PyTorch version:\", torch.__version__)\n",
99
- " print(\"Torchvision version:\", torchvision.__version__)\n",
100
- " print(\"CUDA is available:\", torch.cuda.is_available())\n",
101
- " import sys\n",
102
- " !{sys.executable} -m pip install opencv-python matplotlib onnx onnxruntime\n",
103
- " !{sys.executable} -m pip install 'git+https://github.com/facebookresearch/segment-anything.git'\n",
104
- " \n",
105
- " !mkdir images\n",
106
- " !wget -P images https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/truck.jpg\n",
107
- " \n",
108
- " !wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
109
- ]
110
- },
111
- {
112
- "cell_type": "markdown",
113
- "id": "dc4a58be",
114
- "metadata": {},
115
- "source": [
116
- "## Set-up"
117
- ]
118
- },
119
- {
120
- "cell_type": "markdown",
121
- "id": "42396e8d",
122
- "metadata": {},
123
- "source": [
124
- "Note that this notebook requires both the `onnx` and `onnxruntime` optional dependencies, in addition to `opencv-python` and `matplotlib` for visualization."
125
- ]
126
- },
127
- {
128
- "cell_type": "code",
129
- "execution_count": null,
130
- "id": "2c712610",
131
- "metadata": {},
132
- "outputs": [],
133
- "source": [
134
- "import torch\n",
135
- "import numpy as np\n",
136
- "import cv2\n",
137
- "import matplotlib.pyplot as plt\n",
138
- "from segment_anything import sam_model_registry, SamPredictor\n",
139
- "from segment_anything.utils.onnx import SamOnnxModel\n",
140
- "\n",
141
- "import onnxruntime\n",
142
- "from onnxruntime.quantization import QuantType\n",
143
- "from onnxruntime.quantization.quantize import quantize_dynamic"
144
- ]
145
- },
146
- {
147
- "cell_type": "code",
148
- "execution_count": null,
149
- "id": "f29441b9",
150
- "metadata": {},
151
- "outputs": [],
152
- "source": [
153
- "def show_mask(mask, ax):\n",
154
- " color = np.array([30/255, 144/255, 255/255, 0.6])\n",
155
- " h, w = mask.shape[-2:]\n",
156
- " mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
157
- " ax.imshow(mask_image)\n",
158
- " \n",
159
- "def show_points(coords, labels, ax, marker_size=375):\n",
160
- " pos_points = coords[labels==1]\n",
161
- " neg_points = coords[labels==0]\n",
162
- " ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
163
- " ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25) \n",
164
- " \n",
165
- "def show_box(box, ax):\n",
166
- " x0, y0 = box[0], box[1]\n",
167
- " w, h = box[2] - box[0], box[3] - box[1]\n",
168
- " ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2)) "
169
- ]
170
- },
171
- {
172
- "cell_type": "markdown",
173
- "id": "bd0f6b2b",
174
- "metadata": {},
175
- "source": [
176
- "## Export an ONNX model"
177
- ]
178
- },
179
- {
180
- "cell_type": "markdown",
181
- "id": "1540f719",
182
- "metadata": {},
183
- "source": [
184
- "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."
185
- ]
186
- },
187
- {
188
- "cell_type": "code",
189
- "execution_count": null,
190
- "id": "76fc53f4",
191
- "metadata": {},
192
- "outputs": [],
193
- "source": [
194
- "checkpoint = \"sam_vit_h_4b8939.pth\"\n",
195
- "model_type = \"default\""
196
- ]
197
- },
198
- {
199
- "cell_type": "code",
200
- "execution_count": null,
201
- "id": "11bfc8aa",
202
- "metadata": {},
203
- "outputs": [],
204
- "source": [
205
- "sam = sam_model_registry[model_type](checkpoint=checkpoint)"
206
- ]
207
- },
208
- {
209
- "cell_type": "markdown",
210
- "id": "450c089c",
211
- "metadata": {},
212
- "source": [
213
- "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`."
214
- ]
215
- },
216
- {
217
- "cell_type": "code",
218
- "execution_count": null,
219
- "id": "38a8add8",
220
- "metadata": {},
221
- "outputs": [],
222
- "source": [
223
- "onnx_model_path = None # Set to use an already exported model, then skip to the next section."
224
- ]
225
- },
226
- {
227
- "cell_type": "code",
228
- "execution_count": null,
229
- "id": "7da638ba",
230
- "metadata": {
231
- "scrolled": false
232
- },
233
- "outputs": [],
234
- "source": [
235
- "import warnings\n",
236
- "\n",
237
- "onnx_model_path = \"sam_onnx_example.onnx\"\n",
238
- "\n",
239
- "onnx_model = SamOnnxModel(sam, return_single_mask=True)\n",
240
- "\n",
241
- "dynamic_axes = {\n",
242
- " \"point_coords\": {1: \"num_points\"},\n",
243
- " \"point_labels\": {1: \"num_points\"},\n",
244
- "}\n",
245
- "\n",
246
- "embed_dim = sam.prompt_encoder.embed_dim\n",
247
- "embed_size = sam.prompt_encoder.image_embedding_size\n",
248
- "mask_input_size = [4 * x for x in embed_size]\n",
249
- "dummy_inputs = {\n",
250
- " \"image_embeddings\": torch.randn(1, embed_dim, *embed_size, dtype=torch.float),\n",
251
- " \"point_coords\": torch.randint(low=0, high=1024, size=(1, 5, 2), dtype=torch.float),\n",
252
- " \"point_labels\": torch.randint(low=0, high=4, size=(1, 5), dtype=torch.float),\n",
253
- " \"mask_input\": torch.randn(1, 1, *mask_input_size, dtype=torch.float),\n",
254
- " \"has_mask_input\": torch.tensor([1], dtype=torch.float),\n",
255
- " \"orig_im_size\": torch.tensor([1500, 2250], dtype=torch.float),\n",
256
- "}\n",
257
- "output_names = [\"masks\", \"iou_predictions\", \"low_res_masks\"]\n",
258
- "\n",
259
- "with warnings.catch_warnings():\n",
260
- " warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)\n",
261
- " warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
262
- " with open(onnx_model_path, \"wb\") as f:\n",
263
- " torch.onnx.export(\n",
264
- " onnx_model,\n",
265
- " tuple(dummy_inputs.values()),\n",
266
- " f,\n",
267
- " export_params=True,\n",
268
- " verbose=False,\n",
269
- " opset_version=17,\n",
270
- " do_constant_folding=True,\n",
271
- " input_names=list(dummy_inputs.keys()),\n",
272
- " output_names=output_names,\n",
273
- " dynamic_axes=dynamic_axes,\n",
274
- " ) "
275
- ]
276
- },
277
- {
278
- "cell_type": "markdown",
279
- "id": "c450cf1a",
280
- "metadata": {},
281
- "source": [
282
- "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."
283
- ]
284
- },
285
- {
286
- "cell_type": "code",
287
- "execution_count": null,
288
- "id": "235d39fe",
289
- "metadata": {},
290
- "outputs": [],
291
- "source": [
292
- "onnx_model_quantized_path = \"sam_onnx_quantized_example.onnx\"\n",
293
- "quantize_dynamic(\n",
294
- " model_input=onnx_model_path,\n",
295
- " model_output=onnx_model_quantized_path,\n",
296
- " optimize_model=True,\n",
297
- " per_channel=False,\n",
298
- " reduce_range=False,\n",
299
- " weight_type=QuantType.QUInt8,\n",
300
- ")\n",
301
- "onnx_model_path = onnx_model_quantized_path"
302
- ]
303
- },
304
- {
305
- "cell_type": "markdown",
306
- "id": "927a928b",
307
- "metadata": {},
308
- "source": [
309
- "## Example Image"
310
- ]
311
- },
312
- {
313
- "cell_type": "code",
314
- "execution_count": null,
315
- "id": "6be6eb55",
316
- "metadata": {},
317
- "outputs": [],
318
- "source": [
319
- "image = cv2.imread('images/truck.jpg')\n",
320
- "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)"
321
- ]
322
- },
323
- {
324
- "cell_type": "code",
325
- "execution_count": null,
326
- "id": "b7e9a27a",
327
- "metadata": {},
328
- "outputs": [],
329
- "source": [
330
- "plt.figure(figsize=(10,10))\n",
331
- "plt.imshow(image)\n",
332
- "plt.axis('on')\n",
333
- "plt.show()"
334
- ]
335
- },
336
- {
337
- "cell_type": "markdown",
338
- "id": "027b177b",
339
- "metadata": {},
340
- "source": [
341
- "## Using an ONNX model"
342
- ]
343
- },
344
- {
345
- "cell_type": "markdown",
346
- "id": "778d4593",
347
- "metadata": {},
348
- "source": [
349
- "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:"
350
- ]
351
- },
352
- {
353
- "cell_type": "code",
354
- "execution_count": null,
355
- "id": "9689b1bf",
356
- "metadata": {},
357
- "outputs": [],
358
- "source": [
359
- "ort_session = onnxruntime.InferenceSession(onnx_model_path)"
360
- ]
361
- },
362
- {
363
- "cell_type": "markdown",
364
- "id": "7708ead6",
365
- "metadata": {},
366
- "source": [
367
- "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."
368
- ]
369
- },
370
- {
371
- "cell_type": "code",
372
- "execution_count": null,
373
- "id": "26e067b4",
374
- "metadata": {},
375
- "outputs": [],
376
- "source": [
377
- "sam.to(device='cuda')\n",
378
- "predictor = SamPredictor(sam)"
379
- ]
380
- },
381
- {
382
- "cell_type": "code",
383
- "execution_count": null,
384
- "id": "7ad3f0d6",
385
- "metadata": {},
386
- "outputs": [],
387
- "source": [
388
- "predictor.set_image(image)"
389
- ]
390
- },
391
- {
392
- "cell_type": "code",
393
- "execution_count": null,
394
- "id": "8a6f0f07",
395
- "metadata": {},
396
- "outputs": [],
397
- "source": [
398
- "image_embedding = predictor.get_image_embedding().cpu().numpy()"
399
- ]
400
- },
401
- {
402
- "cell_type": "code",
403
- "execution_count": null,
404
- "id": "5e112f33",
405
- "metadata": {},
406
- "outputs": [],
407
- "source": [
408
- "image_embedding.shape"
409
- ]
410
- },
411
- {
412
- "cell_type": "markdown",
413
- "id": "6337b654",
414
- "metadata": {},
415
- "source": [
416
- "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",
417
- "* `image_embeddings`: The image embedding from `predictor.get_image_embedding()`. Has a batch index of length 1.\n",
418
- "* `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",
419
- "* `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",
420
- "* `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",
421
- "* `has_mask_input`: An indicator for the mask input. 1 indicates a mask input, 0 indicates no mask input.\n",
422
- "* `orig_im_size`: The size of the input image in (H,W) format, before any transformation. \n",
423
- "\n",
424
- "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)."
425
- ]
426
- },
427
- {
428
- "cell_type": "markdown",
429
- "id": "bf5a9f55",
430
- "metadata": {},
431
- "source": [
432
- "### Example point input"
433
- ]
434
- },
435
- {
436
- "cell_type": "code",
437
- "execution_count": null,
438
- "id": "1c0deef0",
439
- "metadata": {},
440
- "outputs": [],
441
- "source": [
442
- "input_point = np.array([[500, 375]])\n",
443
- "input_label = np.array([1])"
444
- ]
445
- },
446
- {
447
- "cell_type": "markdown",
448
- "id": "7256394c",
449
- "metadata": {},
450
- "source": [
451
- "Add a batch index, concatenate a padding point, and transform."
452
- ]
453
- },
454
- {
455
- "cell_type": "code",
456
- "execution_count": null,
457
- "id": "4f69903e",
458
- "metadata": {},
459
- "outputs": [],
460
- "source": [
461
- "onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]\n",
462
- "onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)\n",
463
- "\n",
464
- "onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)\n"
465
- ]
466
- },
467
- {
468
- "cell_type": "markdown",
469
- "id": "b188dc53",
470
- "metadata": {},
471
- "source": [
472
- "Create an empty mask input and an indicator for no mask."
473
- ]
474
- },
475
- {
476
- "cell_type": "code",
477
- "execution_count": null,
478
- "id": "5cb52bcf",
479
- "metadata": {},
480
- "outputs": [],
481
- "source": [
482
- "onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)\n",
483
- "onnx_has_mask_input = np.zeros(1, dtype=np.float32)"
484
- ]
485
- },
486
- {
487
- "cell_type": "markdown",
488
- "id": "a99c2cc5",
489
- "metadata": {},
490
- "source": [
491
- "Package the inputs to run in the onnx model"
492
- ]
493
- },
494
- {
495
- "cell_type": "code",
496
- "execution_count": null,
497
- "id": "b1d7ea11",
498
- "metadata": {},
499
- "outputs": [],
500
- "source": [
501
- "ort_inputs = {\n",
502
- " \"image_embeddings\": image_embedding,\n",
503
- " \"point_coords\": onnx_coord,\n",
504
- " \"point_labels\": onnx_label,\n",
505
- " \"mask_input\": onnx_mask_input,\n",
506
- " \"has_mask_input\": onnx_has_mask_input,\n",
507
- " \"orig_im_size\": np.array(image.shape[:2], dtype=np.float32)\n",
508
- "}"
509
- ]
510
- },
511
- {
512
- "cell_type": "markdown",
513
- "id": "4b6409c9",
514
- "metadata": {},
515
- "source": [
516
- "Predict a mask and threshold it."
517
- ]
518
- },
519
- {
520
- "cell_type": "code",
521
- "execution_count": null,
522
- "id": "dc4cc082",
523
- "metadata": {
524
- "scrolled": false
525
- },
526
- "outputs": [],
527
- "source": [
528
- "masks, _, low_res_logits = ort_session.run(None, ort_inputs)\n",
529
- "masks = masks > predictor.model.mask_threshold"
530
- ]
531
- },
532
- {
533
- "cell_type": "code",
534
- "execution_count": null,
535
- "id": "d778a8fb",
536
- "metadata": {},
537
- "outputs": [],
538
- "source": [
539
- "masks.shape"
540
- ]
541
- },
542
- {
543
- "cell_type": "code",
544
- "execution_count": null,
545
- "id": "badb1175",
546
- "metadata": {},
547
- "outputs": [],
548
- "source": [
549
- "plt.figure(figsize=(10,10))\n",
550
- "plt.imshow(image)\n",
551
- "show_mask(masks, plt.gca())\n",
552
- "show_points(input_point, input_label, plt.gca())\n",
553
- "plt.axis('off')\n",
554
- "plt.show() "
555
- ]
556
- },
557
- {
558
- "cell_type": "markdown",
559
- "id": "1f1d4d15",
560
- "metadata": {},
561
- "source": [
562
- "### Example mask input"
563
- ]
564
- },
565
- {
566
- "cell_type": "code",
567
- "execution_count": null,
568
- "id": "b319da82",
569
- "metadata": {},
570
- "outputs": [],
571
- "source": [
572
- "input_point = np.array([[500, 375], [1125, 625]])\n",
573
- "input_label = np.array([1, 1])\n",
574
- "\n",
575
- "# Use the mask output from the previous run. It is already in the correct form for input to the ONNX model.\n",
576
- "onnx_mask_input = low_res_logits"
577
- ]
578
- },
579
- {
580
- "cell_type": "markdown",
581
- "id": "b1823b37",
582
- "metadata": {},
583
- "source": [
584
- "Transform the points as in the previous example."
585
- ]
586
- },
587
- {
588
- "cell_type": "code",
589
- "execution_count": null,
590
- "id": "8885130f",
591
- "metadata": {},
592
- "outputs": [],
593
- "source": [
594
- "onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]\n",
595
- "onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)\n",
596
- "\n",
597
- "onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)"
598
- ]
599
- },
600
- {
601
- "cell_type": "markdown",
602
- "id": "28e47b69",
603
- "metadata": {},
604
- "source": [
605
- "The `has_mask_input` indicator is now 1."
606
- ]
607
- },
608
- {
609
- "cell_type": "code",
610
- "execution_count": null,
611
- "id": "3ab4483a",
612
- "metadata": {},
613
- "outputs": [],
614
- "source": [
615
- "onnx_has_mask_input = np.ones(1, dtype=np.float32)"
616
- ]
617
- },
618
- {
619
- "cell_type": "markdown",
620
- "id": "d3781955",
621
- "metadata": {},
622
- "source": [
623
- "Package inputs, then predict and threshold the mask."
624
- ]
625
- },
626
- {
627
- "cell_type": "code",
628
- "execution_count": null,
629
- "id": "0c1ec096",
630
- "metadata": {},
631
- "outputs": [],
632
- "source": [
633
- "ort_inputs = {\n",
634
- " \"image_embeddings\": image_embedding,\n",
635
- " \"point_coords\": onnx_coord,\n",
636
- " \"point_labels\": onnx_label,\n",
637
- " \"mask_input\": onnx_mask_input,\n",
638
- " \"has_mask_input\": onnx_has_mask_input,\n",
639
- " \"orig_im_size\": np.array(image.shape[:2], dtype=np.float32)\n",
640
- "}\n",
641
- "\n",
642
- "masks, _, _ = ort_session.run(None, ort_inputs)\n",
643
- "masks = masks > predictor.model.mask_threshold"
644
- ]
645
- },
646
- {
647
- "cell_type": "code",
648
- "execution_count": null,
649
- "id": "1e36554b",
650
- "metadata": {},
651
- "outputs": [],
652
- "source": [
653
- "plt.figure(figsize=(10,10))\n",
654
- "plt.imshow(image)\n",
655
- "show_mask(masks, plt.gca())\n",
656
- "show_points(input_point, input_label, plt.gca())\n",
657
- "plt.axis('off')\n",
658
- "plt.show() "
659
- ]
660
- },
661
- {
662
- "cell_type": "markdown",
663
- "id": "2ef211d0",
664
- "metadata": {},
665
- "source": [
666
- "### Example box and point input"
667
- ]
668
- },
669
- {
670
- "cell_type": "code",
671
- "execution_count": null,
672
- "id": "51e58d2e",
673
- "metadata": {},
674
- "outputs": [],
675
- "source": [
676
- "input_box = np.array([425, 600, 700, 875])\n",
677
- "input_point = np.array([[575, 750]])\n",
678
- "input_label = np.array([0])"
679
- ]
680
- },
681
- {
682
- "cell_type": "markdown",
683
- "id": "6e119dcb",
684
- "metadata": {},
685
- "source": [
686
- "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."
687
- ]
688
- },
689
- {
690
- "cell_type": "code",
691
- "execution_count": null,
692
- "id": "bfbe4911",
693
- "metadata": {},
694
- "outputs": [],
695
- "source": [
696
- "onnx_box_coords = input_box.reshape(2, 2)\n",
697
- "onnx_box_labels = np.array([2,3])\n",
698
- "\n",
699
- "onnx_coord = np.concatenate([input_point, onnx_box_coords], axis=0)[None, :, :]\n",
700
- "onnx_label = np.concatenate([input_label, onnx_box_labels], axis=0)[None, :].astype(np.float32)\n",
701
- "\n",
702
- "onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)"
703
- ]
704
- },
705
- {
706
- "cell_type": "markdown",
707
- "id": "65edabd2",
708
- "metadata": {},
709
- "source": [
710
- "Package inputs, then predict and threshold the mask."
711
- ]
712
- },
713
- {
714
- "cell_type": "code",
715
- "execution_count": null,
716
- "id": "2abfba56",
717
- "metadata": {},
718
- "outputs": [],
719
- "source": [
720
- "onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)\n",
721
- "onnx_has_mask_input = np.zeros(1, dtype=np.float32)\n",
722
- "\n",
723
- "ort_inputs = {\n",
724
- " \"image_embeddings\": image_embedding,\n",
725
- " \"point_coords\": onnx_coord,\n",
726
- " \"point_labels\": onnx_label,\n",
727
- " \"mask_input\": onnx_mask_input,\n",
728
- " \"has_mask_input\": onnx_has_mask_input,\n",
729
- " \"orig_im_size\": np.array(image.shape[:2], dtype=np.float32)\n",
730
- "}\n",
731
- "\n",
732
- "masks, _, _ = ort_session.run(None, ort_inputs)\n",
733
- "masks = masks > predictor.model.mask_threshold"
734
- ]
735
- },
736
- {
737
- "cell_type": "code",
738
- "execution_count": null,
739
- "id": "8301bf33",
740
- "metadata": {},
741
- "outputs": [],
742
- "source": [
743
- "plt.figure(figsize=(10, 10))\n",
744
- "plt.imshow(image)\n",
745
- "show_mask(masks[0], plt.gca())\n",
746
- "show_box(input_box, plt.gca())\n",
747
- "show_points(input_point, input_label, plt.gca())\n",
748
- "plt.axis('off')\n",
749
- "plt.show()"
750
- ]
751
- }
752
- ],
753
- "metadata": {
754
- "kernelspec": {
755
- "display_name": "Python 3 (ipykernel)",
756
- "language": "python",
757
- "name": "python3"
758
- },
759
- "language_info": {
760
- "codemirror_mode": {
761
- "name": "ipython",
762
- "version": 3
763
- },
764
- "file_extension": ".py",
765
- "mimetype": "text/x-python",
766
- "name": "python",
767
- "nbconvert_exporter": "python",
768
- "pygments_lexer": "ipython3",
769
- "version": "3.10.10"
770
- }
771
- },
772
- "nbformat": 4,
773
- "nbformat_minor": 5
774
- }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/notebooks/predictor_example.ipynb DELETED
The diff for this file is too large to render. See raw diff
 
segment_anything/scripts/amg.py DELETED
@@ -1,238 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import cv2 # type: ignore
8
-
9
- from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
10
-
11
- import argparse
12
- import json
13
- import os
14
- from typing import Any, Dict, List
15
-
16
- parser = argparse.ArgumentParser(
17
- description=(
18
- "Runs automatic mask generation on an input image or directory of images, "
19
- "and outputs masks as either PNGs or COCO-style RLEs. Requires open-cv, "
20
- "as well as pycocotools if saving in RLE format."
21
- )
22
- )
23
-
24
- parser.add_argument(
25
- "--input",
26
- type=str,
27
- required=True,
28
- help="Path to either a single input image or folder of images.",
29
- )
30
-
31
- parser.add_argument(
32
- "--output",
33
- type=str,
34
- required=True,
35
- help=(
36
- "Path to the directory where masks will be output. Output will be either a folder "
37
- "of PNGs per image or a single json with COCO-style masks."
38
- ),
39
- )
40
-
41
- parser.add_argument(
42
- "--model-type",
43
- type=str,
44
- default="default",
45
- help="The type of model to load, in ['default', 'vit_l', 'vit_b']",
46
- )
47
-
48
- parser.add_argument(
49
- "--checkpoint",
50
- type=str,
51
- required=True,
52
- help="The path to the SAM checkpoint to use for mask generation.",
53
- )
54
-
55
- parser.add_argument("--device", type=str, default="cuda", help="The device to run generation on.")
56
-
57
- parser.add_argument(
58
- "--convert-to-rle",
59
- action="store_true",
60
- help=(
61
- "Save masks as COCO RLEs in a single json instead of as a folder of PNGs. "
62
- "Requires pycocotools."
63
- ),
64
- )
65
-
66
- amg_settings = parser.add_argument_group("AMG Settings")
67
-
68
- amg_settings.add_argument(
69
- "--points-per-side",
70
- type=int,
71
- default=None,
72
- help="Generate masks by sampling a grid over the image with this many points to a side.",
73
- )
74
-
75
- amg_settings.add_argument(
76
- "--points-per-batch",
77
- type=int,
78
- default=None,
79
- help="How many input points to process simultaneously in one batch.",
80
- )
81
-
82
- amg_settings.add_argument(
83
- "--pred-iou-thresh",
84
- type=float,
85
- default=None,
86
- help="Exclude masks with a predicted score from the model that is lower than this threshold.",
87
- )
88
-
89
- amg_settings.add_argument(
90
- "--stability-score-thresh",
91
- type=float,
92
- default=None,
93
- help="Exclude masks with a stability score lower than this threshold.",
94
- )
95
-
96
- amg_settings.add_argument(
97
- "--stability-score-offset",
98
- type=float,
99
- default=None,
100
- help="Larger values perturb the mask more when measuring stability score.",
101
- )
102
-
103
- amg_settings.add_argument(
104
- "--box-nms-thresh",
105
- type=float,
106
- default=None,
107
- help="The overlap threshold for excluding a duplicate mask.",
108
- )
109
-
110
- amg_settings.add_argument(
111
- "--crop-n-layers",
112
- type=int,
113
- default=None,
114
- help=(
115
- "If >0, mask generation is run on smaller crops of the image to generate more masks. "
116
- "The value sets how many different scales to crop at."
117
- ),
118
- )
119
-
120
- amg_settings.add_argument(
121
- "--crop-nms-thresh",
122
- type=float,
123
- default=None,
124
- help="The overlap threshold for excluding duplicate masks across different crops.",
125
- )
126
-
127
- amg_settings.add_argument(
128
- "--crop-overlap-ratio",
129
- type=int,
130
- default=None,
131
- help="Larger numbers mean image crops will overlap more.",
132
- )
133
-
134
- amg_settings.add_argument(
135
- "--crop-n-points-downscale-factor",
136
- type=int,
137
- default=None,
138
- help="The number of points-per-side in each layer of crop is reduced by this factor.",
139
- )
140
-
141
- amg_settings.add_argument(
142
- "--min-mask-region-area",
143
- type=int,
144
- default=None,
145
- help=(
146
- "Disconnected mask regions or holes with area smaller than this value "
147
- "in pixels are removed by postprocessing."
148
- ),
149
- )
150
-
151
-
152
- def write_masks_to_folder(masks: List[Dict[str, Any]], path: str) -> None:
153
- header = "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" # noqa
154
- metadata = [header]
155
- for i, mask_data in enumerate(masks):
156
- mask = mask_data["segmentation"]
157
- filename = f"{i}.png"
158
- cv2.imwrite(os.path.join(path, filename), mask * 255)
159
- mask_metadata = [
160
- str(i),
161
- str(mask_data["area"]),
162
- *[str(x) for x in mask_data["bbox"]],
163
- *[str(x) for x in mask_data["point_coords"][0]],
164
- str(mask_data["predicted_iou"]),
165
- str(mask_data["stability_score"]),
166
- *[str(x) for x in mask_data["crop_box"]],
167
- ]
168
- row = ",".join(mask_metadata)
169
- metadata.append(row)
170
- metadata_path = os.path.join(path, "metadata.csv")
171
- with open(metadata_path, "w") as f:
172
- f.write("\n".join(metadata))
173
-
174
- return
175
-
176
-
177
- def get_amg_kwargs(args):
178
- amg_kwargs = {
179
- "points_per_side": args.points_per_side,
180
- "points_per_batch": args.points_per_batch,
181
- "pred_iou_thresh": args.pred_iou_thresh,
182
- "stability_score_thresh": args.stability_score_thresh,
183
- "stability_score_offset": args.stability_score_offset,
184
- "box_nms_thresh": args.box_nms_thresh,
185
- "crop_n_layers": args.crop_n_layers,
186
- "crop_nms_thresh": args.crop_nms_thresh,
187
- "crop_overlap_ratio": args.crop_overlap_ratio,
188
- "crop_n_points_downscale_factor": args.crop_n_points_downscale_factor,
189
- "min_mask_region_area": args.min_mask_region_area,
190
- }
191
- amg_kwargs = {k: v for k, v in amg_kwargs.items() if v is not None}
192
- return amg_kwargs
193
-
194
-
195
- def main(args: argparse.Namespace) -> None:
196
- print("Loading model...")
197
- sam = sam_model_registry[args.model_type](checkpoint=args.checkpoint)
198
- _ = sam.to(device=args.device)
199
- output_mode = "coco_rle" if args.convert_to_rle else "binary_mask"
200
- amg_kwargs = get_amg_kwargs(args)
201
- generator = SamAutomaticMaskGenerator(sam, output_mode=output_mode, **amg_kwargs)
202
-
203
- if not os.path.isdir(args.input):
204
- targets = [args.input]
205
- else:
206
- targets = [
207
- f for f in os.listdir(args.input) if not os.path.isdir(os.path.join(args.input, f))
208
- ]
209
- targets = [os.path.join(args.input, f) for f in targets]
210
-
211
- os.makedirs(args.output, exist_ok=True)
212
-
213
- for t in targets:
214
- print(f"Processing '{t}'...")
215
- image = cv2.imread(t)
216
- if image is None:
217
- print(f"Could not load '{t}' as an image, skipping...")
218
- continue
219
- image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
220
-
221
- masks = generator.generate(image)
222
-
223
- base = os.path.basename(t)
224
- base = os.path.splitext(base)[0]
225
- save_base = os.path.join(args.output, base)
226
- if output_mode == "binary_mask":
227
- os.makedirs(save_base, exist_ok=False)
228
- write_masks_to_folder(masks, save_base)
229
- else:
230
- save_file = save_base + ".json"
231
- with open(save_file, "w") as f:
232
- json.dump(masks, f)
233
- print("Done!")
234
-
235
-
236
- if __name__ == "__main__":
237
- args = parser.parse_args()
238
- main(args)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/scripts/export_onnx_model.py DELETED
@@ -1,204 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import torch
8
-
9
- from segment_anything import build_sam, build_sam_vit_b, build_sam_vit_l
10
- from segment_anything.utils.onnx import SamOnnxModel
11
-
12
- import argparse
13
- import warnings
14
-
15
- try:
16
- import onnxruntime # type: ignore
17
-
18
- onnxruntime_exists = True
19
- except ImportError:
20
- onnxruntime_exists = False
21
-
22
- parser = argparse.ArgumentParser(
23
- description="Export the SAM prompt encoder and mask decoder to an ONNX model."
24
- )
25
-
26
- parser.add_argument(
27
- "--checkpoint", type=str, required=True, help="The path to the SAM model checkpoint."
28
- )
29
-
30
- parser.add_argument(
31
- "--output", type=str, required=True, help="The filename to save the ONNX model to."
32
- )
33
-
34
- parser.add_argument(
35
- "--model-type",
36
- type=str,
37
- default="default",
38
- help="In ['default', 'vit_b', 'vit_l']. Which type of SAM model to export.",
39
- )
40
-
41
- parser.add_argument(
42
- "--return-single-mask",
43
- action="store_true",
44
- help=(
45
- "If true, the exported ONNX model will only return the best mask, "
46
- "instead of returning multiple masks. For high resolution images "
47
- "this can improve runtime when upscaling masks is expensive."
48
- ),
49
- )
50
-
51
- parser.add_argument(
52
- "--opset",
53
- type=int,
54
- default=17,
55
- help="The ONNX opset version to use. Must be >=11",
56
- )
57
-
58
- parser.add_argument(
59
- "--quantize-out",
60
- type=str,
61
- default=None,
62
- help=(
63
- "If set, will quantize the model and save it with this name. "
64
- "Quantization is performed with quantize_dynamic from onnxruntime.quantization.quantize."
65
- ),
66
- )
67
-
68
- parser.add_argument(
69
- "--gelu-approximate",
70
- action="store_true",
71
- help=(
72
- "Replace GELU operations with approximations using tanh. Useful "
73
- "for some runtimes that have slow or unimplemented erf ops, used in GELU."
74
- ),
75
- )
76
-
77
- parser.add_argument(
78
- "--use-stability-score",
79
- action="store_true",
80
- help=(
81
- "Replaces the model's predicted mask quality score with the stability "
82
- "score calculated on the low resolution masks using an offset of 1.0. "
83
- ),
84
- )
85
-
86
- parser.add_argument(
87
- "--return-extra-metrics",
88
- action="store_true",
89
- help=(
90
- "The model will return five results: (masks, scores, stability_scores, "
91
- "areas, low_res_logits) instead of the usual three. This can be "
92
- "significantly slower for high resolution outputs."
93
- ),
94
- )
95
-
96
-
97
- def run_export(
98
- model_type: str,
99
- checkpoint: str,
100
- output: str,
101
- opset: int,
102
- return_single_mask: bool,
103
- gelu_approximate: bool = False,
104
- use_stability_score: bool = False,
105
- return_extra_metrics=False,
106
- ):
107
- print("Loading model...")
108
- if model_type == "vit_b":
109
- sam = build_sam_vit_b(checkpoint)
110
- elif model_type == "vit_l":
111
- sam = build_sam_vit_l(checkpoint)
112
- else:
113
- sam = build_sam(checkpoint)
114
-
115
- onnx_model = SamOnnxModel(
116
- model=sam,
117
- return_single_mask=return_single_mask,
118
- use_stability_score=use_stability_score,
119
- return_extra_metrics=return_extra_metrics,
120
- )
121
-
122
- if gelu_approximate:
123
- for n, m in onnx_model.named_modules():
124
- if isinstance(m, torch.nn.GELU):
125
- m.approximate = "tanh"
126
-
127
- dynamic_axes = {
128
- "point_coords": {1: "num_points"},
129
- "point_labels": {1: "num_points"},
130
- }
131
-
132
- embed_dim = sam.prompt_encoder.embed_dim
133
- embed_size = sam.prompt_encoder.image_embedding_size
134
- mask_input_size = [4 * x for x in embed_size]
135
- dummy_inputs = {
136
- "image_embeddings": torch.randn(1, embed_dim, *embed_size, dtype=torch.float),
137
- "point_coords": torch.randint(low=0, high=1024, size=(1, 5, 2), dtype=torch.float),
138
- "point_labels": torch.randint(low=0, high=4, size=(1, 5), dtype=torch.float),
139
- "mask_input": torch.randn(1, 1, *mask_input_size, dtype=torch.float),
140
- "has_mask_input": torch.tensor([1], dtype=torch.float),
141
- "orig_im_size": torch.tensor([1500, 2250], dtype=torch.float),
142
- }
143
-
144
- _ = onnx_model(**dummy_inputs)
145
-
146
- output_names = ["masks", "iou_predictions", "low_res_masks"]
147
-
148
- with warnings.catch_warnings():
149
- warnings.filterwarnings("ignore", category=torch.jit.TracerWarning)
150
- warnings.filterwarnings("ignore", category=UserWarning)
151
- with open(output, "wb") as f:
152
- print(f"Exporing onnx model to {output}...")
153
- torch.onnx.export(
154
- onnx_model,
155
- tuple(dummy_inputs.values()),
156
- f,
157
- export_params=True,
158
- verbose=False,
159
- opset_version=opset,
160
- do_constant_folding=True,
161
- input_names=list(dummy_inputs.keys()),
162
- output_names=output_names,
163
- dynamic_axes=dynamic_axes,
164
- )
165
-
166
- if onnxruntime_exists:
167
- ort_inputs = {k: to_numpy(v) for k, v in dummy_inputs.items()}
168
- ort_session = onnxruntime.InferenceSession(output)
169
- _ = ort_session.run(None, ort_inputs)
170
- print("Model has successfully been run with ONNXRuntime.")
171
-
172
-
173
- def to_numpy(tensor):
174
- return tensor.cpu().numpy()
175
-
176
-
177
- if __name__ == "__main__":
178
- args = parser.parse_args()
179
- run_export(
180
- model_type=args.model_type,
181
- checkpoint=args.checkpoint,
182
- output=args.output,
183
- opset=args.opset,
184
- return_single_mask=args.return_single_mask,
185
- gelu_approximate=args.gelu_approximate,
186
- use_stability_score=args.use_stability_score,
187
- return_extra_metrics=args.return_extra_metrics,
188
- )
189
-
190
- if args.quantize_out is not None:
191
- assert onnxruntime_exists, "onnxruntime is required to quantize the model."
192
- from onnxruntime.quantization import QuantType # type: ignore
193
- from onnxruntime.quantization.quantize import quantize_dynamic # type: ignore
194
-
195
- print(f"Quantizing model and writing to {args.quantize_out}...")
196
- quantize_dynamic(
197
- model_input=args.output,
198
- model_output=args.quantize_out,
199
- optimize_model=True,
200
- per_channel=False,
201
- reduce_range=False,
202
- weight_type=QuantType.QUInt8,
203
- )
204
- print("Done!")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything.egg-info/PKG-INFO DELETED
@@ -1,13 +0,0 @@
1
- Metadata-Version: 2.1
2
- Name: segment-anything
3
- Version: 1.0
4
- Summary: UNKNOWN
5
- Home-page: UNKNOWN
6
- License: UNKNOWN
7
- Platform: UNKNOWN
8
- Provides-Extra: all
9
- Provides-Extra: dev
10
- License-File: LICENSE
11
-
12
- UNKNOWN
13
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything.egg-info/SOURCES.txt DELETED
@@ -1,24 +0,0 @@
1
- LICENSE
2
- README.md
3
- setup.cfg
4
- setup.py
5
- segment_anything/__init__.py
6
- segment_anything/automatic_mask_generator.py
7
- segment_anything/build_sam.py
8
- segment_anything/predictor.py
9
- segment_anything.egg-info/PKG-INFO
10
- segment_anything.egg-info/SOURCES.txt
11
- segment_anything.egg-info/dependency_links.txt
12
- segment_anything.egg-info/requires.txt
13
- segment_anything.egg-info/top_level.txt
14
- segment_anything/modeling/__init__.py
15
- segment_anything/modeling/common.py
16
- segment_anything/modeling/image_encoder.py
17
- segment_anything/modeling/mask_decoder.py
18
- segment_anything/modeling/prompt_encoder.py
19
- segment_anything/modeling/sam.py
20
- segment_anything/modeling/transformer.py
21
- segment_anything/utils/__init__.py
22
- segment_anything/utils/amg.py
23
- segment_anything/utils/onnx.py
24
- segment_anything/utils/transforms.py
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything.egg-info/dependency_links.txt DELETED
@@ -1 +0,0 @@
1
-
 
 
segment_anything/segment_anything.egg-info/requires.txt DELETED
@@ -1,13 +0,0 @@
1
-
2
- [all]
3
- matplotlib
4
- pycocotools
5
- opencv-python
6
- onnx
7
- onnxruntime
8
-
9
- [dev]
10
- flake8
11
- isort
12
- black
13
- mypy
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything.egg-info/top_level.txt DELETED
@@ -1 +0,0 @@
1
- segment_anything
 
 
segment_anything/segment_anything/__init__.py DELETED
@@ -1,15 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- from .build_sam import (
8
- build_sam,
9
- build_sam_vit_h,
10
- build_sam_vit_l,
11
- build_sam_vit_b,
12
- sam_model_registry,
13
- )
14
- from .predictor import SamPredictor
15
- from .automatic_mask_generator import SamAutomaticMaskGenerator
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/__pycache__/__init__.cpython-39.pyc DELETED
Binary file (472 Bytes)
 
segment_anything/segment_anything/__pycache__/automatic_mask_generator.cpython-39.pyc DELETED
Binary file (11.4 kB)
 
segment_anything/segment_anything/__pycache__/build_sam.cpython-39.pyc DELETED
Binary file (2.27 kB)
 
segment_anything/segment_anything/__pycache__/predictor.cpython-39.pyc DELETED
Binary file (9.95 kB)
 
segment_anything/segment_anything/automatic_mask_generator.py DELETED
@@ -1,372 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import numpy as np
8
- import torch
9
- from torchvision.ops.boxes import batched_nms, box_area # type: ignore
10
-
11
- from typing import Any, Dict, List, Optional, Tuple
12
-
13
- from .modeling import Sam
14
- from .predictor import SamPredictor
15
- from .utils.amg import (
16
- MaskData,
17
- area_from_rle,
18
- batch_iterator,
19
- batched_mask_to_box,
20
- box_xyxy_to_xywh,
21
- build_all_layer_point_grids,
22
- calculate_stability_score,
23
- coco_encode_rle,
24
- generate_crop_boxes,
25
- is_box_near_crop_edge,
26
- mask_to_rle_pytorch,
27
- remove_small_regions,
28
- rle_to_mask,
29
- uncrop_boxes_xyxy,
30
- uncrop_masks,
31
- uncrop_points,
32
- )
33
-
34
-
35
- class SamAutomaticMaskGenerator:
36
- def __init__(
37
- self,
38
- model: Sam,
39
- points_per_side: Optional[int] = 32,
40
- points_per_batch: int = 64,
41
- pred_iou_thresh: float = 0.88,
42
- stability_score_thresh: float = 0.95,
43
- stability_score_offset: float = 1.0,
44
- box_nms_thresh: float = 0.7,
45
- crop_n_layers: int = 0,
46
- crop_nms_thresh: float = 0.7,
47
- crop_overlap_ratio: float = 512 / 1500,
48
- crop_n_points_downscale_factor: int = 1,
49
- point_grids: Optional[List[np.ndarray]] = None,
50
- min_mask_region_area: int = 0,
51
- output_mode: str = "binary_mask",
52
- ) -> None:
53
- """
54
- Using a SAM model, generates masks for the entire image.
55
- Generates a grid of point prompts over the image, then filters
56
- low quality and duplicate masks. The default settings are chosen
57
- for SAM with a ViT-H backbone.
58
-
59
- Arguments:
60
- model (Sam): The SAM model to use for mask prediction.
61
- points_per_side (int or None): The number of points to be sampled
62
- along one side of the image. The total number of points is
63
- points_per_side**2. If None, 'point_grids' must provide explicit
64
- point sampling.
65
- points_per_batch (int): Sets the number of points run simultaneously
66
- by the model. Higher numbers may be faster but use more GPU memory.
67
- pred_iou_thresh (float): A filtering threshold in [0,1], using the
68
- model's predicted mask quality.
69
- stability_score_thresh (float): A filtering threshold in [0,1], using
70
- the stability of the mask under changes to the cutoff used to binarize
71
- the model's mask predictions.
72
- stability_score_offset (float): The amount to shift the cutoff when
73
- calculated the stability score.
74
- box_nms_thresh (float): The box IoU cutoff used by non-maximal
75
- suppression to filter duplicate masks.
76
- crops_n_layers (int): If >0, mask prediction will be run again on
77
- crops of the image. Sets the number of layers to run, where each
78
- layer has 2**i_layer number of image crops.
79
- crops_nms_thresh (float): The box IoU cutoff used by non-maximal
80
- suppression to filter duplicate masks between different crops.
81
- crop_overlap_ratio (float): Sets the degree to which crops overlap.
82
- In the first crop layer, crops will overlap by this fraction of
83
- the image length. Later layers with more crops scale down this overlap.
84
- crop_n_points_downscale_factor (int): The number of points-per-side
85
- sampled in layer n is scaled down by crop_n_points_downscale_factor**n.
86
- point_grids (list(np.ndarray) or None): A list over explicit grids
87
- of points used for sampling, normalized to [0,1]. The nth grid in the
88
- list is used in the nth crop layer. Exclusive with points_per_side.
89
- min_mask_region_area (int): If >0, postprocessing will be applied
90
- to remove disconnected regions and holes in masks with area smaller
91
- than min_mask_region_area. Requires opencv.
92
- output_mode (str): The form masks are returned in. Can be 'binary_mask',
93
- 'uncompressed_rle', or 'coco_rle'. 'coco_rle' requires pycocotools.
94
- For large resolutions, 'binary_mask' may consume large amounts of
95
- memory.
96
- """
97
-
98
- assert (points_per_side is None) != (
99
- point_grids is None
100
- ), "Exactly one of points_per_side or point_grid must be provided."
101
- if points_per_side is not None:
102
- self.point_grids = build_all_layer_point_grids(
103
- points_per_side,
104
- crop_n_layers,
105
- crop_n_points_downscale_factor,
106
- )
107
- elif point_grids is not None:
108
- self.point_grids = point_grids
109
- else:
110
- raise ValueError("Can't have both points_per_side and point_grid be None.")
111
-
112
- assert output_mode in [
113
- "binary_mask",
114
- "uncompressed_rle",
115
- "coco_rle",
116
- ], f"Unknown output_mode {output_mode}."
117
- if output_mode == "coco_rle":
118
- from pycocotools import mask as mask_utils # type: ignore # noqa: F401
119
-
120
- if min_mask_region_area > 0:
121
- import cv2 # type: ignore # noqa: F401
122
-
123
- self.predictor = SamPredictor(model)
124
- self.points_per_batch = points_per_batch
125
- self.pred_iou_thresh = pred_iou_thresh
126
- self.stability_score_thresh = stability_score_thresh
127
- self.stability_score_offset = stability_score_offset
128
- self.box_nms_thresh = box_nms_thresh
129
- self.crop_n_layers = crop_n_layers
130
- self.crop_nms_thresh = crop_nms_thresh
131
- self.crop_overlap_ratio = crop_overlap_ratio
132
- self.crop_n_points_downscale_factor = crop_n_points_downscale_factor
133
- self.min_mask_region_area = min_mask_region_area
134
- self.output_mode = output_mode
135
-
136
- @torch.no_grad()
137
- def generate(self, image: np.ndarray) -> List[Dict[str, Any]]:
138
- """
139
- Generates masks for the given image.
140
-
141
- Arguments:
142
- image (np.ndarray): The image to generate masks for, in HWC uint8 format.
143
-
144
- Returns:
145
- list(dict(str, any)): A list over records for masks. Each record is
146
- a dict containing the following keys:
147
- segmentation (dict(str, any) or np.ndarray): The mask. If
148
- output_mode='binary_mask', is an array of shape HW. Otherwise,
149
- is a dictionary containing the RLE.
150
- bbox (list(float)): The box around the mask, in XYWH format.
151
- area (int): The area in pixels of the mask.
152
- predicted_iou (float): The model's own prediction of the mask's
153
- quality. This is filtered by the pred_iou_thresh parameter.
154
- point_coords (list(list(float))): The point coordinates input
155
- to the model to generate this mask.
156
- stability_score (float): A measure of the mask's quality. This
157
- is filtered on using the stability_score_thresh parameter.
158
- crop_box (list(float)): The crop of the image used to generate
159
- the mask, given in XYWH format.
160
- """
161
-
162
- # Generate masks
163
- mask_data = self._generate_masks(image)
164
-
165
- # Filter small disconnected regions and holes in masks
166
- if self.min_mask_region_area > 0:
167
- mask_data = self.postprocess_small_regions(
168
- mask_data,
169
- self.min_mask_region_area,
170
- max(self.box_nms_thresh, self.crop_nms_thresh),
171
- )
172
-
173
- # Encode masks
174
- if self.output_mode == "coco_rle":
175
- mask_data["segmentations"] = [coco_encode_rle(rle) for rle in mask_data["rles"]]
176
- elif self.output_mode == "binary_mask":
177
- mask_data["segmentations"] = [rle_to_mask(rle) for rle in mask_data["rles"]]
178
- else:
179
- mask_data["segmentations"] = mask_data["rles"]
180
-
181
- # Write mask records
182
- curr_anns = []
183
- for idx in range(len(mask_data["segmentations"])):
184
- ann = {
185
- "segmentation": mask_data["segmentations"][idx],
186
- "area": area_from_rle(mask_data["rles"][idx]),
187
- "bbox": box_xyxy_to_xywh(mask_data["boxes"][idx]).tolist(),
188
- "predicted_iou": mask_data["iou_preds"][idx].item(),
189
- "point_coords": [mask_data["points"][idx].tolist()],
190
- "stability_score": mask_data["stability_score"][idx].item(),
191
- "crop_box": box_xyxy_to_xywh(mask_data["crop_boxes"][idx]).tolist(),
192
- }
193
- curr_anns.append(ann)
194
-
195
- return curr_anns
196
-
197
- def _generate_masks(self, image: np.ndarray) -> MaskData:
198
- orig_size = image.shape[:2]
199
- crop_boxes, layer_idxs = generate_crop_boxes(
200
- orig_size, self.crop_n_layers, self.crop_overlap_ratio
201
- )
202
-
203
- # Iterate over image crops
204
- data = MaskData()
205
- for crop_box, layer_idx in zip(crop_boxes, layer_idxs):
206
- crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)
207
- data.cat(crop_data)
208
-
209
- # Remove duplicate masks between crops
210
- if len(crop_boxes) > 1:
211
- # Prefer masks from smaller crops
212
- scores = 1 / box_area(data["crop_boxes"])
213
- scores = scores.to(data["boxes"].device)
214
- keep_by_nms = batched_nms(
215
- data["boxes"].float(),
216
- scores,
217
- torch.zeros(len(data["boxes"])), # categories
218
- iou_threshold=self.crop_nms_thresh,
219
- )
220
- data.filter(keep_by_nms)
221
-
222
- data.to_numpy()
223
- return data
224
-
225
- def _process_crop(
226
- self,
227
- image: np.ndarray,
228
- crop_box: List[int],
229
- crop_layer_idx: int,
230
- orig_size: Tuple[int, ...],
231
- ) -> MaskData:
232
- # Crop the image and calculate embeddings
233
- x0, y0, x1, y1 = crop_box
234
- cropped_im = image[y0:y1, x0:x1, :]
235
- cropped_im_size = cropped_im.shape[:2]
236
- self.predictor.set_image(cropped_im)
237
-
238
- # Get points for this crop
239
- points_scale = np.array(cropped_im_size)[None, ::-1]
240
- points_for_image = self.point_grids[crop_layer_idx] * points_scale
241
-
242
- # Generate masks for this crop in batches
243
- data = MaskData()
244
- for (points,) in batch_iterator(self.points_per_batch, points_for_image):
245
- batch_data = self._process_batch(points, cropped_im_size, crop_box, orig_size)
246
- data.cat(batch_data)
247
- del batch_data
248
- self.predictor.reset_image()
249
-
250
- # Remove duplicates within this crop.
251
- keep_by_nms = batched_nms(
252
- data["boxes"].float(),
253
- data["iou_preds"],
254
- torch.zeros(len(data["boxes"])), # categories
255
- iou_threshold=self.box_nms_thresh,
256
- )
257
- data.filter(keep_by_nms)
258
-
259
- # Return to the original image frame
260
- data["boxes"] = uncrop_boxes_xyxy(data["boxes"], crop_box)
261
- data["points"] = uncrop_points(data["points"], crop_box)
262
- data["crop_boxes"] = torch.tensor([crop_box for _ in range(len(data["rles"]))])
263
-
264
- return data
265
-
266
- def _process_batch(
267
- self,
268
- points: np.ndarray,
269
- im_size: Tuple[int, ...],
270
- crop_box: List[int],
271
- orig_size: Tuple[int, ...],
272
- ) -> MaskData:
273
- orig_h, orig_w = orig_size
274
-
275
- # Run model on this batch
276
- transformed_points = self.predictor.transform.apply_coords(points, im_size)
277
- in_points = torch.as_tensor(transformed_points, device=self.predictor.device)
278
- in_labels = torch.ones(in_points.shape[0], dtype=torch.int, device=in_points.device)
279
- masks, iou_preds, _ = self.predictor.predict_torch(
280
- in_points[:, None, :],
281
- in_labels[:, None],
282
- multimask_output=True,
283
- return_logits=True,
284
- )
285
-
286
- # Serialize predictions and store in MaskData
287
- data = MaskData(
288
- masks=masks.flatten(0, 1),
289
- iou_preds=iou_preds.flatten(0, 1),
290
- points=torch.as_tensor(points.repeat(masks.shape[1], axis=0)),
291
- )
292
- del masks
293
-
294
- # Filter by predicted IoU
295
- if self.pred_iou_thresh > 0.0:
296
- keep_mask = data["iou_preds"] > self.pred_iou_thresh
297
- data.filter(keep_mask)
298
-
299
- # Calculate stability score
300
- data["stability_score"] = calculate_stability_score(
301
- data["masks"], self.predictor.model.mask_threshold, self.stability_score_offset
302
- )
303
- if self.stability_score_thresh > 0.0:
304
- keep_mask = data["stability_score"] >= self.stability_score_thresh
305
- data.filter(keep_mask)
306
-
307
- # Threshold masks and calculate boxes
308
- data["masks"] = data["masks"] > self.predictor.model.mask_threshold
309
- data["boxes"] = batched_mask_to_box(data["masks"])
310
-
311
- # Filter boxes that touch crop boundaries
312
- keep_mask = ~is_box_near_crop_edge(data["boxes"], crop_box, [0, 0, orig_w, orig_h])
313
- if not torch.all(keep_mask):
314
- data.filter(keep_mask)
315
-
316
- # Compress to RLE
317
- data["masks"] = uncrop_masks(data["masks"], crop_box, orig_h, orig_w)
318
- data["rles"] = mask_to_rle_pytorch(data["masks"])
319
- del data["masks"]
320
-
321
- return data
322
-
323
- @staticmethod
324
- def postprocess_small_regions(
325
- mask_data: MaskData, min_area: int, nms_thresh: float
326
- ) -> MaskData:
327
- """
328
- Removes small disconnected regions and holes in masks, then reruns
329
- box NMS to remove any new duplicates.
330
-
331
- Edits mask_data in place.
332
-
333
- Requires open-cv as a dependency.
334
- """
335
- if len(mask_data["rles"]) == 0:
336
- return mask_data
337
-
338
- # Filter small disconnected regions and holes
339
- new_masks = []
340
- scores = []
341
- for rle in mask_data["rles"]:
342
- mask = rle_to_mask(rle)
343
-
344
- mask, changed = remove_small_regions(mask, min_area, mode="holes")
345
- unchanged = not changed
346
- mask, changed = remove_small_regions(mask, min_area, mode="islands")
347
- unchanged = unchanged and not changed
348
-
349
- new_masks.append(torch.as_tensor(mask).unsqueeze(0))
350
- # Give score=0 to changed masks and score=1 to unchanged masks
351
- # so NMS will prefer ones that didn't need postprocessing
352
- scores.append(float(unchanged))
353
-
354
- # Recalculate boxes and remove any new duplicates
355
- masks = torch.cat(new_masks, dim=0)
356
- boxes = batched_mask_to_box(masks)
357
- keep_by_nms = batched_nms(
358
- boxes.float(),
359
- torch.as_tensor(scores),
360
- torch.zeros(len(boxes)), # categories
361
- iou_threshold=nms_thresh,
362
- )
363
-
364
- # Only recalculate RLEs for masks that have changed
365
- for i_mask in keep_by_nms:
366
- if scores[i_mask] == 0.0:
367
- mask_torch = masks[i_mask].unsqueeze(0)
368
- mask_data["rles"][i_mask] = mask_to_rle_pytorch(mask_torch)[0]
369
- mask_data["boxes"][i_mask] = boxes[i_mask] # update res directly
370
- mask_data.filter(keep_by_nms)
371
-
372
- return mask_data
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/build_sam.py DELETED
@@ -1,107 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import torch
8
-
9
- from functools import partial
10
-
11
- from .modeling import ImageEncoderViT, MaskDecoder, PromptEncoder, Sam, TwoWayTransformer
12
-
13
-
14
- def build_sam_vit_h(checkpoint=None):
15
- return _build_sam(
16
- encoder_embed_dim=1280,
17
- encoder_depth=32,
18
- encoder_num_heads=16,
19
- encoder_global_attn_indexes=[7, 15, 23, 31],
20
- checkpoint=checkpoint,
21
- )
22
-
23
-
24
- build_sam = build_sam_vit_h
25
-
26
-
27
- def build_sam_vit_l(checkpoint=None):
28
- return _build_sam(
29
- encoder_embed_dim=1024,
30
- encoder_depth=24,
31
- encoder_num_heads=16,
32
- encoder_global_attn_indexes=[5, 11, 17, 23],
33
- checkpoint=checkpoint,
34
- )
35
-
36
-
37
- def build_sam_vit_b(checkpoint=None):
38
- return _build_sam(
39
- encoder_embed_dim=768,
40
- encoder_depth=12,
41
- encoder_num_heads=12,
42
- encoder_global_attn_indexes=[2, 5, 8, 11],
43
- checkpoint=checkpoint,
44
- )
45
-
46
-
47
- sam_model_registry = {
48
- "default": build_sam,
49
- "vit_h": build_sam,
50
- "vit_l": build_sam_vit_l,
51
- "vit_b": build_sam_vit_b,
52
- }
53
-
54
-
55
- def _build_sam(
56
- encoder_embed_dim,
57
- encoder_depth,
58
- encoder_num_heads,
59
- encoder_global_attn_indexes,
60
- checkpoint=None,
61
- ):
62
- prompt_embed_dim = 256
63
- image_size = 1024
64
- vit_patch_size = 16
65
- image_embedding_size = image_size // vit_patch_size
66
- sam = Sam(
67
- image_encoder=ImageEncoderViT(
68
- depth=encoder_depth,
69
- embed_dim=encoder_embed_dim,
70
- img_size=image_size,
71
- mlp_ratio=4,
72
- norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
73
- num_heads=encoder_num_heads,
74
- patch_size=vit_patch_size,
75
- qkv_bias=True,
76
- use_rel_pos=True,
77
- global_attn_indexes=encoder_global_attn_indexes,
78
- window_size=14,
79
- out_chans=prompt_embed_dim,
80
- ),
81
- prompt_encoder=PromptEncoder(
82
- embed_dim=prompt_embed_dim,
83
- image_embedding_size=(image_embedding_size, image_embedding_size),
84
- input_image_size=(image_size, image_size),
85
- mask_in_chans=16,
86
- ),
87
- mask_decoder=MaskDecoder(
88
- num_multimask_outputs=3,
89
- transformer=TwoWayTransformer(
90
- depth=2,
91
- embedding_dim=prompt_embed_dim,
92
- mlp_dim=2048,
93
- num_heads=8,
94
- ),
95
- transformer_dim=prompt_embed_dim,
96
- iou_head_depth=3,
97
- iou_head_hidden_dim=256,
98
- ),
99
- pixel_mean=[123.675, 116.28, 103.53],
100
- pixel_std=[58.395, 57.12, 57.375],
101
- )
102
- sam.eval()
103
- if checkpoint is not None:
104
- with open(checkpoint, "rb") as f:
105
- state_dict = torch.load(f)
106
- sam.load_state_dict(state_dict)
107
- return sam
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/modeling/__init__.py DELETED
@@ -1,11 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- from .sam import Sam
8
- from .image_encoder import ImageEncoderViT
9
- from .mask_decoder import MaskDecoder
10
- from .prompt_encoder import PromptEncoder
11
- from .transformer import TwoWayTransformer
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/modeling/__pycache__/__init__.cpython-39.pyc DELETED
Binary file (459 Bytes)
 
segment_anything/segment_anything/modeling/__pycache__/common.cpython-39.pyc DELETED
Binary file (1.8 kB)
 
segment_anything/segment_anything/modeling/__pycache__/image_encoder.cpython-39.pyc DELETED
Binary file (12.5 kB)
 
segment_anything/segment_anything/modeling/__pycache__/mask_decoder.cpython-39.pyc DELETED
Binary file (5.47 kB)
 
segment_anything/segment_anything/modeling/__pycache__/prompt_encoder.cpython-39.pyc DELETED
Binary file (7.72 kB)
 
segment_anything/segment_anything/modeling/__pycache__/sam.cpython-39.pyc DELETED
Binary file (6.7 kB)
 
segment_anything/segment_anything/modeling/__pycache__/transformer.cpython-39.pyc DELETED
Binary file (6.6 kB)
 
segment_anything/segment_anything/modeling/common.py DELETED
@@ -1,43 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import torch
8
- import torch.nn as nn
9
-
10
- from typing import Type
11
-
12
-
13
- class MLPBlock(nn.Module):
14
- def __init__(
15
- self,
16
- embedding_dim: int,
17
- mlp_dim: int,
18
- act: Type[nn.Module] = nn.GELU,
19
- ) -> None:
20
- super().__init__()
21
- self.lin1 = nn.Linear(embedding_dim, mlp_dim)
22
- self.lin2 = nn.Linear(mlp_dim, embedding_dim)
23
- self.act = act()
24
-
25
- def forward(self, x: torch.Tensor) -> torch.Tensor:
26
- return self.lin2(self.act(self.lin1(x)))
27
-
28
-
29
- # From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
30
- # Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa
31
- class LayerNorm2d(nn.Module):
32
- def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
33
- super().__init__()
34
- self.weight = nn.Parameter(torch.ones(num_channels))
35
- self.bias = nn.Parameter(torch.zeros(num_channels))
36
- self.eps = eps
37
-
38
- def forward(self, x: torch.Tensor) -> torch.Tensor:
39
- u = x.mean(1, keepdim=True)
40
- s = (x - u).pow(2).mean(1, keepdim=True)
41
- x = (x - u) / torch.sqrt(s + self.eps)
42
- x = self.weight[:, None, None] * x + self.bias[:, None, None]
43
- return x
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/modeling/image_encoder.py DELETED
@@ -1,395 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import torch
8
- import torch.nn as nn
9
- import torch.nn.functional as F
10
-
11
- from typing import Optional, Tuple, Type
12
-
13
- from .common import LayerNorm2d, MLPBlock
14
-
15
-
16
- # This class and its supporting functions below lightly adapted from the ViTDet backbone available at: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py # noqa
17
- class ImageEncoderViT(nn.Module):
18
- def __init__(
19
- self,
20
- img_size: int = 1024,
21
- patch_size: int = 16,
22
- in_chans: int = 3,
23
- embed_dim: int = 768,
24
- depth: int = 12,
25
- num_heads: int = 12,
26
- mlp_ratio: float = 4.0,
27
- out_chans: int = 256,
28
- qkv_bias: bool = True,
29
- norm_layer: Type[nn.Module] = nn.LayerNorm,
30
- act_layer: Type[nn.Module] = nn.GELU,
31
- use_abs_pos: bool = True,
32
- use_rel_pos: bool = False,
33
- rel_pos_zero_init: bool = True,
34
- window_size: int = 0,
35
- global_attn_indexes: Tuple[int, ...] = (),
36
- ) -> None:
37
- """
38
- Args:
39
- img_size (int): Input image size.
40
- patch_size (int): Patch size.
41
- in_chans (int): Number of input image channels.
42
- embed_dim (int): Patch embedding dimension.
43
- depth (int): Depth of ViT.
44
- num_heads (int): Number of attention heads in each ViT block.
45
- mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
46
- qkv_bias (bool): If True, add a learnable bias to query, key, value.
47
- norm_layer (nn.Module): Normalization layer.
48
- act_layer (nn.Module): Activation layer.
49
- use_abs_pos (bool): If True, use absolute positional embeddings.
50
- use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
51
- rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
52
- window_size (int): Window size for window attention blocks.
53
- global_attn_indexes (list): Indexes for blocks using global attention.
54
- """
55
- super().__init__()
56
- self.img_size = img_size
57
-
58
- self.patch_embed = PatchEmbed(
59
- kernel_size=(patch_size, patch_size),
60
- stride=(patch_size, patch_size),
61
- in_chans=in_chans,
62
- embed_dim=embed_dim,
63
- )
64
-
65
- self.pos_embed: Optional[nn.Parameter] = None
66
- if use_abs_pos:
67
- # Initialize absolute positional embedding with pretrain image size.
68
- self.pos_embed = nn.Parameter(
69
- torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim)
70
- )
71
-
72
- self.blocks = nn.ModuleList()
73
- for i in range(depth):
74
- block = Block(
75
- dim=embed_dim,
76
- num_heads=num_heads,
77
- mlp_ratio=mlp_ratio,
78
- qkv_bias=qkv_bias,
79
- norm_layer=norm_layer,
80
- act_layer=act_layer,
81
- use_rel_pos=use_rel_pos,
82
- rel_pos_zero_init=rel_pos_zero_init,
83
- window_size=window_size if i not in global_attn_indexes else 0,
84
- input_size=(img_size // patch_size, img_size // patch_size),
85
- )
86
- self.blocks.append(block)
87
-
88
- self.neck = nn.Sequential(
89
- nn.Conv2d(
90
- embed_dim,
91
- out_chans,
92
- kernel_size=1,
93
- bias=False,
94
- ),
95
- LayerNorm2d(out_chans),
96
- nn.Conv2d(
97
- out_chans,
98
- out_chans,
99
- kernel_size=3,
100
- padding=1,
101
- bias=False,
102
- ),
103
- LayerNorm2d(out_chans),
104
- )
105
-
106
- def forward(self, x: torch.Tensor) -> torch.Tensor:
107
- x = self.patch_embed(x)
108
- if self.pos_embed is not None:
109
- x = x + self.pos_embed
110
-
111
- for blk in self.blocks:
112
- x = blk(x)
113
-
114
- x = self.neck(x.permute(0, 3, 1, 2))
115
-
116
- return x
117
-
118
-
119
- class Block(nn.Module):
120
- """Transformer blocks with support of window attention and residual propagation blocks"""
121
-
122
- def __init__(
123
- self,
124
- dim: int,
125
- num_heads: int,
126
- mlp_ratio: float = 4.0,
127
- qkv_bias: bool = True,
128
- norm_layer: Type[nn.Module] = nn.LayerNorm,
129
- act_layer: Type[nn.Module] = nn.GELU,
130
- use_rel_pos: bool = False,
131
- rel_pos_zero_init: bool = True,
132
- window_size: int = 0,
133
- input_size: Optional[Tuple[int, int]] = None,
134
- ) -> None:
135
- """
136
- Args:
137
- dim (int): Number of input channels.
138
- num_heads (int): Number of attention heads in each ViT block.
139
- mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
140
- qkv_bias (bool): If True, add a learnable bias to query, key, value.
141
- norm_layer (nn.Module): Normalization layer.
142
- act_layer (nn.Module): Activation layer.
143
- use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
144
- rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
145
- window_size (int): Window size for window attention blocks. If it equals 0, then
146
- use global attention.
147
- input_size (int or None): Input resolution for calculating the relative positional
148
- parameter size.
149
- """
150
- super().__init__()
151
- self.norm1 = norm_layer(dim)
152
- self.attn = Attention(
153
- dim,
154
- num_heads=num_heads,
155
- qkv_bias=qkv_bias,
156
- use_rel_pos=use_rel_pos,
157
- rel_pos_zero_init=rel_pos_zero_init,
158
- input_size=input_size if window_size == 0 else (window_size, window_size),
159
- )
160
-
161
- self.norm2 = norm_layer(dim)
162
- self.mlp = MLPBlock(embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer)
163
-
164
- self.window_size = window_size
165
-
166
- def forward(self, x: torch.Tensor) -> torch.Tensor:
167
- shortcut = x
168
- x = self.norm1(x)
169
- # Window partition
170
- if self.window_size > 0:
171
- H, W = x.shape[1], x.shape[2]
172
- x, pad_hw = window_partition(x, self.window_size)
173
-
174
- x = self.attn(x)
175
- # Reverse window partition
176
- if self.window_size > 0:
177
- x = window_unpartition(x, self.window_size, pad_hw, (H, W))
178
-
179
- x = shortcut + x
180
- x = x + self.mlp(self.norm2(x))
181
-
182
- return x
183
-
184
-
185
- class Attention(nn.Module):
186
- """Multi-head Attention block with relative position embeddings."""
187
-
188
- def __init__(
189
- self,
190
- dim: int,
191
- num_heads: int = 8,
192
- qkv_bias: bool = True,
193
- use_rel_pos: bool = False,
194
- rel_pos_zero_init: bool = True,
195
- input_size: Optional[Tuple[int, int]] = None,
196
- ) -> None:
197
- """
198
- Args:
199
- dim (int): Number of input channels.
200
- num_heads (int): Number of attention heads.
201
- qkv_bias (bool: If True, add a learnable bias to query, key, value.
202
- rel_pos (bool): If True, add relative positional embeddings to the attention map.
203
- rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
204
- input_size (int or None): Input resolution for calculating the relative positional
205
- parameter size.
206
- """
207
- super().__init__()
208
- self.num_heads = num_heads
209
- head_dim = dim // num_heads
210
- self.scale = head_dim**-0.5
211
-
212
- self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
213
- self.proj = nn.Linear(dim, dim)
214
-
215
- self.use_rel_pos = use_rel_pos
216
- if self.use_rel_pos:
217
- assert (
218
- input_size is not None
219
- ), "Input size must be provided if using relative positional encoding."
220
- # initialize relative positional embeddings
221
- self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))
222
- self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))
223
-
224
- def forward(self, x: torch.Tensor) -> torch.Tensor:
225
- B, H, W, _ = x.shape
226
- # qkv with shape (3, B, nHead, H * W, C)
227
- qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
228
- # q, k, v with shape (B * nHead, H * W, C)
229
- q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)
230
-
231
- attn = (q * self.scale) @ k.transpose(-2, -1)
232
-
233
- if self.use_rel_pos:
234
- attn = add_decomposed_rel_pos(attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W))
235
-
236
- attn = attn.softmax(dim=-1)
237
- x = (attn @ v).view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1)
238
- x = self.proj(x)
239
-
240
- return x
241
-
242
-
243
- def window_partition(x: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]:
244
- """
245
- Partition into non-overlapping windows with padding if needed.
246
- Args:
247
- x (tensor): input tokens with [B, H, W, C].
248
- window_size (int): window size.
249
-
250
- Returns:
251
- windows: windows after partition with [B * num_windows, window_size, window_size, C].
252
- (Hp, Wp): padded height and width before partition
253
- """
254
- B, H, W, C = x.shape
255
-
256
- pad_h = (window_size - H % window_size) % window_size
257
- pad_w = (window_size - W % window_size) % window_size
258
- if pad_h > 0 or pad_w > 0:
259
- x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
260
- Hp, Wp = H + pad_h, W + pad_w
261
-
262
- x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)
263
- windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
264
- return windows, (Hp, Wp)
265
-
266
-
267
- def window_unpartition(
268
- windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int]
269
- ) -> torch.Tensor:
270
- """
271
- Window unpartition into original sequences and removing padding.
272
- Args:
273
- x (tensor): input tokens with [B * num_windows, window_size, window_size, C].
274
- window_size (int): window size.
275
- pad_hw (Tuple): padded height and width (Hp, Wp).
276
- hw (Tuple): original height and width (H, W) before padding.
277
-
278
- Returns:
279
- x: unpartitioned sequences with [B, H, W, C].
280
- """
281
- Hp, Wp = pad_hw
282
- H, W = hw
283
- B = windows.shape[0] // (Hp * Wp // window_size // window_size)
284
- x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)
285
- x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1)
286
-
287
- if Hp > H or Wp > W:
288
- x = x[:, :H, :W, :].contiguous()
289
- return x
290
-
291
-
292
- def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
293
- """
294
- Get relative positional embeddings according to the relative positions of
295
- query and key sizes.
296
- Args:
297
- q_size (int): size of query q.
298
- k_size (int): size of key k.
299
- rel_pos (Tensor): relative position embeddings (L, C).
300
-
301
- Returns:
302
- Extracted positional embeddings according to relative positions.
303
- """
304
- max_rel_dist = int(2 * max(q_size, k_size) - 1)
305
- # Interpolate rel pos if needed.
306
- if rel_pos.shape[0] != max_rel_dist:
307
- # Interpolate rel pos.
308
- rel_pos_resized = F.interpolate(
309
- rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
310
- size=max_rel_dist,
311
- mode="linear",
312
- )
313
- rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
314
- else:
315
- rel_pos_resized = rel_pos
316
-
317
- # Scale the coords with short length if shapes for q and k are different.
318
- q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
319
- k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
320
- relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
321
-
322
- return rel_pos_resized[relative_coords.long()]
323
-
324
-
325
- def add_decomposed_rel_pos(
326
- attn: torch.Tensor,
327
- q: torch.Tensor,
328
- rel_pos_h: torch.Tensor,
329
- rel_pos_w: torch.Tensor,
330
- q_size: Tuple[int, int],
331
- k_size: Tuple[int, int],
332
- ) -> torch.Tensor:
333
- """
334
- Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
335
- https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py # noqa B950
336
- Args:
337
- attn (Tensor): attention map.
338
- q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).
339
- rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.
340
- rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.
341
- q_size (Tuple): spatial sequence size of query q with (q_h, q_w).
342
- k_size (Tuple): spatial sequence size of key k with (k_h, k_w).
343
-
344
- Returns:
345
- attn (Tensor): attention map with added relative positional embeddings.
346
- """
347
- q_h, q_w = q_size
348
- k_h, k_w = k_size
349
- Rh = get_rel_pos(q_h, k_h, rel_pos_h)
350
- Rw = get_rel_pos(q_w, k_w, rel_pos_w)
351
-
352
- B, _, dim = q.shape
353
- r_q = q.reshape(B, q_h, q_w, dim)
354
- rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh)
355
- rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw)
356
-
357
- attn = (
358
- attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]
359
- ).view(B, q_h * q_w, k_h * k_w)
360
-
361
- return attn
362
-
363
-
364
- class PatchEmbed(nn.Module):
365
- """
366
- Image to Patch Embedding.
367
- """
368
-
369
- def __init__(
370
- self,
371
- kernel_size: Tuple[int, int] = (16, 16),
372
- stride: Tuple[int, int] = (16, 16),
373
- padding: Tuple[int, int] = (0, 0),
374
- in_chans: int = 3,
375
- embed_dim: int = 768,
376
- ) -> None:
377
- """
378
- Args:
379
- kernel_size (Tuple): kernel size of the projection layer.
380
- stride (Tuple): stride of the projection layer.
381
- padding (Tuple): padding size of the projection layer.
382
- in_chans (int): Number of input image channels.
383
- embed_dim (int): embed_dim (int): Patch embedding dimension.
384
- """
385
- super().__init__()
386
-
387
- self.proj = nn.Conv2d(
388
- in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding
389
- )
390
-
391
- def forward(self, x: torch.Tensor) -> torch.Tensor:
392
- x = self.proj(x)
393
- # B C H W -> B H W C
394
- x = x.permute(0, 2, 3, 1)
395
- return x
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/modeling/mask_decoder.py DELETED
@@ -1,176 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import torch
8
- from torch import nn
9
- from torch.nn import functional as F
10
-
11
- from typing import List, Tuple, Type
12
-
13
- from .common import LayerNorm2d
14
-
15
-
16
- class MaskDecoder(nn.Module):
17
- def __init__(
18
- self,
19
- *,
20
- transformer_dim: int,
21
- transformer: nn.Module,
22
- num_multimask_outputs: int = 3,
23
- activation: Type[nn.Module] = nn.GELU,
24
- iou_head_depth: int = 3,
25
- iou_head_hidden_dim: int = 256,
26
- ) -> None:
27
- """
28
- Predicts masks given an image and prompt embeddings, using a
29
- tranformer architecture.
30
-
31
- Arguments:
32
- transformer_dim (int): the channel dimension of the transformer
33
- transformer (nn.Module): the transformer used to predict masks
34
- num_multimask_outputs (int): the number of masks to predict
35
- when disambiguating masks
36
- activation (nn.Module): the type of activation to use when
37
- upscaling masks
38
- iou_head_depth (int): the depth of the MLP used to predict
39
- mask quality
40
- iou_head_hidden_dim (int): the hidden dimension of the MLP
41
- used to predict mask quality
42
- """
43
- super().__init__()
44
- self.transformer_dim = transformer_dim
45
- self.transformer = transformer
46
-
47
- self.num_multimask_outputs = num_multimask_outputs
48
-
49
- self.iou_token = nn.Embedding(1, transformer_dim)
50
- self.num_mask_tokens = num_multimask_outputs + 1
51
- self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)
52
-
53
- self.output_upscaling = nn.Sequential(
54
- nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2),
55
- LayerNorm2d(transformer_dim // 4),
56
- activation(),
57
- nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
58
- activation(),
59
- )
60
- self.output_hypernetworks_mlps = nn.ModuleList(
61
- [
62
- MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3)
63
- for i in range(self.num_mask_tokens)
64
- ]
65
- )
66
-
67
- self.iou_prediction_head = MLP(
68
- transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth
69
- )
70
-
71
- def forward(
72
- self,
73
- image_embeddings: torch.Tensor,
74
- image_pe: torch.Tensor,
75
- sparse_prompt_embeddings: torch.Tensor,
76
- dense_prompt_embeddings: torch.Tensor,
77
- multimask_output: bool,
78
- ) -> Tuple[torch.Tensor, torch.Tensor]:
79
- """
80
- Predict masks given image and prompt embeddings.
81
-
82
- Arguments:
83
- image_embeddings (torch.Tensor): the embeddings from the image encoder
84
- image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
85
- sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
86
- dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
87
- multimask_output (bool): Whether to return multiple masks or a single
88
- mask.
89
-
90
- Returns:
91
- torch.Tensor: batched predicted masks
92
- torch.Tensor: batched predictions of mask quality
93
- """
94
- masks, iou_pred = self.predict_masks(
95
- image_embeddings=image_embeddings,
96
- image_pe=image_pe,
97
- sparse_prompt_embeddings=sparse_prompt_embeddings,
98
- dense_prompt_embeddings=dense_prompt_embeddings,
99
- )
100
-
101
- # Select the correct mask or masks for outptu
102
- if multimask_output:
103
- mask_slice = slice(1, None)
104
- else:
105
- mask_slice = slice(0, 1)
106
- masks = masks[:, mask_slice, :, :]
107
- iou_pred = iou_pred[:, mask_slice]
108
-
109
- # Prepare output
110
- return masks, iou_pred
111
-
112
- def predict_masks(
113
- self,
114
- image_embeddings: torch.Tensor,
115
- image_pe: torch.Tensor,
116
- sparse_prompt_embeddings: torch.Tensor,
117
- dense_prompt_embeddings: torch.Tensor,
118
- ) -> Tuple[torch.Tensor, torch.Tensor]:
119
- """Predicts masks. See 'forward' for more details."""
120
- # Concatenate output tokens
121
- output_tokens = torch.cat([self.iou_token.weight, self.mask_tokens.weight], dim=0)
122
- output_tokens = output_tokens.unsqueeze(0).expand(sparse_prompt_embeddings.size(0), -1, -1)
123
- tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)
124
-
125
- # Expand per-image data in batch direction to be per-mask
126
- src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)
127
- src = src + dense_prompt_embeddings
128
- pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
129
- b, c, h, w = src.shape
130
-
131
- # Run the transformer
132
- hs, src = self.transformer(src, pos_src, tokens)
133
- iou_token_out = hs[:, 0, :]
134
- mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :]
135
-
136
- # Upscale mask embeddings and predict masks using the mask tokens
137
- src = src.transpose(1, 2).view(b, c, h, w)
138
- upscaled_embedding = self.output_upscaling(src)
139
- hyper_in_list: List[torch.Tensor] = []
140
- for i in range(self.num_mask_tokens):
141
- hyper_in_list.append(self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :]))
142
- hyper_in = torch.stack(hyper_in_list, dim=1)
143
- b, c, h, w = upscaled_embedding.shape
144
- masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w)
145
-
146
- # Generate mask quality predictions
147
- iou_pred = self.iou_prediction_head(iou_token_out)
148
-
149
- return masks, iou_pred
150
-
151
-
152
- # Lightly adapted from
153
- # https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa
154
- class MLP(nn.Module):
155
- def __init__(
156
- self,
157
- input_dim: int,
158
- hidden_dim: int,
159
- output_dim: int,
160
- num_layers: int,
161
- sigmoid_output: bool = False,
162
- ) -> None:
163
- super().__init__()
164
- self.num_layers = num_layers
165
- h = [hidden_dim] * (num_layers - 1)
166
- self.layers = nn.ModuleList(
167
- nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
168
- )
169
- self.sigmoid_output = sigmoid_output
170
-
171
- def forward(self, x):
172
- for i, layer in enumerate(self.layers):
173
- x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
174
- if self.sigmoid_output:
175
- x = F.sigmoid(x)
176
- return x
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/modeling/prompt_encoder.py DELETED
@@ -1,214 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import numpy as np
8
- import torch
9
- from torch import nn
10
-
11
- from typing import Any, Optional, Tuple, Type
12
-
13
- from .common import LayerNorm2d
14
-
15
-
16
- class PromptEncoder(nn.Module):
17
- def __init__(
18
- self,
19
- embed_dim: int,
20
- image_embedding_size: Tuple[int, int],
21
- input_image_size: Tuple[int, int],
22
- mask_in_chans: int,
23
- activation: Type[nn.Module] = nn.GELU,
24
- ) -> None:
25
- """
26
- Encodes prompts for input to SAM's mask decoder.
27
-
28
- Arguments:
29
- embed_dim (int): The prompts' embedding dimension
30
- image_embedding_size (tuple(int, int)): The spatial size of the
31
- image embedding, as (H, W).
32
- input_image_size (int): The padded size of the image as input
33
- to the image encoder, as (H, W).
34
- mask_in_chans (int): The number of hidden channels used for
35
- encoding input masks.
36
- activation (nn.Module): The activation to use when encoding
37
- input masks.
38
- """
39
- super().__init__()
40
- self.embed_dim = embed_dim
41
- self.input_image_size = input_image_size
42
- self.image_embedding_size = image_embedding_size
43
- self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)
44
-
45
- self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners
46
- point_embeddings = [nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)]
47
- self.point_embeddings = nn.ModuleList(point_embeddings)
48
- self.not_a_point_embed = nn.Embedding(1, embed_dim)
49
-
50
- self.mask_input_size = (4 * image_embedding_size[0], 4 * image_embedding_size[1])
51
- self.mask_downscaling = nn.Sequential(
52
- nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),
53
- LayerNorm2d(mask_in_chans // 4),
54
- activation(),
55
- nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),
56
- LayerNorm2d(mask_in_chans),
57
- activation(),
58
- nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1),
59
- )
60
- self.no_mask_embed = nn.Embedding(1, embed_dim)
61
-
62
- def get_dense_pe(self) -> torch.Tensor:
63
- """
64
- Returns the positional encoding used to encode point prompts,
65
- applied to a dense set of points the shape of the image encoding.
66
-
67
- Returns:
68
- torch.Tensor: Positional encoding with shape
69
- 1x(embed_dim)x(embedding_h)x(embedding_w)
70
- """
71
- return self.pe_layer(self.image_embedding_size).unsqueeze(0)
72
-
73
- def _embed_points(
74
- self,
75
- points: torch.Tensor,
76
- labels: torch.Tensor,
77
- pad: bool,
78
- ) -> torch.Tensor:
79
- """Embeds point prompts."""
80
- points = points + 0.5 # Shift to center of pixel
81
- if pad:
82
- padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)
83
- padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
84
- points = torch.cat([points, padding_point], dim=1)
85
- labels = torch.cat([labels, padding_label], dim=1)
86
- point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size)
87
- point_embedding[labels == -1] = 0.0
88
- point_embedding[labels == -1] += self.not_a_point_embed.weight
89
- point_embedding[labels == 0] += self.point_embeddings[0].weight
90
- point_embedding[labels == 1] += self.point_embeddings[1].weight
91
- return point_embedding
92
-
93
- def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
94
- """Embeds box prompts."""
95
- boxes = boxes + 0.5 # Shift to center of pixel
96
- coords = boxes.reshape(-1, 2, 2)
97
- corner_embedding = self.pe_layer.forward_with_coords(coords, self.input_image_size)
98
- corner_embedding[:, 0, :] += self.point_embeddings[2].weight
99
- corner_embedding[:, 1, :] += self.point_embeddings[3].weight
100
- return corner_embedding
101
-
102
- def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:
103
- """Embeds mask inputs."""
104
- mask_embedding = self.mask_downscaling(masks)
105
- return mask_embedding
106
-
107
- def _get_batch_size(
108
- self,
109
- points: Optional[Tuple[torch.Tensor, torch.Tensor]],
110
- boxes: Optional[torch.Tensor],
111
- masks: Optional[torch.Tensor],
112
- ) -> int:
113
- """
114
- Gets the batch size of the output given the batch size of the input prompts.
115
- """
116
- if points is not None:
117
- return points[0].shape[0]
118
- elif boxes is not None:
119
- return boxes.shape[0]
120
- elif masks is not None:
121
- return masks.shape[0]
122
- else:
123
- return 1
124
-
125
- def _get_device(self) -> torch.device:
126
- return self.point_embeddings[0].weight.device
127
-
128
- def forward(
129
- self,
130
- points: Optional[Tuple[torch.Tensor, torch.Tensor]],
131
- boxes: Optional[torch.Tensor],
132
- masks: Optional[torch.Tensor],
133
- ) -> Tuple[torch.Tensor, torch.Tensor]:
134
- """
135
- Embeds different types of prompts, returning both sparse and dense
136
- embeddings.
137
-
138
- Arguments:
139
- points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates
140
- and labels to embed.
141
- boxes (torch.Tensor or none): boxes to embed
142
- masks (torch.Tensor or none): masks to embed
143
-
144
- Returns:
145
- torch.Tensor: sparse embeddings for the points and boxes, with shape
146
- BxNx(embed_dim), where N is determined by the number of input points
147
- and boxes.
148
- torch.Tensor: dense embeddings for the masks, in the shape
149
- Bx(embed_dim)x(embed_H)x(embed_W)
150
- """
151
- bs = self._get_batch_size(points, boxes, masks)
152
- sparse_embeddings = torch.empty((bs, 0, self.embed_dim), device=self._get_device())
153
- if points is not None:
154
- coords, labels = points
155
- point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))
156
- sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
157
- if boxes is not None:
158
- box_embeddings = self._embed_boxes(boxes)
159
- sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
160
-
161
- if masks is not None:
162
- dense_embeddings = self._embed_masks(masks)
163
- else:
164
- dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
165
- bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]
166
- )
167
-
168
- return sparse_embeddings, dense_embeddings
169
-
170
-
171
- class PositionEmbeddingRandom(nn.Module):
172
- """
173
- Positional encoding using random spatial frequencies.
174
- """
175
-
176
- def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:
177
- super().__init__()
178
- if scale is None or scale <= 0.0:
179
- scale = 1.0
180
- self.register_buffer(
181
- "positional_encoding_gaussian_matrix",
182
- scale * torch.randn((2, num_pos_feats)),
183
- )
184
-
185
- def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
186
- """Positionally encode points that are normalized to [0,1]."""
187
- # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
188
- coords = 2 * coords - 1
189
- coords = coords @ self.positional_encoding_gaussian_matrix
190
- coords = 2 * np.pi * coords
191
- # outputs d_1 x ... x d_n x C shape
192
- return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
193
-
194
- def forward(self, size: Tuple[int, int]) -> torch.Tensor:
195
- """Generate positional encoding for a grid of the specified size."""
196
- h, w = size
197
- device: Any = self.positional_encoding_gaussian_matrix.device
198
- grid = torch.ones((h, w), device=device, dtype=torch.float32)
199
- y_embed = grid.cumsum(dim=0) - 0.5
200
- x_embed = grid.cumsum(dim=1) - 0.5
201
- y_embed = y_embed / h
202
- x_embed = x_embed / w
203
-
204
- pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
205
- return pe.permute(2, 0, 1) # C x H x W
206
-
207
- def forward_with_coords(
208
- self, coords_input: torch.Tensor, image_size: Tuple[int, int]
209
- ) -> torch.Tensor:
210
- """Positionally encode points that are not normalized to [0,1]."""
211
- coords = coords_input.clone()
212
- coords[:, :, 0] = coords[:, :, 0] / image_size[1]
213
- coords[:, :, 1] = coords[:, :, 1] / image_size[0]
214
- return self._pe_encoding(coords.to(torch.float)) # B x N x C
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/modeling/sam.py DELETED
@@ -1,174 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import torch
8
- from torch import nn
9
- from torch.nn import functional as F
10
-
11
- from typing import Any, Dict, List, Tuple
12
-
13
- from .image_encoder import ImageEncoderViT
14
- from .mask_decoder import MaskDecoder
15
- from .prompt_encoder import PromptEncoder
16
-
17
-
18
- class Sam(nn.Module):
19
- mask_threshold: float = 0.0
20
- image_format: str = "RGB"
21
-
22
- def __init__(
23
- self,
24
- image_encoder: ImageEncoderViT,
25
- prompt_encoder: PromptEncoder,
26
- mask_decoder: MaskDecoder,
27
- pixel_mean: List[float] = [123.675, 116.28, 103.53],
28
- pixel_std: List[float] = [58.395, 57.12, 57.375],
29
- ) -> None:
30
- """
31
- SAM predicts object masks from an image and input prompts.
32
-
33
- Arguments:
34
- image_encoder (ImageEncoderViT): The backbone used to encode the
35
- image into image embeddings that allow for efficient mask prediction.
36
- prompt_encoder (PromptEncoder): Encodes various types of input prompts.
37
- mask_decoder (MaskDecoder): Predicts masks from the image embeddings
38
- and encoded prompts.
39
- pixel_mean (list(float)): Mean values for normalizing pixels in the input image.
40
- pixel_std (list(float)): Std values for normalizing pixels in the input image.
41
- """
42
- super().__init__()
43
- self.image_encoder = image_encoder
44
- self.prompt_encoder = prompt_encoder
45
- self.mask_decoder = mask_decoder
46
- self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False)
47
- self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False)
48
-
49
- @property
50
- def device(self) -> Any:
51
- return self.pixel_mean.device
52
-
53
- @torch.no_grad()
54
- def forward(
55
- self,
56
- batched_input: List[Dict[str, Any]],
57
- multimask_output: bool,
58
- ) -> List[Dict[str, torch.Tensor]]:
59
- """
60
- Predicts masks end-to-end from provided images and prompts.
61
- If prompts are not known in advance, using SamPredictor is
62
- recommended over calling the model directly.
63
-
64
- Arguments:
65
- batched_input (list(dict)): A list over input images, each a
66
- dictionary with the following keys. A prompt key can be
67
- excluded if it is not present.
68
- 'image': The image as a torch tensor in 3xHxW format,
69
- already transformed for input to the model.
70
- 'original_size': (tuple(int, int)) The original size of
71
- the image before transformation, as (H, W).
72
- 'point_coords': (torch.Tensor) Batched point prompts for
73
- this image, with shape BxNx2. Already transformed to the
74
- input frame of the model.
75
- 'point_labels': (torch.Tensor) Batched labels for point prompts,
76
- with shape BxN.
77
- 'boxes': (torch.Tensor) Batched box inputs, with shape Bx4.
78
- Already transformed to the input frame of the model.
79
- 'mask_inputs': (torch.Tensor) Batched mask inputs to the model,
80
- in the form Bx1xHxW.
81
- multimask_output (bool): Whether the model should predict multiple
82
- disambiguating masks, or return a single mask.
83
-
84
- Returns:
85
- (list(dict)): A list over input images, where each element is
86
- as dictionary with the following keys.
87
- 'masks': (torch.Tensor) Batched binary mask predictions,
88
- with shape BxCxHxW, where B is the number of input promts,
89
- C is determiend by multimask_output, and (H, W) is the
90
- original size of the image.
91
- 'iou_predictions': (torch.Tensor) The model's predictions
92
- of mask quality, in shape BxC.
93
- 'low_res_logits': (torch.Tensor) Low resolution logits with
94
- shape BxCxHxW, where H=W=256. Can be passed as mask input
95
- to subsequent iterations of prediction.
96
- """
97
- input_images = torch.stack([self.preprocess(x["image"]) for x in batched_input], dim=0)
98
- image_embeddings = self.image_encoder(input_images)
99
-
100
- outputs = []
101
- for image_record, curr_embedding in zip(batched_input, image_embeddings):
102
- if "point_coords" in image_record:
103
- points = (image_record["point_coords"], image_record["point_labels"])
104
- else:
105
- points = None
106
- sparse_embeddings, dense_embeddings = self.prompt_encoder(
107
- points=points,
108
- boxes=image_record.get("boxes", None),
109
- masks=image_record.get("mask_inputs", None),
110
- )
111
- low_res_masks, iou_predictions = self.mask_decoder(
112
- image_embeddings=curr_embedding.unsqueeze(0),
113
- image_pe=self.prompt_encoder.get_dense_pe(),
114
- sparse_prompt_embeddings=sparse_embeddings,
115
- dense_prompt_embeddings=dense_embeddings,
116
- multimask_output=multimask_output,
117
- )
118
- masks = self.postprocess_masks(
119
- low_res_masks,
120
- input_size=image_record["image"].shape[-2:],
121
- original_size=image_record["original_size"],
122
- )
123
- masks = masks > self.mask_threshold
124
- outputs.append(
125
- {
126
- "masks": masks,
127
- "iou_predictions": iou_predictions,
128
- "low_res_logits": low_res_masks,
129
- }
130
- )
131
- return outputs
132
-
133
- def postprocess_masks(
134
- self,
135
- masks: torch.Tensor,
136
- input_size: Tuple[int, ...],
137
- original_size: Tuple[int, ...],
138
- ) -> torch.Tensor:
139
- """
140
- Remove padding and upscale masks to the original image size.
141
-
142
- Arguments:
143
- masks (torch.Tensor): Batched masks from the mask_decoder,
144
- in BxCxHxW format.
145
- input_size (tuple(int, int)): The size of the image input to the
146
- model, in (H, W) format. Used to remove padding.
147
- original_size (tuple(int, int)): The original size of the image
148
- before resizing for input to the model, in (H, W) format.
149
-
150
- Returns:
151
- (torch.Tensor): Batched masks in BxCxHxW format, where (H, W)
152
- is given by original_size.
153
- """
154
- masks = F.interpolate(
155
- masks,
156
- (self.image_encoder.img_size, self.image_encoder.img_size),
157
- mode="bilinear",
158
- align_corners=False,
159
- )
160
- masks = masks[..., : input_size[0], : input_size[1]]
161
- masks = F.interpolate(masks, original_size, mode="bilinear", align_corners=False)
162
- return masks
163
-
164
- def preprocess(self, x: torch.Tensor) -> torch.Tensor:
165
- """Normalize pixel values and pad to a square input."""
166
- # Normalize colors
167
- x = (x - self.pixel_mean) / self.pixel_std
168
-
169
- # Pad
170
- h, w = x.shape[-2:]
171
- padh = self.image_encoder.img_size - h
172
- padw = self.image_encoder.img_size - w
173
- x = F.pad(x, (0, padw, 0, padh))
174
- return x
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/modeling/transformer.py DELETED
@@ -1,240 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import torch
8
- from torch import Tensor, nn
9
-
10
- import math
11
- from typing import Tuple, Type
12
-
13
- from .common import MLPBlock
14
-
15
-
16
- class TwoWayTransformer(nn.Module):
17
- def __init__(
18
- self,
19
- depth: int,
20
- embedding_dim: int,
21
- num_heads: int,
22
- mlp_dim: int,
23
- activation: Type[nn.Module] = nn.ReLU,
24
- attention_downsample_rate: int = 2,
25
- ) -> None:
26
- """
27
- A transformer decoder that attends to an input image using
28
- queries whose positional embedding is supplied.
29
-
30
- Args:
31
- depth (int): number of layers in the transformer
32
- embedding_dim (int): the channel dimension for the input embeddings
33
- num_heads (int): the number of heads for multihead attention. Must
34
- divide embedding_dim
35
- mlp_dim (int): the channel dimension internal to the MLP block
36
- activation (nn.Module): the activation to use in the MLP block
37
- """
38
- super().__init__()
39
- self.depth = depth
40
- self.embedding_dim = embedding_dim
41
- self.num_heads = num_heads
42
- self.mlp_dim = mlp_dim
43
- self.layers = nn.ModuleList()
44
-
45
- for i in range(depth):
46
- self.layers.append(
47
- TwoWayAttentionBlock(
48
- embedding_dim=embedding_dim,
49
- num_heads=num_heads,
50
- mlp_dim=mlp_dim,
51
- activation=activation,
52
- attention_downsample_rate=attention_downsample_rate,
53
- skip_first_layer_pe=(i == 0),
54
- )
55
- )
56
-
57
- self.final_attn_token_to_image = Attention(
58
- embedding_dim, num_heads, downsample_rate=attention_downsample_rate
59
- )
60
- self.norm_final_attn = nn.LayerNorm(embedding_dim)
61
-
62
- def forward(
63
- self,
64
- image_embedding: Tensor,
65
- image_pe: Tensor,
66
- point_embedding: Tensor,
67
- ) -> Tuple[Tensor, Tensor]:
68
- """
69
- Args:
70
- image_embedding (torch.Tensor): image to attend to. Should be shape
71
- B x embedding_dim x h x w for any h and w.
72
- image_pe (torch.Tensor): the positional encoding to add to the image. Must
73
- have the same shape as image_embedding.
74
- point_embedding (torch.Tensor): the embedding to add to the query points.
75
- Must have shape B x N_points x embedding_dim for any N_points.
76
-
77
- Returns:
78
- torch.Tensor: the processed point_embedding
79
- torch.Tensor: the processed image_embedding
80
- """
81
- # BxCxHxW -> BxHWxC == B x N_image_tokens x C
82
- bs, c, h, w = image_embedding.shape
83
- image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
84
- image_pe = image_pe.flatten(2).permute(0, 2, 1)
85
-
86
- # Prepare queries
87
- queries = point_embedding
88
- keys = image_embedding
89
-
90
- # Apply transformer blocks and final layernorm
91
- for layer in self.layers:
92
- queries, keys = layer(
93
- queries=queries,
94
- keys=keys,
95
- query_pe=point_embedding,
96
- key_pe=image_pe,
97
- )
98
-
99
- # Apply the final attenion layer from the points to the image
100
- q = queries + point_embedding
101
- k = keys + image_pe
102
- attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
103
- queries = queries + attn_out
104
- queries = self.norm_final_attn(queries)
105
-
106
- return queries, keys
107
-
108
-
109
- class TwoWayAttentionBlock(nn.Module):
110
- def __init__(
111
- self,
112
- embedding_dim: int,
113
- num_heads: int,
114
- mlp_dim: int = 2048,
115
- activation: Type[nn.Module] = nn.ReLU,
116
- attention_downsample_rate: int = 2,
117
- skip_first_layer_pe: bool = False,
118
- ) -> None:
119
- """
120
- A transformer block with four layers: (1) self-attention of sparse
121
- inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp
122
- block on sparse inputs, and (4) cross attention of dense inputs to sparse
123
- inputs.
124
-
125
- Arguments:
126
- embedding_dim (int): the channel dimension of the embeddings
127
- num_heads (int): the number of heads in the attention layers
128
- mlp_dim (int): the hidden dimension of the mlp block
129
- activation (nn.Module): the activation of the mlp block
130
- skip_first_layer_pe (bool): skip the PE on the first layer
131
- """
132
- super().__init__()
133
- self.self_attn = Attention(embedding_dim, num_heads)
134
- self.norm1 = nn.LayerNorm(embedding_dim)
135
-
136
- self.cross_attn_token_to_image = Attention(
137
- embedding_dim, num_heads, downsample_rate=attention_downsample_rate
138
- )
139
- self.norm2 = nn.LayerNorm(embedding_dim)
140
-
141
- self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
142
- self.norm3 = nn.LayerNorm(embedding_dim)
143
-
144
- self.norm4 = nn.LayerNorm(embedding_dim)
145
- self.cross_attn_image_to_token = Attention(
146
- embedding_dim, num_heads, downsample_rate=attention_downsample_rate
147
- )
148
-
149
- self.skip_first_layer_pe = skip_first_layer_pe
150
-
151
- def forward(
152
- self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor
153
- ) -> Tuple[Tensor, Tensor]:
154
- # Self attention block
155
- if self.skip_first_layer_pe:
156
- queries = self.self_attn(q=queries, k=queries, v=queries)
157
- else:
158
- q = queries + query_pe
159
- attn_out = self.self_attn(q=q, k=q, v=queries)
160
- queries = queries + attn_out
161
- queries = self.norm1(queries)
162
-
163
- # Cross attention block, tokens attending to image embedding
164
- q = queries + query_pe
165
- k = keys + key_pe
166
- attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
167
- queries = queries + attn_out
168
- queries = self.norm2(queries)
169
-
170
- # MLP block
171
- mlp_out = self.mlp(queries)
172
- queries = queries + mlp_out
173
- queries = self.norm3(queries)
174
-
175
- # Cross attention block, image embedding attending to tokens
176
- q = queries + query_pe
177
- k = keys + key_pe
178
- attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)
179
- keys = keys + attn_out
180
- keys = self.norm4(keys)
181
-
182
- return queries, keys
183
-
184
-
185
- class Attention(nn.Module):
186
- """
187
- An attention layer that allows for downscaling the size of the embedding
188
- after projection to queries, keys, and values.
189
- """
190
-
191
- def __init__(
192
- self,
193
- embedding_dim: int,
194
- num_heads: int,
195
- downsample_rate: int = 1,
196
- ) -> None:
197
- super().__init__()
198
- self.embedding_dim = embedding_dim
199
- self.internal_dim = embedding_dim // downsample_rate
200
- self.num_heads = num_heads
201
- assert self.internal_dim % num_heads == 0, "num_heads must divide embedding_dim."
202
-
203
- self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
204
- self.k_proj = nn.Linear(embedding_dim, self.internal_dim)
205
- self.v_proj = nn.Linear(embedding_dim, self.internal_dim)
206
- self.out_proj = nn.Linear(self.internal_dim, embedding_dim)
207
-
208
- def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:
209
- b, n, c = x.shape
210
- x = x.reshape(b, n, num_heads, c // num_heads)
211
- return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head
212
-
213
- def _recombine_heads(self, x: Tensor) -> Tensor:
214
- b, n_heads, n_tokens, c_per_head = x.shape
215
- x = x.transpose(1, 2)
216
- return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C
217
-
218
- def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
219
- # Input projections
220
- q = self.q_proj(q)
221
- k = self.k_proj(k)
222
- v = self.v_proj(v)
223
-
224
- # Separate into heads
225
- q = self._separate_heads(q, self.num_heads)
226
- k = self._separate_heads(k, self.num_heads)
227
- v = self._separate_heads(v, self.num_heads)
228
-
229
- # Attention
230
- _, _, _, c_per_head = q.shape
231
- attn = q @ k.permute(0, 1, 3, 2) # B x N_heads x N_tokens x N_tokens
232
- attn = attn / math.sqrt(c_per_head)
233
- attn = torch.softmax(attn, dim=-1)
234
-
235
- # Get output
236
- out = attn @ v
237
- out = self._recombine_heads(out)
238
- out = self.out_proj(out)
239
-
240
- return out
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/predictor.py DELETED
@@ -1,269 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import numpy as np
8
- import torch
9
-
10
- from segment_anything.modeling import Sam
11
-
12
- from typing import Optional, Tuple
13
-
14
- from .utils.transforms import ResizeLongestSide
15
-
16
-
17
- class SamPredictor:
18
- def __init__(
19
- self,
20
- sam_model: Sam,
21
- ) -> None:
22
- """
23
- Uses SAM to calculate the image embedding for an image, and then
24
- allow repeated, efficient mask prediction given prompts.
25
-
26
- Arguments:
27
- sam_model (Sam): The model to use for mask prediction.
28
- """
29
- super().__init__()
30
- self.model = sam_model
31
- self.transform = ResizeLongestSide(sam_model.image_encoder.img_size)
32
- self.reset_image()
33
-
34
- def set_image(
35
- self,
36
- image: np.ndarray,
37
- image_format: str = "RGB",
38
- ) -> None:
39
- """
40
- Calculates the image embeddings for the provided image, allowing
41
- masks to be predicted with the 'predict' method.
42
-
43
- Arguments:
44
- image (np.ndarray): The image for calculating masks. Expects an
45
- image in HWC uint8 format, with pixel values in [0, 255].
46
- image_format (str): The color format of the image, in ['RGB', 'BGR'].
47
- """
48
- assert image_format in [
49
- "RGB",
50
- "BGR",
51
- ], f"image_format must be in ['RGB', 'BGR'], is {image_format}."
52
- if image_format != self.model.image_format:
53
- image = image[..., ::-1]
54
-
55
- # Transform the image to the form expected by the model
56
- input_image = self.transform.apply_image(image)
57
- input_image_torch = torch.as_tensor(input_image, device=self.device)
58
- input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :]
59
-
60
- self.set_torch_image(input_image_torch, image.shape[:2])
61
-
62
- @torch.no_grad()
63
- def set_torch_image(
64
- self,
65
- transformed_image: torch.Tensor,
66
- original_image_size: Tuple[int, ...],
67
- ) -> None:
68
- """
69
- Calculates the image embeddings for the provided image, allowing
70
- masks to be predicted with the 'predict' method. Expects the input
71
- image to be already transformed to the format expected by the model.
72
-
73
- Arguments:
74
- transformed_image (torch.Tensor): The input image, with shape
75
- 1x3xHxW, which has been transformed with ResizeLongestSide.
76
- original_image_size (tuple(int, int)): The size of the image
77
- before transformation, in (H, W) format.
78
- """
79
- assert (
80
- len(transformed_image.shape) == 4
81
- and transformed_image.shape[1] == 3
82
- and max(*transformed_image.shape[2:]) == self.model.image_encoder.img_size
83
- ), f"set_torch_image input must be BCHW with long side {self.model.image_encoder.img_size}."
84
- self.reset_image()
85
-
86
- self.original_size = original_image_size
87
- self.input_size = tuple(transformed_image.shape[-2:])
88
- input_image = self.model.preprocess(transformed_image)
89
- self.features = self.model.image_encoder(input_image)
90
- self.is_image_set = True
91
-
92
- def predict(
93
- self,
94
- point_coords: Optional[np.ndarray] = None,
95
- point_labels: Optional[np.ndarray] = None,
96
- box: Optional[np.ndarray] = None,
97
- mask_input: Optional[np.ndarray] = None,
98
- multimask_output: bool = True,
99
- return_logits: bool = False,
100
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
101
- """
102
- Predict masks for the given input prompts, using the currently set image.
103
-
104
- Arguments:
105
- point_coords (np.ndarray or None): A Nx2 array of point prompts to the
106
- model. Each point is in (X,Y) in pixels.
107
- point_labels (np.ndarray or None): A length N array of labels for the
108
- point prompts. 1 indicates a foreground point and 0 indicates a
109
- background point.
110
- box (np.ndarray or None): A length 4 array given a box prompt to the
111
- model, in XYXY format.
112
- mask_input (np.ndarray): A low resolution mask input to the model, typically
113
- coming from a previous prediction iteration. Has form 1xHxW, where
114
- for SAM, H=W=256.
115
- multimask_output (bool): If true, the model will return three masks.
116
- For ambiguous input prompts (such as a single click), this will often
117
- produce better masks than a single prediction. If only a single
118
- mask is needed, the model's predicted quality score can be used
119
- to select the best mask. For non-ambiguous prompts, such as multiple
120
- input prompts, multimask_output=False can give better results.
121
- return_logits (bool): If true, returns un-thresholded masks logits
122
- instead of a binary mask.
123
-
124
- Returns:
125
- (np.ndarray): The output masks in CxHxW format, where C is the
126
- number of masks, and (H, W) is the original image size.
127
- (np.ndarray): An array of length C containing the model's
128
- predictions for the quality of each mask.
129
- (np.ndarray): An array of shape CxHxW, where C is the number
130
- of masks and H=W=256. These low resolution logits can be passed to
131
- a subsequent iteration as mask input.
132
- """
133
- if not self.is_image_set:
134
- raise RuntimeError("An image must be set with .set_image(...) before mask prediction.")
135
-
136
- # Transform input prompts
137
- coords_torch, labels_torch, box_torch, mask_input_torch = None, None, None, None
138
- if point_coords is not None:
139
- assert (
140
- point_labels is not None
141
- ), "point_labels must be supplied if point_coords is supplied."
142
- point_coords = self.transform.apply_coords(point_coords, self.original_size)
143
- coords_torch = torch.as_tensor(point_coords, dtype=torch.float, device=self.device)
144
- labels_torch = torch.as_tensor(point_labels, dtype=torch.int, device=self.device)
145
- coords_torch, labels_torch = coords_torch[None, :, :], labels_torch[None, :]
146
- if box is not None:
147
- box = self.transform.apply_boxes(box, self.original_size)
148
- box_torch = torch.as_tensor(box, dtype=torch.float, device=self.device)
149
- box_torch = box_torch[None, :]
150
- if mask_input is not None:
151
- mask_input_torch = torch.as_tensor(mask_input, dtype=torch.float, device=self.device)
152
- mask_input_torch = mask_input_torch[None, :, :, :]
153
-
154
- masks, iou_predictions, low_res_masks = self.predict_torch(
155
- coords_torch,
156
- labels_torch,
157
- box_torch,
158
- mask_input_torch,
159
- multimask_output,
160
- return_logits=return_logits,
161
- )
162
-
163
- masks = masks[0].detach().cpu().numpy()
164
- iou_predictions = iou_predictions[0].detach().cpu().numpy()
165
- low_res_masks = low_res_masks[0].detach().cpu().numpy()
166
- return masks, iou_predictions, low_res_masks
167
-
168
- @torch.no_grad()
169
- def predict_torch(
170
- self,
171
- point_coords: Optional[torch.Tensor],
172
- point_labels: Optional[torch.Tensor],
173
- boxes: Optional[torch.Tensor] = None,
174
- mask_input: Optional[torch.Tensor] = None,
175
- multimask_output: bool = True,
176
- return_logits: bool = False,
177
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
178
- """
179
- Predict masks for the given input prompts, using the currently set image.
180
- Input prompts are batched torch tensors and are expected to already be
181
- transformed to the input frame using ResizeLongestSide.
182
-
183
- Arguments:
184
- point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the
185
- model. Each point is in (X,Y) in pixels.
186
- point_labels (torch.Tensor or None): A BxN array of labels for the
187
- point prompts. 1 indicates a foreground point and 0 indicates a
188
- background point.
189
- box (np.ndarray or None): A Bx4 array given a box prompt to the
190
- model, in XYXY format.
191
- mask_input (np.ndarray): A low resolution mask input to the model, typically
192
- coming from a previous prediction iteration. Has form Bx1xHxW, where
193
- for SAM, H=W=256. Masks returned by a previous iteration of the
194
- predict method do not need further transformation.
195
- multimask_output (bool): If true, the model will return three masks.
196
- For ambiguous input prompts (such as a single click), this will often
197
- produce better masks than a single prediction. If only a single
198
- mask is needed, the model's predicted quality score can be used
199
- to select the best mask. For non-ambiguous prompts, such as multiple
200
- input prompts, multimask_output=False can give better results.
201
- return_logits (bool): If true, returns un-thresholded masks logits
202
- instead of a binary mask.
203
-
204
- Returns:
205
- (torch.Tensor): The output masks in BxCxHxW format, where C is the
206
- number of masks, and (H, W) is the original image size.
207
- (torch.Tensor): An array of shape BxC containing the model's
208
- predictions for the quality of each mask.
209
- (torch.Tensor): An array of shape BxCxHxW, where C is the number
210
- of masks and H=W=256. These low res logits can be passed to
211
- a subsequent iteration as mask input.
212
- """
213
- if not self.is_image_set:
214
- raise RuntimeError("An image must be set with .set_image(...) before mask prediction.")
215
-
216
- if point_coords is not None:
217
- points = (point_coords, point_labels)
218
- else:
219
- points = None
220
-
221
- # Embed prompts
222
- sparse_embeddings, dense_embeddings = self.model.prompt_encoder(
223
- points=points,
224
- boxes=boxes,
225
- masks=mask_input,
226
- )
227
-
228
- # Predict masks
229
- low_res_masks, iou_predictions = self.model.mask_decoder(
230
- image_embeddings=self.features,
231
- image_pe=self.model.prompt_encoder.get_dense_pe(),
232
- sparse_prompt_embeddings=sparse_embeddings,
233
- dense_prompt_embeddings=dense_embeddings,
234
- multimask_output=multimask_output,
235
- )
236
-
237
- # Upscale the masks to the original image resolution
238
- masks = self.model.postprocess_masks(low_res_masks, self.input_size, self.original_size)
239
-
240
- if not return_logits:
241
- masks = masks > self.model.mask_threshold
242
-
243
- return masks, iou_predictions, low_res_masks
244
-
245
- def get_image_embedding(self) -> torch.Tensor:
246
- """
247
- Returns the image embeddings for the currently set image, with
248
- shape 1xCxHxW, where C is the embedding dimension and (H,W) are
249
- the embedding spatial dimension of SAM (typically C=256, H=W=64).
250
- """
251
- if not self.is_image_set:
252
- raise RuntimeError(
253
- "An image must be set with .set_image(...) to generate an embedding."
254
- )
255
- assert self.features is not None, "Features must exist if an image has been set."
256
- return self.features
257
-
258
- @property
259
- def device(self) -> torch.device:
260
- return self.model.device
261
-
262
- def reset_image(self) -> None:
263
- """Resets the currently set image."""
264
- self.is_image_set = False
265
- self.features = None
266
- self.orig_h = None
267
- self.orig_w = None
268
- self.input_h = None
269
- self.input_w = None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/utils/__init__.py DELETED
@@ -1,5 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
 
 
 
 
 
 
segment_anything/segment_anything/utils/__pycache__/__init__.cpython-39.pyc DELETED
Binary file (219 Bytes)
 
segment_anything/segment_anything/utils/__pycache__/amg.cpython-39.pyc DELETED
Binary file (12.2 kB)
 
segment_anything/segment_anything/utils/__pycache__/transforms.cpython-39.pyc DELETED
Binary file (4.02 kB)
 
segment_anything/segment_anything/utils/amg.py DELETED
@@ -1,346 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import numpy as np
8
- import torch
9
-
10
- import math
11
- from copy import deepcopy
12
- from itertools import product
13
- from typing import Any, Dict, Generator, ItemsView, List, Tuple
14
-
15
-
16
- class MaskData:
17
- """
18
- A structure for storing masks and their related data in batched format.
19
- Implements basic filtering and concatenation.
20
- """
21
-
22
- def __init__(self, **kwargs) -> None:
23
- for v in kwargs.values():
24
- assert isinstance(
25
- v, (list, np.ndarray, torch.Tensor)
26
- ), "MaskData only supports list, numpy arrays, and torch tensors."
27
- self._stats = dict(**kwargs)
28
-
29
- def __setitem__(self, key: str, item: Any) -> None:
30
- assert isinstance(
31
- item, (list, np.ndarray, torch.Tensor)
32
- ), "MaskData only supports list, numpy arrays, and torch tensors."
33
- self._stats[key] = item
34
-
35
- def __delitem__(self, key: str) -> None:
36
- del self._stats[key]
37
-
38
- def __getitem__(self, key: str) -> Any:
39
- return self._stats[key]
40
-
41
- def items(self) -> ItemsView[str, Any]:
42
- return self._stats.items()
43
-
44
- def filter(self, keep: torch.Tensor) -> None:
45
- for k, v in self._stats.items():
46
- if v is None:
47
- self._stats[k] = None
48
- elif isinstance(v, torch.Tensor):
49
- self._stats[k] = v[torch.as_tensor(keep, device=v.device)]
50
- elif isinstance(v, np.ndarray):
51
- self._stats[k] = v[keep.detach().cpu().numpy()]
52
- elif isinstance(v, list) and keep.dtype == torch.bool:
53
- self._stats[k] = [a for i, a in enumerate(v) if keep[i]]
54
- elif isinstance(v, list):
55
- self._stats[k] = [v[i] for i in keep]
56
- else:
57
- raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.")
58
-
59
- def cat(self, new_stats: "MaskData") -> None:
60
- for k, v in new_stats.items():
61
- if k not in self._stats or self._stats[k] is None:
62
- self._stats[k] = deepcopy(v)
63
- elif isinstance(v, torch.Tensor):
64
- self._stats[k] = torch.cat([self._stats[k], v], dim=0)
65
- elif isinstance(v, np.ndarray):
66
- self._stats[k] = np.concatenate([self._stats[k], v], axis=0)
67
- elif isinstance(v, list):
68
- self._stats[k] = self._stats[k] + deepcopy(v)
69
- else:
70
- raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.")
71
-
72
- def to_numpy(self) -> None:
73
- for k, v in self._stats.items():
74
- if isinstance(v, torch.Tensor):
75
- self._stats[k] = v.detach().cpu().numpy()
76
-
77
-
78
- def is_box_near_crop_edge(
79
- boxes: torch.Tensor, crop_box: List[int], orig_box: List[int], atol: float = 20.0
80
- ) -> torch.Tensor:
81
- """Filter masks at the edge of a crop, but not at the edge of the original image."""
82
- crop_box_torch = torch.as_tensor(crop_box, dtype=torch.float, device=boxes.device)
83
- orig_box_torch = torch.as_tensor(orig_box, dtype=torch.float, device=boxes.device)
84
- boxes = uncrop_boxes_xyxy(boxes, crop_box).float()
85
- near_crop_edge = torch.isclose(boxes, crop_box_torch[None, :], atol=atol, rtol=0)
86
- near_image_edge = torch.isclose(boxes, orig_box_torch[None, :], atol=atol, rtol=0)
87
- near_crop_edge = torch.logical_and(near_crop_edge, ~near_image_edge)
88
- return torch.any(near_crop_edge, dim=1)
89
-
90
-
91
- def box_xyxy_to_xywh(box_xyxy: torch.Tensor) -> torch.Tensor:
92
- box_xywh = deepcopy(box_xyxy)
93
- box_xywh[2] = box_xywh[2] - box_xywh[0]
94
- box_xywh[3] = box_xywh[3] - box_xywh[1]
95
- return box_xywh
96
-
97
-
98
- def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]:
99
- assert len(args) > 0 and all(
100
- len(a) == len(args[0]) for a in args
101
- ), "Batched iteration must have inputs of all the same size."
102
- n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0)
103
- for b in range(n_batches):
104
- yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args]
105
-
106
-
107
- def mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]:
108
- """
109
- Encodes masks to an uncompressed RLE, in the format expected by
110
- pycoco tools.
111
- """
112
- # Put in fortran order and flatten h,w
113
- b, h, w = tensor.shape
114
- tensor = tensor.permute(0, 2, 1).flatten(1)
115
-
116
- # Compute change indices
117
- diff = tensor[:, 1:] ^ tensor[:, :-1]
118
- change_indices = diff.nonzero()
119
-
120
- # Encode run length
121
- out = []
122
- for i in range(b):
123
- cur_idxs = change_indices[change_indices[:, 0] == i, 1]
124
- cur_idxs = torch.cat(
125
- [
126
- torch.tensor([0], dtype=cur_idxs.dtype, device=cur_idxs.device),
127
- cur_idxs + 1,
128
- torch.tensor([h * w], dtype=cur_idxs.dtype, device=cur_idxs.device),
129
- ]
130
- )
131
- btw_idxs = cur_idxs[1:] - cur_idxs[:-1]
132
- counts = [] if tensor[i, 0] == 0 else [0]
133
- counts.extend(btw_idxs.detach().cpu().tolist())
134
- out.append({"size": [h, w], "counts": counts})
135
- return out
136
-
137
-
138
- def rle_to_mask(rle: Dict[str, Any]) -> np.ndarray:
139
- """Compute a binary mask from an uncompressed RLE."""
140
- h, w = rle["size"]
141
- mask = np.empty(h * w, dtype=bool)
142
- idx = 0
143
- parity = False
144
- for count in rle["counts"]:
145
- mask[idx : idx + count] = parity
146
- idx += count
147
- parity ^= True
148
- mask = mask.reshape(w, h)
149
- return mask.transpose() # Put in C order
150
-
151
-
152
- def area_from_rle(rle: Dict[str, Any]) -> int:
153
- return sum(rle["counts"][1::2])
154
-
155
-
156
- def calculate_stability_score(
157
- masks: torch.Tensor, mask_threshold: float, threshold_offset: float
158
- ) -> torch.Tensor:
159
- """
160
- Computes the stability score for a batch of masks. The stability
161
- score is the IoU between the binary masks obtained by thresholding
162
- the predicted mask logits at high and low values.
163
- """
164
- # One mask is always contained inside the other.
165
- # Save memory by preventing unnecesary cast to torch.int64
166
- intersections = (
167
- (masks > (mask_threshold + threshold_offset))
168
- .sum(-1, dtype=torch.int16)
169
- .sum(-1, dtype=torch.int32)
170
- )
171
- unions = (
172
- (masks > (mask_threshold - threshold_offset))
173
- .sum(-1, dtype=torch.int16)
174
- .sum(-1, dtype=torch.int32)
175
- )
176
- return intersections / unions
177
-
178
-
179
- def build_point_grid(n_per_side: int) -> np.ndarray:
180
- """Generates a 2D grid of points evenly spaced in [0,1]x[0,1]."""
181
- offset = 1 / (2 * n_per_side)
182
- points_one_side = np.linspace(offset, 1 - offset, n_per_side)
183
- points_x = np.tile(points_one_side[None, :], (n_per_side, 1))
184
- points_y = np.tile(points_one_side[:, None], (1, n_per_side))
185
- points = np.stack([points_x, points_y], axis=-1).reshape(-1, 2)
186
- return points
187
-
188
-
189
- def build_all_layer_point_grids(
190
- n_per_side: int, n_layers: int, scale_per_layer: int
191
- ) -> List[np.ndarray]:
192
- """Generates point grids for all crop layers."""
193
- points_by_layer = []
194
- for i in range(n_layers + 1):
195
- n_points = int(n_per_side / (scale_per_layer**i))
196
- points_by_layer.append(build_point_grid(n_points))
197
- return points_by_layer
198
-
199
-
200
- def generate_crop_boxes(
201
- im_size: Tuple[int, ...], n_layers: int, overlap_ratio: float
202
- ) -> Tuple[List[List[int]], List[int]]:
203
- """
204
- Generates a list of crop boxes of different sizes. Each layer
205
- has (2**i)**2 boxes for the ith layer.
206
- """
207
- crop_boxes, layer_idxs = [], []
208
- im_h, im_w = im_size
209
- short_side = min(im_h, im_w)
210
-
211
- # Original image
212
- crop_boxes.append([0, 0, im_w, im_h])
213
- layer_idxs.append(0)
214
-
215
- def crop_len(orig_len, n_crops, overlap):
216
- return int(math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops))
217
-
218
- for i_layer in range(n_layers):
219
- n_crops_per_side = 2 ** (i_layer + 1)
220
- overlap = int(overlap_ratio * short_side * (2 / n_crops_per_side))
221
-
222
- crop_w = crop_len(im_w, n_crops_per_side, overlap)
223
- crop_h = crop_len(im_h, n_crops_per_side, overlap)
224
-
225
- crop_box_x0 = [int((crop_w - overlap) * i) for i in range(n_crops_per_side)]
226
- crop_box_y0 = [int((crop_h - overlap) * i) for i in range(n_crops_per_side)]
227
-
228
- # Crops in XYWH format
229
- for x0, y0 in product(crop_box_x0, crop_box_y0):
230
- box = [x0, y0, min(x0 + crop_w, im_w), min(y0 + crop_h, im_h)]
231
- crop_boxes.append(box)
232
- layer_idxs.append(i_layer + 1)
233
-
234
- return crop_boxes, layer_idxs
235
-
236
-
237
- def uncrop_boxes_xyxy(boxes: torch.Tensor, crop_box: List[int]) -> torch.Tensor:
238
- x0, y0, _, _ = crop_box
239
- offset = torch.tensor([[x0, y0, x0, y0]], device=boxes.device)
240
- # Check if boxes has a channel dimension
241
- if len(boxes.shape) == 3:
242
- offset = offset.unsqueeze(1)
243
- return boxes + offset
244
-
245
-
246
- def uncrop_points(points: torch.Tensor, crop_box: List[int]) -> torch.Tensor:
247
- x0, y0, _, _ = crop_box
248
- offset = torch.tensor([[x0, y0]], device=points.device)
249
- # Check if points has a channel dimension
250
- if len(points.shape) == 3:
251
- offset = offset.unsqueeze(1)
252
- return points + offset
253
-
254
-
255
- def uncrop_masks(
256
- masks: torch.Tensor, crop_box: List[int], orig_h: int, orig_w: int
257
- ) -> torch.Tensor:
258
- x0, y0, x1, y1 = crop_box
259
- if x0 == 0 and y0 == 0 and x1 == orig_w and y1 == orig_h:
260
- return masks
261
- # Coordinate transform masks
262
- pad_x, pad_y = orig_w - (x1 - x0), orig_h - (y1 - y0)
263
- pad = (x0, pad_x - x0, y0, pad_y - y0)
264
- return torch.nn.functional.pad(masks, pad, value=0)
265
-
266
-
267
- def remove_small_regions(
268
- mask: np.ndarray, area_thresh: float, mode: str
269
- ) -> Tuple[np.ndarray, bool]:
270
- """
271
- Removes small disconnected regions and holes in a mask. Returns the
272
- mask and an indicator of if the mask has been modified.
273
- """
274
- import cv2 # type: ignore
275
-
276
- assert mode in ["holes", "islands"]
277
- correct_holes = mode == "holes"
278
- working_mask = (correct_holes ^ mask).astype(np.uint8)
279
- n_labels, regions, stats, _ = cv2.connectedComponentsWithStats(working_mask, 8)
280
- sizes = stats[:, -1][1:] # Row 0 is background label
281
- small_regions = [i + 1 for i, s in enumerate(sizes) if s < area_thresh]
282
- if len(small_regions) == 0:
283
- return mask, False
284
- fill_labels = [0] + small_regions
285
- if not correct_holes:
286
- fill_labels = [i for i in range(n_labels) if i not in fill_labels]
287
- # If every region is below threshold, keep largest
288
- if len(fill_labels) == 0:
289
- fill_labels = [int(np.argmax(sizes)) + 1]
290
- mask = np.isin(regions, fill_labels)
291
- return mask, True
292
-
293
-
294
- def coco_encode_rle(uncompressed_rle: Dict[str, Any]) -> Dict[str, Any]:
295
- from pycocotools import mask as mask_utils # type: ignore
296
-
297
- h, w = uncompressed_rle["size"]
298
- rle = mask_utils.frPyObjects(uncompressed_rle, h, w)
299
- rle["counts"] = rle["counts"].decode("utf-8") # Necessary to serialize with json
300
- return rle
301
-
302
-
303
- def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor:
304
- """
305
- Calculates boxes in XYXY format around masks. Return [0,0,0,0] for
306
- an empty mask. For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4.
307
- """
308
- # torch.max below raises an error on empty inputs, just skip in this case
309
- if torch.numel(masks) == 0:
310
- return torch.zeros(*masks.shape[:-2], 4, device=masks.device)
311
-
312
- # Normalize shape to CxHxW
313
- shape = masks.shape
314
- h, w = shape[-2:]
315
- if len(shape) > 2:
316
- masks = masks.flatten(0, -3)
317
- else:
318
- masks = masks.unsqueeze(0)
319
-
320
- # Get top and bottom edges
321
- in_height, _ = torch.max(masks, dim=-1)
322
- in_height_coords = in_height * torch.arange(h, device=in_height.device)[None, :]
323
- bottom_edges, _ = torch.max(in_height_coords, dim=-1)
324
- in_height_coords = in_height_coords + h * (~in_height)
325
- top_edges, _ = torch.min(in_height_coords, dim=-1)
326
-
327
- # Get left and right edges
328
- in_width, _ = torch.max(masks, dim=-2)
329
- in_width_coords = in_width * torch.arange(w, device=in_width.device)[None, :]
330
- right_edges, _ = torch.max(in_width_coords, dim=-1)
331
- in_width_coords = in_width_coords + w * (~in_width)
332
- left_edges, _ = torch.min(in_width_coords, dim=-1)
333
-
334
- # If the mask is empty the right edge will be to the left of the left edge.
335
- # Replace these boxes with [0, 0, 0, 0]
336
- empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges)
337
- out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1)
338
- out = out * (~empty_filter).unsqueeze(-1)
339
-
340
- # Return to original shape
341
- if len(shape) > 2:
342
- out = out.reshape(*shape[:-2], 4)
343
- else:
344
- out = out[0]
345
-
346
- return out
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/utils/onnx.py DELETED
@@ -1,144 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import torch
8
- import torch.nn as nn
9
- from torch.nn import functional as F
10
-
11
- from typing import Tuple
12
-
13
- from ..modeling import Sam
14
- from .amg import calculate_stability_score
15
-
16
-
17
- class SamOnnxModel(nn.Module):
18
- """
19
- This model should not be called directly, but is used in ONNX export.
20
- It combines the prompt encoder, mask decoder, and mask postprocessing of Sam,
21
- with some functions modified to enable model tracing. Also supports extra
22
- options controlling what information. See the ONNX export script for details.
23
- """
24
-
25
- def __init__(
26
- self,
27
- model: Sam,
28
- return_single_mask: bool,
29
- use_stability_score: bool = False,
30
- return_extra_metrics: bool = False,
31
- ) -> None:
32
- super().__init__()
33
- self.mask_decoder = model.mask_decoder
34
- self.model = model
35
- self.img_size = model.image_encoder.img_size
36
- self.return_single_mask = return_single_mask
37
- self.use_stability_score = use_stability_score
38
- self.stability_score_offset = 1.0
39
- self.return_extra_metrics = return_extra_metrics
40
-
41
- @staticmethod
42
- def resize_longest_image_size(
43
- input_image_size: torch.Tensor, longest_side: int
44
- ) -> torch.Tensor:
45
- input_image_size = input_image_size.to(torch.float32)
46
- scale = longest_side / torch.max(input_image_size)
47
- transformed_size = scale * input_image_size
48
- transformed_size = torch.floor(transformed_size + 0.5).to(torch.int64)
49
- return transformed_size
50
-
51
- def _embed_points(self, point_coords: torch.Tensor, point_labels: torch.Tensor) -> torch.Tensor:
52
- point_coords = point_coords + 0.5
53
- point_coords = point_coords / self.img_size
54
- point_embedding = self.model.prompt_encoder.pe_layer._pe_encoding(point_coords)
55
- point_labels = point_labels.unsqueeze(-1).expand_as(point_embedding)
56
-
57
- point_embedding = point_embedding * (point_labels != -1)
58
- point_embedding = point_embedding + self.model.prompt_encoder.not_a_point_embed.weight * (
59
- point_labels == -1
60
- )
61
-
62
- for i in range(self.model.prompt_encoder.num_point_embeddings):
63
- point_embedding = point_embedding + self.model.prompt_encoder.point_embeddings[
64
- i
65
- ].weight * (point_labels == i)
66
-
67
- return point_embedding
68
-
69
- def _embed_masks(self, input_mask: torch.Tensor, has_mask_input: torch.Tensor) -> torch.Tensor:
70
- mask_embedding = has_mask_input * self.model.prompt_encoder.mask_downscaling(input_mask)
71
- mask_embedding = mask_embedding + (
72
- 1 - has_mask_input
73
- ) * self.model.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1)
74
- return mask_embedding
75
-
76
- def mask_postprocessing(self, masks: torch.Tensor, orig_im_size: torch.Tensor) -> torch.Tensor:
77
- masks = F.interpolate(
78
- masks,
79
- size=(self.img_size, self.img_size),
80
- mode="bilinear",
81
- align_corners=False,
82
- )
83
-
84
- prepadded_size = self.resize_longest_image_size(orig_im_size, self.img_size)
85
- masks = masks[..., : int(prepadded_size[0]), : int(prepadded_size[1])]
86
-
87
- orig_im_size = orig_im_size.to(torch.int64)
88
- h, w = orig_im_size[0], orig_im_size[1]
89
- masks = F.interpolate(masks, size=(h, w), mode="bilinear", align_corners=False)
90
- return masks
91
-
92
- def select_masks(
93
- self, masks: torch.Tensor, iou_preds: torch.Tensor, num_points: int
94
- ) -> Tuple[torch.Tensor, torch.Tensor]:
95
- # Determine if we should return the multiclick mask or not from the number of points.
96
- # The reweighting is used to avoid control flow.
97
- score_reweight = torch.tensor(
98
- [[1000] + [0] * (self.model.mask_decoder.num_mask_tokens - 1)]
99
- ).to(iou_preds.device)
100
- score = iou_preds + (num_points - 2.5) * score_reweight
101
- best_idx = torch.argmax(score, dim=1)
102
- masks = masks[torch.arange(masks.shape[0]), best_idx, :, :].unsqueeze(1)
103
- iou_preds = iou_preds[torch.arange(masks.shape[0]), best_idx].unsqueeze(1)
104
-
105
- return masks, iou_preds
106
-
107
- @torch.no_grad()
108
- def forward(
109
- self,
110
- image_embeddings: torch.Tensor,
111
- point_coords: torch.Tensor,
112
- point_labels: torch.Tensor,
113
- mask_input: torch.Tensor,
114
- has_mask_input: torch.Tensor,
115
- orig_im_size: torch.Tensor,
116
- ):
117
- sparse_embedding = self._embed_points(point_coords, point_labels)
118
- dense_embedding = self._embed_masks(mask_input, has_mask_input)
119
-
120
- masks, scores = self.model.mask_decoder.predict_masks(
121
- image_embeddings=image_embeddings,
122
- image_pe=self.model.prompt_encoder.get_dense_pe(),
123
- sparse_prompt_embeddings=sparse_embedding,
124
- dense_prompt_embeddings=dense_embedding,
125
- )
126
-
127
- if self.use_stability_score:
128
- scores = calculate_stability_score(
129
- masks, self.model.mask_threshold, self.stability_score_offset
130
- )
131
-
132
- if self.return_single_mask:
133
- masks, scores = self.select_masks(masks, scores, point_coords.shape[1])
134
-
135
- upscaled_masks = self.mask_postprocessing(masks, orig_im_size)
136
-
137
- if self.return_extra_metrics:
138
- stability_scores = calculate_stability_score(
139
- upscaled_masks, self.model.mask_threshold, self.stability_score_offset
140
- )
141
- areas = (upscaled_masks > self.model.mask_threshold).sum(-1).sum(-1)
142
- return upscaled_masks, scores, stability_scores, areas, masks
143
-
144
- return upscaled_masks, scores, masks
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/segment_anything/utils/transforms.py DELETED
@@ -1,102 +0,0 @@
1
- # Copyright (c) Meta Platforms, Inc. and affiliates.
2
- # All rights reserved.
3
-
4
- # This source code is licensed under the license found in the
5
- # LICENSE file in the root directory of this source tree.
6
-
7
- import numpy as np
8
- import torch
9
- from torch.nn import functional as F
10
- from torchvision.transforms.functional import resize, to_pil_image # type: ignore
11
-
12
- from copy import deepcopy
13
- from typing import Tuple
14
-
15
-
16
- class ResizeLongestSide:
17
- """
18
- Resizes images to longest side 'target_length', as well as provides
19
- methods for resizing coordinates and boxes. Provides methods for
20
- transforming both numpy array and batched torch tensors.
21
- """
22
-
23
- def __init__(self, target_length: int) -> None:
24
- self.target_length = target_length
25
-
26
- def apply_image(self, image: np.ndarray) -> np.ndarray:
27
- """
28
- Expects a numpy array with shape HxWxC in uint8 format.
29
- """
30
- target_size = self.get_preprocess_shape(image.shape[0], image.shape[1], self.target_length)
31
- return np.array(resize(to_pil_image(image), target_size))
32
-
33
- def apply_coords(self, coords: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray:
34
- """
35
- Expects a numpy array of length 2 in the final dimension. Requires the
36
- original image size in (H, W) format.
37
- """
38
- old_h, old_w = original_size
39
- new_h, new_w = self.get_preprocess_shape(
40
- original_size[0], original_size[1], self.target_length
41
- )
42
- coords = deepcopy(coords).astype(float)
43
- coords[..., 0] = coords[..., 0] * (new_w / old_w)
44
- coords[..., 1] = coords[..., 1] * (new_h / old_h)
45
- return coords
46
-
47
- def apply_boxes(self, boxes: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray:
48
- """
49
- Expects a numpy array shape Bx4. Requires the original image size
50
- in (H, W) format.
51
- """
52
- boxes = self.apply_coords(boxes.reshape(-1, 2, 2), original_size)
53
- return boxes.reshape(-1, 4)
54
-
55
- def apply_image_torch(self, image: torch.Tensor) -> torch.Tensor:
56
- """
57
- Expects batched images with shape BxCxHxW and float format. This
58
- transformation may not exactly match apply_image. apply_image is
59
- the transformation expected by the model.
60
- """
61
- # Expects an image in BCHW format. May not exactly match apply_image.
62
- target_size = self.get_preprocess_shape(image.shape[0], image.shape[1], self.target_length)
63
- return F.interpolate(
64
- image, target_size, mode="bilinear", align_corners=False, antialias=True
65
- )
66
-
67
- def apply_coords_torch(
68
- self, coords: torch.Tensor, original_size: Tuple[int, ...]
69
- ) -> torch.Tensor:
70
- """
71
- Expects a torch tensor with length 2 in the last dimension. Requires the
72
- original image size in (H, W) format.
73
- """
74
- old_h, old_w = original_size
75
- new_h, new_w = self.get_preprocess_shape(
76
- original_size[0], original_size[1], self.target_length
77
- )
78
- coords = deepcopy(coords).to(torch.float)
79
- coords[..., 0] = coords[..., 0] * (new_w / old_w)
80
- coords[..., 1] = coords[..., 1] * (new_h / old_h)
81
- return coords
82
-
83
- def apply_boxes_torch(
84
- self, boxes: torch.Tensor, original_size: Tuple[int, ...]
85
- ) -> torch.Tensor:
86
- """
87
- Expects a torch tensor with shape Bx4. Requires the original image
88
- size in (H, W) format.
89
- """
90
- boxes = self.apply_coords_torch(boxes.reshape(-1, 2, 2), original_size)
91
- return boxes.reshape(-1, 4)
92
-
93
- @staticmethod
94
- def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int) -> Tuple[int, int]:
95
- """
96
- Compute the output size given input size and target long side length.
97
- """
98
- scale = long_side_length * 1.0 / max(oldh, oldw)
99
- newh, neww = oldh * scale, oldw * scale
100
- neww = int(neww + 0.5)
101
- newh = int(newh + 0.5)
102
- return (newh, neww)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
segment_anything/setup.cfg DELETED
@@ -1,11 +0,0 @@
1
- [isort]
2
- line_length=100
3
- multi_line_output=3
4
- include_trailing_comma=True
5
- known_standard_library=numpy,setuptools
6
- skip_glob=*/__init__.py
7
- known_myself=segment_anything
8
- known_third_party=matplotlib,cv2,torch,torchvision,pycocotools,onnx,black,isort
9
- no_lines_before=STDLIB,THIRDPARTY
10
- sections=FUTURE,STDLIB,THIRDPARTY,MYSELF,FIRSTPARTY,LOCALFOLDER
11
- default_section=FIRSTPARTY