+
+## Issues
+We use GitHub issues to track public bugs. Please ensure your description is
+clear and has sufficient instructions to be able to reproduce the issue.
+
+Meta has a [bounty program](https://www.facebook.com/whitehat/) for the safe
+disclosure of security bugs. In those cases, please go through the process
+outlined on that page and do not file a public issue.
+
+## License
+By contributing to DINOv2, you agree that your contributions will be licensed
+under the LICENSE file in the root directory of this source tree.
diff --git a/torchhub/facebookresearch_dinov2_main/LICENSE b/torchhub/facebookresearch_dinov2_main/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..a115f899f8d09ef3b1def4a16c7bae1a0bd50fbe
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/LICENSE
@@ -0,0 +1,400 @@
+
+Attribution-NonCommercial 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+ Considerations for licensors: Our public licenses are
+ intended for use by those authorized to give the public
+ permission to use material in ways otherwise restricted by
+ copyright and certain other rights. Our licenses are
+ irrevocable. Licensors should read and understand the terms
+ and conditions of the license they choose before applying it.
+ Licensors should also secure all rights necessary before
+ applying our licenses so that the public can reuse the
+ material as expected. Licensors should clearly mark any
+ material not subject to the license. This includes other CC-
+ licensed material, or material used under an exception or
+ limitation to copyright. More considerations for licensors:
+ wiki.creativecommons.org/Considerations_for_licensors
+
+ Considerations for the public: By using one of our public
+ licenses, a licensor grants the public permission to use the
+ licensed material under specified terms and conditions. If
+ the licensor's permission is not necessary for any reason--for
+ example, because of any applicable exception or limitation to
+ copyright--then that use is not regulated by the license. Our
+ licenses grant only permissions under copyright and certain
+ other rights that a licensor has authority to grant. Use of
+ the licensed material may still be restricted for other
+ reasons, including because others have copyright or other
+ rights in the material. A licensor may make special requests,
+ such as asking that all changes be marked or described.
+ Although not required by our licenses, you are encouraged to
+ respect those requests where reasonable. More_considerations
+ for the public:
+ wiki.creativecommons.org/Considerations_for_licensees
+
+=======================================================================
+
+Creative Commons Attribution-NonCommercial 4.0 International Public
+License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial 4.0 International Public License ("Public
+License"). To the extent this Public License may be interpreted as a
+contract, You are granted the Licensed Rights in consideration of Your
+acceptance of these terms and conditions, and the Licensor grants You
+such rights in consideration of benefits the Licensor receives from
+making the Licensed Material available under these terms and
+conditions.
+
+Section 1 -- Definitions.
+
+ a. Adapted Material means material subject to Copyright and Similar
+ Rights that is derived from or based upon the Licensed Material
+ and in which the Licensed Material is translated, altered,
+ arranged, transformed, or otherwise modified in a manner requiring
+ permission under the Copyright and Similar Rights held by the
+ Licensor. For purposes of this Public License, where the Licensed
+ Material is a musical work, performance, or sound recording,
+ Adapted Material is always produced where the Licensed Material is
+ synched in timed relation with a moving image.
+
+ b. Adapter's License means the license You apply to Your Copyright
+ and Similar Rights in Your contributions to Adapted Material in
+ accordance with the terms and conditions of this Public License.
+
+ c. Copyright and Similar Rights means copyright and/or similar rights
+ closely related to copyright including, without limitation,
+ performance, broadcast, sound recording, and Sui Generis Database
+ Rights, without regard to how the rights are labeled or
+ categorized. For purposes of this Public License, the rights
+ specified in Section 2(b)(1)-(2) are not Copyright and Similar
+ Rights.
+ d. Effective Technological Measures means those measures that, in the
+ absence of proper authority, may not be circumvented under laws
+ fulfilling obligations under Article 11 of the WIPO Copyright
+ Treaty adopted on December 20, 1996, and/or similar international
+ agreements.
+
+ e. Exceptions and Limitations means fair use, fair dealing, and/or
+ any other exception or limitation to Copyright and Similar Rights
+ that applies to Your use of the Licensed Material.
+
+ f. Licensed Material means the artistic or literary work, database,
+ or other material to which the Licensor applied this Public
+ License.
+
+ g. Licensed Rights means the rights granted to You subject to the
+ terms and conditions of this Public License, which are limited to
+ all Copyright and Similar Rights that apply to Your use of the
+ Licensed Material and that the Licensor has authority to license.
+
+ h. Licensor means the individual(s) or entity(ies) granting rights
+ under this Public License.
+
+ i. NonCommercial means not primarily intended for or directed towards
+ commercial advantage or monetary compensation. For purposes of
+ this Public License, the exchange of the Licensed Material for
+ other material subject to Copyright and Similar Rights by digital
+ file-sharing or similar means is NonCommercial provided there is
+ no payment of monetary compensation in connection with the
+ exchange.
+
+ j. Share means to provide material to the public by any means or
+ process that requires permission under the Licensed Rights, such
+ as reproduction, public display, public performance, distribution,
+ dissemination, communication, or importation, and to make material
+ available to the public including in ways that members of the
+ public may access the material from a place and at a time
+ individually chosen by them.
+
+ k. Sui Generis Database Rights means rights other than copyright
+ resulting from Directive 96/9/EC of the European Parliament and of
+ the Council of 11 March 1996 on the legal protection of databases,
+ as amended and/or succeeded, as well as other essentially
+ equivalent rights anywhere in the world.
+
+ l. You means the individual or entity exercising the Licensed Rights
+ under this Public License. Your has a corresponding meaning.
+
+Section 2 -- Scope.
+
+ a. License grant.
+
+ 1. Subject to the terms and conditions of this Public License,
+ the Licensor hereby grants You a worldwide, royalty-free,
+ non-sublicensable, non-exclusive, irrevocable license to
+ exercise the Licensed Rights in the Licensed Material to:
+
+ a. reproduce and Share the Licensed Material, in whole or
+ in part, for NonCommercial purposes only; and
+
+ b. produce, reproduce, and Share Adapted Material for
+ NonCommercial purposes only.
+
+ 2. Exceptions and Limitations. For the avoidance of doubt, where
+ Exceptions and Limitations apply to Your use, this Public
+ License does not apply, and You do not need to comply with
+ its terms and conditions.
+
+ 3. Term. The term of this Public License is specified in Section
+ 6(a).
+
+ 4. Media and formats; technical modifications allowed. The
+ Licensor authorizes You to exercise the Licensed Rights in
+ all media and formats whether now known or hereafter created,
+ and to make technical modifications necessary to do so. The
+ Licensor waives and/or agrees not to assert any right or
+ authority to forbid You from making technical modifications
+ necessary to exercise the Licensed Rights, including
+ technical modifications necessary to circumvent Effective
+ Technological Measures. For purposes of this Public License,
+ simply making modifications authorized by this Section 2(a)
+ (4) never produces Adapted Material.
+
+ 5. Downstream recipients.
+
+ a. Offer from the Licensor -- Licensed Material. Every
+ recipient of the Licensed Material automatically
+ receives an offer from the Licensor to exercise the
+ Licensed Rights under the terms and conditions of this
+ Public License.
+
+ b. No downstream restrictions. You may not offer or impose
+ any additional or different terms or conditions on, or
+ apply any Effective Technological Measures to, the
+ Licensed Material if doing so restricts exercise of the
+ Licensed Rights by any recipient of the Licensed
+ Material.
+
+ 6. No endorsement. Nothing in this Public License constitutes or
+ may be construed as permission to assert or imply that You
+ are, or that Your use of the Licensed Material is, connected
+ with, or sponsored, endorsed, or granted official status by,
+ the Licensor or others designated to receive attribution as
+ provided in Section 3(a)(1)(A)(i).
+
+ b. Other rights.
+
+ 1. Moral rights, such as the right of integrity, are not
+ licensed under this Public License, nor are publicity,
+ privacy, and/or other similar personality rights; however, to
+ the extent possible, the Licensor waives and/or agrees not to
+ assert any such rights held by the Licensor to the limited
+ extent necessary to allow You to exercise the Licensed
+ Rights, but not otherwise.
+
+ 2. Patent and trademark rights are not licensed under this
+ Public License.
+
+ 3. To the extent possible, the Licensor waives any right to
+ collect royalties from You for the exercise of the Licensed
+ Rights, whether directly or through a collecting society
+ under any voluntary or waivable statutory or compulsory
+ licensing scheme. In all other cases the Licensor expressly
+ reserves any right to collect such royalties, including when
+ the Licensed Material is used other than for NonCommercial
+ purposes.
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+ a. Attribution.
+
+ 1. If You Share the Licensed Material (including in modified
+ form), You must:
+
+ a. retain the following if it is supplied by the Licensor
+ with the Licensed Material:
+
+ i. identification of the creator(s) of the Licensed
+ Material and any others designated to receive
+ attribution, in any reasonable manner requested by
+ the Licensor (including by pseudonym if
+ designated);
+
+ ii. a copyright notice;
+
+ iii. a notice that refers to this Public License;
+
+ iv. a notice that refers to the disclaimer of
+ warranties;
+
+ v. a URI or hyperlink to the Licensed Material to the
+ extent reasonably practicable;
+
+ b. indicate if You modified the Licensed Material and
+ retain an indication of any previous modifications; and
+
+ c. indicate the Licensed Material is licensed under this
+ Public License, and include the text of, or the URI or
+ hyperlink to, this Public License.
+
+ 2. You may satisfy the conditions in Section 3(a)(1) in any
+ reasonable manner based on the medium, means, and context in
+ which You Share the Licensed Material. For example, it may be
+ reasonable to satisfy the conditions by providing a URI or
+ hyperlink to a resource that includes the required
+ information.
+
+ 3. If requested by the Licensor, You must remove any of the
+ information required by Section 3(a)(1)(A) to the extent
+ reasonably practicable.
+
+ 4. If You Share Adapted Material You produce, the Adapter's
+ License You apply must not prevent recipients of the Adapted
+ Material from complying with this Public License.
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+ a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+ to extract, reuse, reproduce, and Share all or a substantial
+ portion of the contents of the database for NonCommercial purposes
+ only;
+
+ b. if You include all or a substantial portion of the database
+ contents in a database in which You have Sui Generis Database
+ Rights, then the database in which You have Sui Generis Database
+ Rights (but not its individual contents) is Adapted Material; and
+
+ c. You must comply with the conditions in Section 3(a) if You Share
+ all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+ a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+ EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+ AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+ ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+ IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+ WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+ PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+ ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+ KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+ ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+ b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+ TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+ NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+ INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+ COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+ USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+ ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+ DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+ IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+ c. The disclaimer of warranties and limitation of liability provided
+ above shall be interpreted in a manner that, to the extent
+ possible, most closely approximates an absolute disclaimer and
+ waiver of all liability.
+
+Section 6 -- Term and Termination.
+
+ a. This Public License applies for the term of the Copyright and
+ Similar Rights licensed here. However, if You fail to comply with
+ this Public License, then Your rights under this Public License
+ terminate automatically.
+
+ b. Where Your right to use the Licensed Material has terminated under
+ Section 6(a), it reinstates:
+
+ 1. automatically as of the date the violation is cured, provided
+ it is cured within 30 days of Your discovery of the
+ violation; or
+
+ 2. upon express reinstatement by the Licensor.
+
+ For the avoidance of doubt, this Section 6(b) does not affect any
+ right the Licensor may have to seek remedies for Your violations
+ of this Public License.
+
+ c. For the avoidance of doubt, the Licensor may also offer the
+ Licensed Material under separate terms or conditions or stop
+ distributing the Licensed Material at any time; however, doing so
+ will not terminate this Public License.
+
+ d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+ License.
+
+Section 7 -- Other Terms and Conditions.
+
+ a. The Licensor shall not be bound by any additional or different
+ terms or conditions communicated by You unless expressly agreed.
+
+ b. Any arrangements, understandings, or agreements regarding the
+ Licensed Material not stated herein are separate from and
+ independent of the terms and conditions of this Public License.
+
+Section 8 -- Interpretation.
+
+ a. For the avoidance of doubt, this Public License does not, and
+ shall not be interpreted to, reduce, limit, restrict, or impose
+ conditions on any use of the Licensed Material that could lawfully
+ be made without permission under this Public License.
+
+ b. To the extent possible, if any provision of this Public License is
+ deemed unenforceable, it shall be automatically reformed to the
+ minimum extent necessary to make it enforceable. If the provision
+ cannot be reformed, it shall be severed from this Public License
+ without affecting the enforceability of the remaining terms and
+ conditions.
+
+ c. No term or condition of this Public License will be waived and no
+ failure to comply consented to unless expressly agreed to by the
+ Licensor.
+
+ d. Nothing in this Public License constitutes or may be interpreted
+ as a limitation upon, or waiver of, any privileges and immunities
+ that apply to the Licensor or You, including from the legal
+ processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the “Licensor.” The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.
diff --git a/torchhub/facebookresearch_dinov2_main/MODEL_CARD.md b/torchhub/facebookresearch_dinov2_main/MODEL_CARD.md
new file mode 100644
index 0000000000000000000000000000000000000000..5cd35748eb3c5d8f607f83ff068367a0102117c5
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/MODEL_CARD.md
@@ -0,0 +1,201 @@
+# Model Card for DINOv2-S/B/L/g
+
+These are Vision Transformer models trained following the method described in the paper:
+"DINOv2: Learning Robust Visual Features without Supervision"
+
+We provide 4 models: 1 ViT-g trained from scratch, and 3 ViT-S/B/L models distilled from the ViT-g.
+
+## Model Details
+The model takes an image as input and returns a class token and patch tokens.
+
+The embedding dimension is:
+- 384 for ViT-S.
+- 768 for ViT-B.
+- 1024 for ViT-L.
+- 1536 for ViT-g.
+
+The models follow a Transformer architecture, with a patch size of 14.
+
+For a 224x224 image, this results in 1 class token + 256 patch tokens.
+
+The models can accept larger images provided the image shapes are multiples of the patch size (14).
+If this condition is not verified, the model will crop to the closest smaller multiple of the patch size.
+
+### Model Description
+
+- **Developed by:** Meta AI
+- **Model type:** Vision Transformer
+- **License:** CC-BY-NC
+
+- **Repository:** https://github.com/facebookresearch/dinov2
+- **Paper:** https://arxiv.org/abs/2304.07193
+- **Demo:** https://dinov2.metademolab.com/
+
+## Uses
+
+The models are vision backbones providing multi-purpose features for downstream tasks.
+
+### Direct Use
+
+The models can be used without fine-tuning, with downstream classifiers as simple as linear layers, to obtain competitive results:
+- on depth estimation, semantic segmentation, using linear layers.
+- on image classification, using k-NN classifiers on the class token.
+- on image classification, with logistic regression classifiers applied on the class token.
+- on image classification, with a linear layer applied on the class token and the average of the patch tokens.
+- on image retrieval using nearest neighbors.
+
+### Downstream Use
+
+It is technically possible to perform fine-tuning on the models, for small gains (we measured +2% on ImageNet-1k classification).
+We recommend keeping this as a very last step and only when necessary, as the features already provide good performance out-of-the-box.
+
+## Bias, Risks, and Limitations
+
+Despite improvements thanks to the training method not using annotations, we still observe significant biases in our models toward rich households from Western countries.
+
+### Recommendations
+
+We expect fine-tuning will increase the biases in the features produced by the model as they will be tuned to the fine-tuning labels.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+```python
+import torch
+dinov2_vits14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')
+dinov2_vitb14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14')
+dinov2_vitl14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14')
+dinov2_vitg14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14')
+```
+
+## Training Details
+
+### Training Data
+
+- **Training data:** LVD-142M (see paper)
+- **Training regime:** fp16 using PyTorch-FSDP mixed-precision.
+
+### Training Procedure
+
+- **Training objective:**
+ - DINO self-distillation loss with multi-crop
+ - iBOT masked-image modeling loss
+ - KoLeo regularization on [CLS] tokens
+- **Architectures:**
+ - ViT-S (21M params): Patch size 14, embedding dimension 384, 6 heads, MLP FFN
+ - ViT-B (86M params): Patch size 14, embedding dimension 768, 12 heads, MLP FFN
+ - ViT-L (0.3B params): Patch size 14, embedding dimension 1024, 16 heads, MLP FFN
+ - ViT-g (1.1B params): Patch size 14, embedding dimension 1536, 24 heads, SwiGLU FFN
+- **Distillation:**
+ - Distillation follows the standard DINOv2 pretraining procedure, except the teacher is a pretrained ViT-g, frozen.
+
+## Evaluation
+
+We refer users to the associated paper for the evaluation protocols.
+
+
+
+ | model |
+ ImageNet-1k |
+ NYU-Depth v2 |
+ SUN-RGBD |
+ ADE20k |
+ iNaturalist 2018 |
+ Oxford-H |
+
+
+ | task |
+ classif. (acc) |
+ classif. (acc) |
+ classif. V2 (acc) |
+ depth (RMSE) |
+ depth (RMSE) |
+ segm. (mAP) |
+ classif. (acc) |
+ retrieval (mAP) |
+
+
+
+ | k-NN |
+ linear |
+ linear |
+ linear
4 layers |
+ NYU-D transfer |
+ multiscale |
+ linear |
+ nearest neighbor |
+
+
+ | ViT-S/14 |
+ 79.0% |
+ 81.1% |
+ 70.8% |
+ 0.417 |
+ 0.431 |
+ 47.2 |
+ 69.5% |
+ 43.2 |
+
+
+ | ViT-B/14 |
+ 82.1% |
+ 84.5% |
+ 74.9% |
+ 0.362 |
+ 0.400 |
+ 51.3 |
+ 76.3% |
+ 49.5 |
+
+
+ | ViT-L/14 |
+ 83.5% |
+ 86.3% |
+ 77.6% |
+ 0.333 |
+ 0.396 |
+ 53.1 |
+ 79.8% |
+ 54.0 |
+
+
+ | ViT-g/14 |
+ 83.5% |
+ 86.5% |
+ 78.4% |
+ 0.298 |
+ 0.362 |
+ 53.0 |
+ 81.6% |
+ 52.3 |
+
+
+
+## Environmental Impact
+
+- **Hardware Type:** Nvidia A100
+- **Hours used:** 22,000 for ViT-g, 4,500 for ViT-S distillation, 5,300 for ViT-B distillation, 8,000 for ViT-L distillation
+- **Cloud Provider:** Private infra
+- **Compute Region:** USA
+- **Carbon Emitted:** 7t CO2eq
+
+#### Hardware
+
+Nvidia A100 GPUs
+
+#### Software
+
+PyTorch 2.0,
+xFormers 0.0.18
+
+**BibTeX**
+
+```
+@misc{oquab2023dinov2,
+ title={DINOv2: Learning Robust Visual Features without Supervision},
+ author={Oquab, Maxime and Darcet, Timothée and Moutakanni, Theo and Vo, Huy and Szafraniec, Marc and Khalidov, Vasil and Fernandez, Pierre and Haziza, Daniel and Massa, Francisco and El-Nouby, Alaaeldin and Howes, Russell and Huang, Po-Yao and Xu, Hu and Sharma, Vasu and Li, Shang-Wen and Galuba, Wojciech and Rabbat, Mike and Assran, Mido and Ballas, Nicolas and Synnaeve, Gabriel and Misra, Ishan and Jegou, Herve and Mairal, Julien and Labatut, Patrick and Joulin, Armand and Bojanowski, Piotr},
+ journal={arXiv:2304.07193},
+ year={2023}
+}
+```
diff --git a/torchhub/facebookresearch_dinov2_main/README.md b/torchhub/facebookresearch_dinov2_main/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..96453e567dee10148be83b5e92d91f347f8521d5
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/README.md
@@ -0,0 +1,277 @@
+# DINOv2: Learning Robust Visual Features without Supervision
+
+**[Meta AI Research, FAIR](https://ai.facebook.com/research/)**
+
+Maxime Oquab,
+Timothée Darcet,
+Théo Moutakanni,
+Huy V. Vo,
+Marc Szafraniec,
+Vasil Khalidov,
+Patrick Labatut,
+Armand Joulin,
+Piotr Bojanowski
+
+[[`Paper`](https://arxiv.org/abs/2304.07193)] [[`Blog`](https://ai.facebook.com/blog/dino-v2-computer-vision-self-supervised-learning/)] [[`Demo`](https://dinov2.metademolab.com)] [[`BibTeX`](#citing-dinov2)]
+
+PyTorch implementation and pretrained models for DINOv2. For details, see the paper: **[DINOv2: Learning Robust Visual Features without Supervision](https://arxiv.org/abs/2304.07193)**.
+
+DINOv2 models produce high-performance visual features that can be directly employed with classifiers as simple as linear layers on a variety of computer vision tasks; these visual features are robust and perform well across domains without any requirement for fine-tuning. The models were pretrained on a dataset of 142 M images without using any labels or annotations.
+
+https://github.com/facebookresearch/dinov2/assets/60359573/f168823e-7922-415a-b429-578badf5c356
+
+
+ Visualization of the three first principal components of the patch features of all frames, mapped to RGB values.
+
+
+## Pretrained models
+
+
+
+ | model |
+ # of
params |
+ ImageNet
k-NN |
+ ImageNet
linear |
+ download |
+
+
+ | ViT-S/14 distilled |
+ 21 M |
+ 79.0% |
+ 81.1% |
+ backbone only |
+
+
+ | ViT-B/14 distilled |
+ 86 M |
+ 82.1% |
+ 84.5% |
+ backbone only |
+
+
+ | ViT-L/14 distilled |
+ 300 M |
+ 83.5% |
+ 86.3% |
+ backbone only |
+
+
+ | ViT-g/14 |
+ 1,100 M |
+ 83.5% |
+ 86.5% |
+ backbone only |
+
+
+
+### Pretrained models via PyTorch Hub
+
+Please follow the instructions [here](https://pytorch.org/get-started/locally/) to install PyTorch (the only required dependency for loading the model). Installing PyTorch with CUDA support is strongly recommended.
+
+A corresponding [model card](MODEL_CARD.md) is included in the repository.
+
+```python
+import torch
+
+dinov2_vits14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')
+dinov2_vitb14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14')
+dinov2_vitl14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14')
+dinov2_vitg14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14')
+```
+
+## Installation
+
+The training and evaluation code requires PyTorch 2.0 and [xFormers](https://github.com/facebookresearch/xformers) 0.0.18 as well as a number of other 3rd party packages. Note that the code has only been tested with the specified versions and also expects a Linux environment. To setup all the required dependencies for training and evaluation, please follow the instructions below:
+
+*[conda](https://docs.conda.io/projects/conda/en/latest/user-guide/getting-started.html)* **(Recommended)** - Clone the repository and then create and activate a `dinov2` conda environment using the provided environment definition:
+
+```shell
+conda env create -f conda.yaml
+conda activate dinov2
+```
+
+*[pip](https://pip.pypa.io/en/stable/getting-started/)* - Clone the repository and then use the provided `requirements.txt` to install the dependencies:
+
+```shell
+pip install -r requirements.txt
+```
+
+## Data preparation
+
+### ImageNet-1k
+
+The root directory of the dataset should hold the following contents:
+
+- `/test/ILSVRC2012_test_00000001.JPEG`
+- `/test/[..]`
+- `/test/ILSVRC2012_test_00100000.JPEG`
+- `/train/n01440764/n01440764_10026.JPEG`
+- `/train/[...]`
+- `/train/n15075141/n15075141_9993.JPEG`
+- `/val/n01440764/ILSVRC2012_val_00000293.JPEG`
+- `/val/[...]`
+- `/val/n15075141/ILSVRC2012_val_00049174.JPEG`
+- `/labels.txt`
+
+The provided dataset implementation expects a few additional metadata files to be present under the extra directory:
+
+- `/class-ids-TRAIN.npy`
+- `/class-ids-VAL.npy`
+- `/class-names-TRAIN.npy`
+- `/class-names-VAL.npy`
+- `/entries-TEST.npy`
+- `/entries-TRAIN.npy`
+- `/entries-VAL.npy`
+
+These metadata files can be generated (once) with the following lines of Python code:
+
+```python
+from dinov2.data.datasets import ImageNet
+
+for split in ImageNet.Split:
+ dataset = ImageNet(split=split, root="", extra="")
+ dataset.dump_extra()
+```
+
+Note that the root and extra directories do not have to be distinct directories.
+
+### ImageNet-22k
+
+Please adapt the [dataset class](dinov2/data/datasets/image_net_22k.py) to match your local setup.
+
+
+
+:warning: To execute the commands provided in the next sections for training and evaluation, the `dinov2` package should be included in the Python module search path, i.e. simply prefix the command to run with `PYTHONPATH=.`.
+
+## Training
+
+### Fast setup: training DINOv2 ViT-L/16 on ImageNet-1k
+
+Run DINOv2 training on 4 A100-80GB nodes (32 GPUs) in a SLURM cluster environment with submitit:
+
+```shell
+python dinov2/run/train/train.py \
+ --nodes 4 \
+ --config-file dinov2/configs/train/vitl16_short.yaml \
+ --output-dir \
+ train.dataset_path=ImageNet:split=TRAIN:root=:extra=
+```
+
+Training time is approximately 1 day and the resulting checkpoint should reach 81.6% on k-NN eval and 82.9% on linear eval.
+
+The training code saves the weights of the teacher in the `eval` folder every 12500 iterations for evaluation.
+
+### Long setup: training DINOv2 ViT-L/14 on ImageNet-22k
+
+Run DINOv2 training on 12 A100-80GB nodes (96 GPUs) in a SLURM cluster environment with submitit:
+
+```shell
+python dinov2/run/train/train.py \
+ --nodes 12 \
+ --config-file dinov2/configs/train/vitl14.yaml \
+ --output-dir \
+ train.dataset_path=ImageNet22k:root=:extra=
+```
+
+Training time is approximately 3.3 days and the resulting checkpoint should reach 82.0% on k-NN eval and 84.5% on linear eval.
+
+The training code saves the weights of the teacher in the `eval` folder every 12500 iterations for evaluation.
+
+
+## Evaluation
+
+The training code regularly saves the teacher weights. In order to evaluate the model, run the following evaluation on a single node:
+
+### k-NN classification on ImageNet-1k
+
+```shell
+python dinov2/run/eval/knn.py \
+ --config-file /config.yaml \
+ --pretrained-weights /eval/training_24999/teacher_checkpoint.pth \
+ --output-dir /eval/training_24999/knn \
+ --train-dataset ImageNet:split=TRAIN:root=:extra= \
+ --val-dataset ImageNet:split=VAL:root=:extra=
+```
+
+### Logistic regression classification on ImageNet-1k
+
+```shell
+python dinov2/run/eval/log_regression.py \
+ --config-file /config.yaml \
+ --pretrained-weights /eval/training_24999/teacher_checkpoint.pth \
+ --output-dir /eval/training_24999/logreg \
+ --train-dataset ImageNet:split=TRAIN:root=:extra= \
+ --val-dataset ImageNet:split=VAL:root=:extra=
+```
+
+### Linear classification with data augmentation on ImageNet-1k
+
+```shell
+python dinov2/run/eval/linear.py \
+ --config-file /config.yaml \
+ --pretrained-weights /eval/training_24999/teacher_checkpoint.pth \
+ --output-dir /eval/training_24999/linear \
+ --train-dataset ImageNet:split=TRAIN:root=:extra= \
+ --val-dataset ImageNet:split=VAL:root=:extra=
+```
+
+We release the weights from evaluating the different models:
+
+
+
+The performance of the provided pretrained model weights can be evaluated as follows on ImageNet-1k:
+
+```shell
+python dinov2/run/eval/linear.py \
+ --config-file dinov2/configs/eval/vitg14_pretrain.yaml \
+ --pretrained-weights https://dl.fbaipublicfiles.com/dinov2/dinov2_vitg14/dinov2_vitg14_pretrain.pth \
+ --train-dataset ImageNet:split=TRAIN:root=:extra= \
+ --val-dataset ImageNet:split=VAL:root=:extra=
+```
+
+## License
+
+DINOv2 code and model weights are released under the CC-BY-NC 4.0 license. See [LICENSE](LICENSE) for additional details.
+
+## Contributing
+
+See [contributing](CONTRIBUTING.md) and the [code of conduct](CODE_OF_CONDUCT.md).
+
+## Citing DINOv2
+
+If you find this repository useful, please consider giving a star :star: and citation :t-rex::
+
+```
+@misc{oquab2023dinov2,
+ title={DINOv2: Learning Robust Visual Features without Supervision},
+ author={Oquab, Maxime and Darcet, Timothée and Moutakanni, Theo and Vo, Huy V. and Szafraniec, Marc and Khalidov, Vasil and Fernandez, Pierre and Haziza, Daniel and Massa, Francisco and El-Nouby, Alaaeldin and Howes, Russell and Huang, Po-Yao and Xu, Hu and Sharma, Vasu and Li, Shang-Wen and Galuba, Wojciech and Rabbat, Mike and Assran, Mido and Ballas, Nicolas and Synnaeve, Gabriel and Misra, Ishan and Jegou, Herve and Mairal, Julien and Labatut, Patrick and Joulin, Armand and Bojanowski, Piotr},
+ journal={arXiv:2304.07193},
+ year={2023}
+}
+```
diff --git a/torchhub/facebookresearch_dinov2_main/conda.yaml b/torchhub/facebookresearch_dinov2_main/conda.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..35dfc30adc275da51b58ff2340dd1d53d2cb9250
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/conda.yaml
@@ -0,0 +1,22 @@
+name: dinov2
+channels:
+ - defaults
+ - pytorch
+ - nvidia
+ - xformers
+ - conda-forge
+dependencies:
+ - python=3.9
+ - pytorch::pytorch=2.0.0
+ - pytorch::pytorch-cuda=11.7.0
+ - pytorch::torchvision=0.15.0
+ - omegaconf
+ - torchmetrics=0.10.3
+ - fvcore
+ - iopath
+ - xformers::xformers=0.0.18
+ - pip
+ - pip:
+ - git+https://github.com/facebookincubator/submitit
+ - --extra-index-url https://pypi.nvidia.com
+ - cuml-cu11
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5b4afb514783786adf76744f9b97f3e1db1d6081
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+__version__ = "0.0.1"
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/configs/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/configs/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..033c35660450afec6612adb342c7c30e1ccd15ee
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/configs/__init__.py
@@ -0,0 +1,23 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import pathlib
+
+from omegaconf import OmegaConf
+
+
+def load_config(config_name: str):
+ config_filename = config_name + ".yaml"
+ return OmegaConf.load(pathlib.Path(__file__).parent.resolve() / config_filename)
+
+
+dinov2_default_config = load_config("ssl_default_config")
+
+
+def load_and_merge_config(config_name: str):
+ default_config = OmegaConf.create(dinov2_default_config)
+ loaded_config = load_config(config_name)
+ return OmegaConf.merge(default_config, loaded_config)
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vitb14_pretrain.yaml b/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vitb14_pretrain.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..117d0f027ca26cd8ce6c010bb78d5a8fac42c70e
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vitb14_pretrain.yaml
@@ -0,0 +1,6 @@
+student:
+ arch: vit_base
+ patch_size: 14
+crops:
+ global_crops_size: 518 # this is to set up the position embeddings properly
+ local_crops_size: 98
\ No newline at end of file
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vitg14_pretrain.yaml b/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vitg14_pretrain.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..a96dd5b117b4d59ee210b65037821f1b3e3f16e3
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vitg14_pretrain.yaml
@@ -0,0 +1,7 @@
+student:
+ arch: vit_giant2
+ patch_size: 14
+ ffn_layer: swiglufused
+crops:
+ global_crops_size: 518 # this is to set up the position embeddings properly
+ local_crops_size: 98
\ No newline at end of file
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vitl14_pretrain.yaml b/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vitl14_pretrain.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..7a984548bd034f762d455419d7193917fa462dd8
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vitl14_pretrain.yaml
@@ -0,0 +1,6 @@
+student:
+ arch: vit_large
+ patch_size: 14
+crops:
+ global_crops_size: 518 # this is to set up the position embeddings properly
+ local_crops_size: 98
\ No newline at end of file
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vits14_pretrain.yaml b/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vits14_pretrain.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..afbdb4ba14f1c97130a25b579360f4d817cda495
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/configs/eval/vits14_pretrain.yaml
@@ -0,0 +1,6 @@
+student:
+ arch: vit_small
+ patch_size: 14
+crops:
+ global_crops_size: 518 # this is to set up the position embeddings properly
+ local_crops_size: 98
\ No newline at end of file
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/configs/ssl_default_config.yaml b/torchhub/facebookresearch_dinov2_main/dinov2/configs/ssl_default_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..a4ef04545ce9d6cc52b5179236008adc8a9bbda2
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/configs/ssl_default_config.yaml
@@ -0,0 +1,115 @@
+MODEL:
+ WEIGHTS: ''
+compute_precision:
+ grad_scaler: true
+ teacher:
+ backbone:
+ sharding_strategy: SHARD_GRAD_OP
+ mixed_precision:
+ param_dtype: fp16
+ reduce_dtype: fp16
+ buffer_dtype: fp32
+ dino_head:
+ sharding_strategy: SHARD_GRAD_OP
+ mixed_precision:
+ param_dtype: fp16
+ reduce_dtype: fp16
+ buffer_dtype: fp32
+ ibot_head:
+ sharding_strategy: SHARD_GRAD_OP
+ mixed_precision:
+ param_dtype: fp16
+ reduce_dtype: fp16
+ buffer_dtype: fp32
+ student:
+ backbone:
+ sharding_strategy: SHARD_GRAD_OP
+ mixed_precision:
+ param_dtype: fp16
+ reduce_dtype: fp16
+ buffer_dtype: fp32
+ dino_head:
+ sharding_strategy: SHARD_GRAD_OP
+ mixed_precision:
+ param_dtype: fp16
+ reduce_dtype: fp32
+ buffer_dtype: fp32
+ ibot_head:
+ sharding_strategy: SHARD_GRAD_OP
+ mixed_precision:
+ param_dtype: fp16
+ reduce_dtype: fp32
+ buffer_dtype: fp32
+dino:
+ loss_weight: 1.0
+ head_n_prototypes: 65536
+ head_bottleneck_dim: 256
+ head_nlayers: 3
+ head_hidden_dim: 2048
+ koleo_loss_weight: 0.1
+ibot:
+ loss_weight: 1.0
+ mask_sample_probability: 0.5
+ mask_ratio_min_max:
+ - 0.1
+ - 0.5
+ separate_head: false
+ head_n_prototypes: 65536
+ head_bottleneck_dim: 256
+ head_nlayers: 3
+ head_hidden_dim: 2048
+train:
+ batch_size_per_gpu: 64
+ dataset_path: ImageNet:split=TRAIN
+ output_dir: .
+ saveckp_freq: 20
+ seed: 0
+ num_workers: 10
+ OFFICIAL_EPOCH_LENGTH: 1250
+ cache_dataset: true
+ centering: "centering" # or "sinkhorn_knopp"
+student:
+ arch: vit_large
+ patch_size: 16
+ drop_path_rate: 0.3
+ layerscale: 1.0e-05
+ drop_path_uniform: true
+ pretrained_weights: ''
+ ffn_layer: "mlp"
+ block_chunks: 0
+ qkv_bias: true
+ proj_bias: true
+ ffn_bias: true
+teacher:
+ momentum_teacher: 0.992
+ final_momentum_teacher: 1
+ warmup_teacher_temp: 0.04
+ teacher_temp: 0.07
+ warmup_teacher_temp_epochs: 30
+optim:
+ epochs: 100
+ weight_decay: 0.04
+ weight_decay_end: 0.4
+ base_lr: 0.004 # learning rate for a batch size of 1024
+ lr: 0. # will be set after applying scaling rule
+ warmup_epochs: 10
+ min_lr: 1.0e-06
+ clip_grad: 3.0
+ freeze_last_layer_epochs: 1
+ scaling_rule: sqrt_wrt_1024
+ patch_embed_lr_mult: 0.2
+ layerwise_decay: 0.9
+ adamw_beta1: 0.9
+ adamw_beta2: 0.999
+crops:
+ global_crops_scale:
+ - 0.32
+ - 1.0
+ local_crops_number: 8
+ local_crops_scale:
+ - 0.05
+ - 0.32
+ global_crops_size: 224
+ local_crops_size: 96
+evaluation:
+ eval_period_iterations: 12500
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/configs/train/vitg14.yaml b/torchhub/facebookresearch_dinov2_main/dinov2/configs/train/vitg14.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..d05cf0d59e07ac6e4a2b0f9bdcb6131d7c508962
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/configs/train/vitg14.yaml
@@ -0,0 +1,26 @@
+dino:
+ head_n_prototypes: 131072
+ head_bottleneck_dim: 384
+ibot:
+ separate_head: true
+ head_n_prototypes: 131072
+train:
+ batch_size_per_gpu: 12
+ dataset_path: ImageNet22k
+ centering: sinkhorn_knopp
+student:
+ arch: vit_giant2
+ patch_size: 14
+ drop_path_rate: 0.4
+ ffn_layer: swiglufused
+ block_chunks: 4
+teacher:
+ momentum_teacher: 0.994
+optim:
+ epochs: 500
+ weight_decay_end: 0.2
+ base_lr: 2.0e-04 # learning rate for a batch size of 1024
+ warmup_epochs: 80
+ layerwise_decay: 1.0
+crops:
+ local_crops_size: 98
\ No newline at end of file
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/configs/train/vitl14.yaml b/torchhub/facebookresearch_dinov2_main/dinov2/configs/train/vitl14.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..d9b491dcc6a522c71328fc2933dd0501123c8f6b
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/configs/train/vitl14.yaml
@@ -0,0 +1,26 @@
+dino:
+ head_n_prototypes: 131072
+ head_bottleneck_dim: 384
+ibot:
+ separate_head: true
+ head_n_prototypes: 131072
+train:
+ batch_size_per_gpu: 32
+ dataset_path: ImageNet22k
+ centering: sinkhorn_knopp
+student:
+ arch: vit_large
+ patch_size: 14
+ drop_path_rate: 0.4
+ ffn_layer: swiglufused
+ block_chunks: 4
+teacher:
+ momentum_teacher: 0.994
+optim:
+ epochs: 500
+ weight_decay_end: 0.2
+ base_lr: 2.0e-04 # learning rate for a batch size of 1024
+ warmup_epochs: 80
+ layerwise_decay: 1.0
+crops:
+ local_crops_size: 98
\ No newline at end of file
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/configs/train/vitl16_short.yaml b/torchhub/facebookresearch_dinov2_main/dinov2/configs/train/vitl16_short.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..3e7e72864c92175a1354142ac1d64da8070d1e5e
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/configs/train/vitl16_short.yaml
@@ -0,0 +1,6 @@
+# this corresponds to the default config
+train:
+ dataset_path: ImageNet:split=TRAIN
+ batch_size_per_gpu: 64
+student:
+ block_chunks: 4
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/data/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/data/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..357db5c542c5810391ba2bd45a60c13c01c3737a
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/data/__init__.py
@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .adapters import DatasetWithEnumeratedTargets
+from .loaders import make_data_loader, make_dataset, SamplerType
+from .collate import collate_data_and_cast
+from .masking import MaskingGenerator
+from .augmentations import DataAugmentationDINO
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/data/adapters.py b/torchhub/facebookresearch_dinov2_main/dinov2/data/adapters.py
new file mode 100644
index 0000000000000000000000000000000000000000..7dcbc68e046f03460d5867f1388d5380d9c6f603
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/data/adapters.py
@@ -0,0 +1,29 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Any, Tuple
+
+from torch.utils.data import Dataset
+
+
+class DatasetWithEnumeratedTargets(Dataset):
+ def __init__(self, dataset):
+ self._dataset = dataset
+
+ def get_image_data(self, index: int) -> bytes:
+ return self._dataset.get_image_data(index)
+
+ def get_target(self, index: int) -> Tuple[Any, int]:
+ target = self._dataset.get_target(index)
+ return (index, target)
+
+ def __getitem__(self, index: int) -> Tuple[Any, Tuple[Any, int]]:
+ image, target = self._dataset[index]
+ target = index if target is None else target
+ return image, (index, target)
+
+ def __len__(self) -> int:
+ return len(self._dataset)
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/data/augmentations.py b/torchhub/facebookresearch_dinov2_main/dinov2/data/augmentations.py
new file mode 100644
index 0000000000000000000000000000000000000000..7ca28cb59a4de2566a6c9ef9c301cbbb4e54b5ee
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/data/augmentations.py
@@ -0,0 +1,119 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+
+from torchvision import transforms
+
+from .transforms import (
+ GaussianBlur,
+ make_normalize_transform,
+)
+
+
+logger = logging.getLogger("dinov2")
+
+
+class DataAugmentationDINO(object):
+ def __init__(
+ self,
+ global_crops_scale,
+ local_crops_scale,
+ local_crops_number,
+ global_crops_size=224,
+ local_crops_size=96,
+ ):
+ self.global_crops_scale = global_crops_scale
+ self.local_crops_scale = local_crops_scale
+ self.local_crops_number = local_crops_number
+ self.global_crops_size = global_crops_size
+ self.local_crops_size = local_crops_size
+
+ logger.info("###################################")
+ logger.info("Using data augmentation parameters:")
+ logger.info(f"global_crops_scale: {global_crops_scale}")
+ logger.info(f"local_crops_scale: {local_crops_scale}")
+ logger.info(f"local_crops_number: {local_crops_number}")
+ logger.info(f"global_crops_size: {global_crops_size}")
+ logger.info(f"local_crops_size: {local_crops_size}")
+ logger.info("###################################")
+
+ # random resized crop and flip
+ self.geometric_augmentation_global = transforms.Compose(
+ [
+ transforms.RandomResizedCrop(
+ global_crops_size, scale=global_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC
+ ),
+ transforms.RandomHorizontalFlip(p=0.5),
+ ]
+ )
+
+ self.geometric_augmentation_local = transforms.Compose(
+ [
+ transforms.RandomResizedCrop(
+ local_crops_size, scale=local_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC
+ ),
+ transforms.RandomHorizontalFlip(p=0.5),
+ ]
+ )
+
+ # color distorsions / blurring
+ color_jittering = transforms.Compose(
+ [
+ transforms.RandomApply(
+ [transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1)],
+ p=0.8,
+ ),
+ transforms.RandomGrayscale(p=0.2),
+ ]
+ )
+
+ global_transfo1_extra = GaussianBlur(p=1.0)
+
+ global_transfo2_extra = transforms.Compose(
+ [
+ GaussianBlur(p=0.1),
+ transforms.RandomSolarize(threshold=128, p=0.2),
+ ]
+ )
+
+ local_transfo_extra = GaussianBlur(p=0.5)
+
+ # normalization
+ self.normalize = transforms.Compose(
+ [
+ transforms.ToTensor(),
+ make_normalize_transform(),
+ ]
+ )
+
+ self.global_transfo1 = transforms.Compose([color_jittering, global_transfo1_extra, self.normalize])
+ self.global_transfo2 = transforms.Compose([color_jittering, global_transfo2_extra, self.normalize])
+ self.local_transfo = transforms.Compose([color_jittering, local_transfo_extra, self.normalize])
+
+ def __call__(self, image):
+ output = {}
+
+ # global crops:
+ im1_base = self.geometric_augmentation_global(image)
+ global_crop_1 = self.global_transfo1(im1_base)
+
+ im2_base = self.geometric_augmentation_global(image)
+ global_crop_2 = self.global_transfo2(im2_base)
+
+ output["global_crops"] = [global_crop_1, global_crop_2]
+
+ # global crops for teacher:
+ output["global_crops_teacher"] = [global_crop_1, global_crop_2]
+
+ # local crops:
+ local_crops = [
+ self.local_transfo(self.geometric_augmentation_local(image)) for _ in range(self.local_crops_number)
+ ]
+ output["local_crops"] = local_crops
+ output["offsets"] = ()
+
+ return output
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/data/collate.py b/torchhub/facebookresearch_dinov2_main/dinov2/data/collate.py
new file mode 100644
index 0000000000000000000000000000000000000000..9f0d98906808ed326dff4486d95b3ec04f8a5e75
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/data/collate.py
@@ -0,0 +1,50 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import random
+
+
+def collate_data_and_cast(samples_list, mask_ratio_tuple, mask_probability, dtype, n_tokens=None, mask_generator=None):
+ # dtype = torch.half # TODO: Remove
+
+ n_global_crops = len(samples_list[0][0]["global_crops"])
+ n_local_crops = len(samples_list[0][0]["local_crops"])
+
+ collated_global_crops = torch.stack([s[0]["global_crops"][i] for i in range(n_global_crops) for s in samples_list])
+
+ collated_local_crops = torch.stack([s[0]["local_crops"][i] for i in range(n_local_crops) for s in samples_list])
+
+ B = len(collated_global_crops)
+ N = n_tokens
+ n_samples_masked = int(B * mask_probability)
+ probs = torch.linspace(*mask_ratio_tuple, n_samples_masked + 1)
+ upperbound = 0
+ masks_list = []
+ for i in range(0, n_samples_masked):
+ prob_min = probs[i]
+ prob_max = probs[i + 1]
+ masks_list.append(torch.BoolTensor(mask_generator(int(N * random.uniform(prob_min, prob_max)))))
+ upperbound += int(N * prob_max)
+ for i in range(n_samples_masked, B):
+ masks_list.append(torch.BoolTensor(mask_generator(0)))
+
+ random.shuffle(masks_list)
+
+ collated_masks = torch.stack(masks_list).flatten(1)
+ mask_indices_list = collated_masks.flatten().nonzero().flatten()
+
+ masks_weight = (1 / collated_masks.sum(-1).clamp(min=1.0)).unsqueeze(-1).expand_as(collated_masks)[collated_masks]
+
+ return {
+ "collated_global_crops": collated_global_crops.to(dtype),
+ "collated_local_crops": collated_local_crops.to(dtype),
+ "collated_masks": collated_masks,
+ "mask_indices_list": mask_indices_list,
+ "masks_weight": masks_weight,
+ "upperbound": upperbound,
+ "n_masked_patches": torch.full((1,), fill_value=mask_indices_list.shape[0], dtype=torch.long),
+ }
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/data/loaders.py b/torchhub/facebookresearch_dinov2_main/dinov2/data/loaders.py
new file mode 100644
index 0000000000000000000000000000000000000000..9fb6f25a0a3c3251b803f48d0a515aa0b9591226
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/data/loaders.py
@@ -0,0 +1,223 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from enum import Enum
+from typing import Any, Callable, List, Optional, TypeVar
+
+import torch
+from torch.utils.data import Sampler
+
+from .datasets import ImageNet, ImageNet22k
+from .samplers import EpochSampler, InfiniteSampler, ShardedInfiniteSampler
+
+
+logger = logging.getLogger("dinov2")
+
+
+class SamplerType(Enum):
+ DISTRIBUTED = 0
+ EPOCH = 1
+ INFINITE = 2
+ SHARDED_INFINITE = 3
+ SHARDED_INFINITE_NEW = 4
+
+
+def _make_bool_str(b: bool) -> str:
+ return "yes" if b else "no"
+
+
+def _make_sample_transform(image_transform: Optional[Callable] = None, target_transform: Optional[Callable] = None):
+ def transform(sample):
+ image, target = sample
+ if image_transform is not None:
+ image = image_transform(image)
+ if target_transform is not None:
+ target = target_transform(target)
+ return image, target
+
+ return transform
+
+
+def _parse_dataset_str(dataset_str: str):
+ tokens = dataset_str.split(":")
+
+ name = tokens[0]
+ kwargs = {}
+
+ for token in tokens[1:]:
+ key, value = token.split("=")
+ assert key in ("root", "extra", "split")
+ kwargs[key] = value
+
+ if name == "ImageNet":
+ class_ = ImageNet
+ if "split" in kwargs:
+ kwargs["split"] = ImageNet.Split[kwargs["split"]]
+ elif name == "ImageNet22k":
+ class_ = ImageNet22k
+ else:
+ raise ValueError(f'Unsupported dataset "{name}"')
+
+ return class_, kwargs
+
+
+def make_dataset(
+ *,
+ dataset_str: str,
+ transform: Optional[Callable] = None,
+ target_transform: Optional[Callable] = None,
+):
+ """
+ Creates a dataset with the specified parameters.
+
+ Args:
+ dataset_str: A dataset string description (e.g. ImageNet:split=TRAIN).
+ transform: A transform to apply to images.
+ target_transform: A transform to apply to targets.
+
+ Returns:
+ The created dataset.
+ """
+ logger.info(f'using dataset: "{dataset_str}"')
+
+ class_, kwargs = _parse_dataset_str(dataset_str)
+ dataset = class_(transform=transform, target_transform=target_transform, **kwargs)
+
+ logger.info(f"# of dataset samples: {len(dataset):,d}")
+
+ # Aggregated datasets do not expose (yet) these attributes, so add them.
+ if not hasattr(dataset, "transform"):
+ setattr(dataset, "transform", transform)
+ if not hasattr(dataset, "target_transform"):
+ setattr(dataset, "target_transform", target_transform)
+
+ return dataset
+
+
+def _make_sampler(
+ *,
+ dataset,
+ type: Optional[SamplerType] = None,
+ shuffle: bool = False,
+ seed: int = 0,
+ size: int = -1,
+ advance: int = 0,
+) -> Optional[Sampler]:
+ sample_count = len(dataset)
+
+ if type == SamplerType.INFINITE:
+ logger.info("sampler: infinite")
+ if size > 0:
+ raise ValueError("sampler size > 0 is invalid")
+ return InfiniteSampler(
+ sample_count=sample_count,
+ shuffle=shuffle,
+ seed=seed,
+ advance=advance,
+ )
+ elif type in (SamplerType.SHARDED_INFINITE, SamplerType.SHARDED_INFINITE_NEW):
+ logger.info("sampler: sharded infinite")
+ if size > 0:
+ raise ValueError("sampler size > 0 is invalid")
+ # TODO: Remove support for old shuffling
+ use_new_shuffle_tensor_slice = type == SamplerType.SHARDED_INFINITE_NEW
+ return ShardedInfiniteSampler(
+ sample_count=sample_count,
+ shuffle=shuffle,
+ seed=seed,
+ advance=advance,
+ use_new_shuffle_tensor_slice=use_new_shuffle_tensor_slice,
+ )
+ elif type == SamplerType.EPOCH:
+ logger.info("sampler: epoch")
+ if advance > 0:
+ raise NotImplementedError("sampler advance > 0 is not supported")
+ size = size if size > 0 else sample_count
+ logger.info(f"# of samples / epoch: {size:,d}")
+ return EpochSampler(
+ size=size,
+ sample_count=sample_count,
+ shuffle=shuffle,
+ seed=seed,
+ )
+ elif type == SamplerType.DISTRIBUTED:
+ logger.info("sampler: distributed")
+ if size > 0:
+ raise ValueError("sampler size > 0 is invalid")
+ if advance > 0:
+ raise ValueError("sampler advance > 0 is invalid")
+ return torch.utils.data.DistributedSampler(
+ dataset=dataset,
+ shuffle=shuffle,
+ seed=seed,
+ drop_last=False,
+ )
+
+ logger.info("sampler: none")
+ return None
+
+
+T = TypeVar("T")
+
+
+def make_data_loader(
+ *,
+ dataset,
+ batch_size: int,
+ num_workers: int,
+ shuffle: bool = True,
+ seed: int = 0,
+ sampler_type: Optional[SamplerType] = SamplerType.INFINITE,
+ sampler_size: int = -1,
+ sampler_advance: int = 0,
+ drop_last: bool = True,
+ persistent_workers: bool = False,
+ collate_fn: Optional[Callable[[List[T]], Any]] = None,
+):
+ """
+ Creates a data loader with the specified parameters.
+
+ Args:
+ dataset: A dataset (third party, LaViDa or WebDataset).
+ batch_size: The size of batches to generate.
+ num_workers: The number of workers to use.
+ shuffle: Whether to shuffle samples.
+ seed: The random seed to use.
+ sampler_type: Which sampler to use: EPOCH, INFINITE, SHARDED_INFINITE, SHARDED_INFINITE_NEW, DISTRIBUTED or None.
+ sampler_size: The number of images per epoch (when applicable) or -1 for the entire dataset.
+ sampler_advance: How many samples to skip (when applicable).
+ drop_last: Whether the last non-full batch of data should be dropped.
+ persistent_workers: maintain the workers Dataset instances alive after a dataset has been consumed once.
+ collate_fn: Function that performs batch collation
+ """
+
+ sampler = _make_sampler(
+ dataset=dataset,
+ type=sampler_type,
+ shuffle=shuffle,
+ seed=seed,
+ size=sampler_size,
+ advance=sampler_advance,
+ )
+
+ logger.info("using PyTorch data loader")
+ data_loader = torch.utils.data.DataLoader(
+ dataset,
+ sampler=sampler,
+ batch_size=batch_size,
+ num_workers=num_workers,
+ pin_memory=True,
+ drop_last=drop_last,
+ persistent_workers=persistent_workers,
+ collate_fn=collate_fn,
+ )
+
+ try:
+ logger.info(f"# of batches: {len(data_loader):,d}")
+ except TypeError: # data loader has no length
+ logger.info("infinite data loader")
+ return data_loader
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/data/masking.py b/torchhub/facebookresearch_dinov2_main/dinov2/data/masking.py
new file mode 100644
index 0000000000000000000000000000000000000000..dc3c72648c3e440dcdb284366b98d2df12ad1272
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/data/masking.py
@@ -0,0 +1,87 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import random
+import math
+import numpy as np
+
+
+class MaskingGenerator:
+ def __init__(
+ self,
+ input_size,
+ num_masking_patches=None,
+ min_num_patches=4,
+ max_num_patches=None,
+ min_aspect=0.3,
+ max_aspect=None,
+ ):
+ if not isinstance(input_size, tuple):
+ input_size = (input_size,) * 2
+ self.height, self.width = input_size
+
+ self.num_patches = self.height * self.width
+ self.num_masking_patches = num_masking_patches
+
+ self.min_num_patches = min_num_patches
+ self.max_num_patches = num_masking_patches if max_num_patches is None else max_num_patches
+
+ max_aspect = max_aspect or 1 / min_aspect
+ self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect))
+
+ def __repr__(self):
+ repr_str = "Generator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % (
+ self.height,
+ self.width,
+ self.min_num_patches,
+ self.max_num_patches,
+ self.num_masking_patches,
+ self.log_aspect_ratio[0],
+ self.log_aspect_ratio[1],
+ )
+ return repr_str
+
+ def get_shape(self):
+ return self.height, self.width
+
+ def _mask(self, mask, max_mask_patches):
+ delta = 0
+ for _ in range(10):
+ target_area = random.uniform(self.min_num_patches, max_mask_patches)
+ aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
+ h = int(round(math.sqrt(target_area * aspect_ratio)))
+ w = int(round(math.sqrt(target_area / aspect_ratio)))
+ if w < self.width and h < self.height:
+ top = random.randint(0, self.height - h)
+ left = random.randint(0, self.width - w)
+
+ num_masked = mask[top : top + h, left : left + w].sum()
+ # Overlap
+ if 0 < h * w - num_masked <= max_mask_patches:
+ for i in range(top, top + h):
+ for j in range(left, left + w):
+ if mask[i, j] == 0:
+ mask[i, j] = 1
+ delta += 1
+
+ if delta > 0:
+ break
+ return delta
+
+ def __call__(self, num_masking_patches=0):
+ mask = np.zeros(shape=self.get_shape(), dtype=bool)
+ mask_count = 0
+ while mask_count < num_masking_patches:
+ max_mask_patches = num_masking_patches - mask_count
+ max_mask_patches = min(max_mask_patches, self.max_num_patches)
+
+ delta = self._mask(mask, max_mask_patches)
+ if delta == 0:
+ break
+ else:
+ mask_count += delta
+
+ return mask
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/data/samplers.py b/torchhub/facebookresearch_dinov2_main/dinov2/data/samplers.py
new file mode 100644
index 0000000000000000000000000000000000000000..e356edf603a33ce2d18a388fd799694e22d1980f
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/data/samplers.py
@@ -0,0 +1,230 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import itertools
+from typing import Any, Optional
+import warnings
+
+import numpy as np
+import torch
+from torch.utils.data.sampler import Sampler
+
+import dinov2.distributed as distributed
+
+
+class EpochSampler(Sampler):
+ def __init__(
+ self,
+ *,
+ size: int,
+ sample_count: int,
+ shuffle: bool = False,
+ seed: int = 0,
+ start: Optional[int] = None,
+ step: Optional[int] = None,
+ ):
+ self._size = size
+ self._sample_count = sample_count
+ self._shuffle = shuffle
+ self._seed = seed
+ self._start = distributed.get_global_rank() if start is None else start
+ self._step = distributed.get_global_size() if step is None else step
+ self._epoch = 0
+
+ def __iter__(self):
+ count = (self._size + self._sample_count - 1) // self._sample_count
+ tiled_indices = np.tile(np.arange(self._sample_count), count)
+ if self._shuffle:
+ seed = self._seed * self._epoch if self._seed != 0 else self._epoch
+ rng = np.random.default_rng(seed)
+ iterable = rng.choice(tiled_indices, self._size, replace=False)
+ else:
+ iterable = tiled_indices[: self._size]
+
+ yield from itertools.islice(iterable, self._start, None, self._step)
+
+ def __len__(self):
+ return (self._size - self._start + self._step - 1) // self._step
+
+ def set_epoch(self, epoch):
+ self._epoch = epoch
+
+
+def _get_numpy_dtype(size: int) -> Any:
+ return np.int32 if size <= 2**31 else np.int64
+
+
+def _get_torch_dtype(size: int) -> Any:
+ return torch.int32 if size <= 2**31 else torch.int64
+
+
+def _generate_randperm_indices(*, size: int, generator: torch.Generator):
+ """Generate the indices of a random permutation."""
+ dtype = _get_torch_dtype(size)
+ # This is actually matching PyTorch's CPU implementation, see: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/TensorFactories.cpp#L900-L921
+ perm = torch.arange(size, dtype=dtype)
+ for i in range(size):
+ j = torch.randint(i, size, size=(1,), generator=generator).item()
+
+ # Always swap even if no-op
+ value = perm[j].item()
+ perm[j] = perm[i].item()
+ perm[i] = value
+ yield value
+
+
+class InfiniteSampler(Sampler):
+ def __init__(
+ self,
+ *,
+ sample_count: int,
+ shuffle: bool = False,
+ seed: int = 0,
+ start: Optional[int] = None,
+ step: Optional[int] = None,
+ advance: int = 0,
+ ):
+ self._sample_count = sample_count
+ self._seed = seed
+ self._shuffle = shuffle
+ self._start = distributed.get_global_rank() if start is None else start
+ self._step = distributed.get_global_size() if step is None else step
+ self._advance = advance
+
+ def __iter__(self):
+ if self._shuffle:
+ iterator = self._shuffled_iterator()
+ else:
+ iterator = self._iterator()
+
+ yield from itertools.islice(iterator, self._advance, None)
+
+ def _iterator(self):
+ assert not self._shuffle
+
+ while True:
+ iterable = range(self._sample_count)
+ yield from itertools.islice(iterable, self._start, None, self._step)
+
+ def _shuffled_iterator(self):
+ assert self._shuffle
+
+ # Instantiate a generator here (rather than in the ctor) to keep the class
+ # picklable (requirement of mp.spawn)
+ generator = torch.Generator().manual_seed(self._seed)
+
+ while True:
+ iterable = _generate_randperm_indices(size=self._sample_count, generator=generator)
+ yield from itertools.islice(iterable, self._start, None, self._step)
+
+
+# The following function is somewhat equivalent to _new_shuffle_tensor_slice below,
+# but avoids a full in-place random permutation generation.
+def _shuffle_tensor_slice(
+ *, tensor: torch.Tensor, start: int = 0, step: int = 1, generator: torch.Generator
+) -> np.ndarray:
+ stop = len(tensor)
+ count = stop // step
+ drop_count = stop - step * count
+ if drop_count:
+ warnings.warn(f"# of dropped samples: {drop_count}")
+
+ dtype = _get_numpy_dtype(stop)
+ result = np.empty(count, dtype=dtype)
+
+ for i in range(count):
+ j = torch.randint(0, i + 1, size=(1,), generator=generator).item() if i > 0 else 0
+
+ result[i] = result[j]
+ result[j] = tensor[start + i * step].item()
+
+ return result
+
+
+def _new_shuffle_tensor_slice(
+ *, tensor: torch.Tensor, start: int = 0, step: int = 1, generator: torch.Generator
+) -> np.ndarray:
+ stop = len(tensor)
+ count = stop // step
+ dtype = torch.int64 # Needed for using randperm result as indices
+ count = stop // step
+ drop_count = stop - step * count
+ if drop_count:
+ warnings.warn(f"# of dropped samples: {drop_count}")
+ indices = torch.randperm(count, dtype=dtype, generator=generator)
+ return tensor[start::step][indices].numpy()
+
+
+def _make_seed(seed: int, start: int, iter_count: int) -> int:
+ # NOTE: Tried a few variants (including iter_count << 32), this one worked best.
+ return seed + start + (iter_count << 24)
+
+
+class ShardedInfiniteSampler(Sampler):
+ def __init__(
+ self,
+ *,
+ sample_count: int,
+ shuffle: bool = False,
+ seed: int = 0,
+ start: Optional[int] = None,
+ step: Optional[int] = None,
+ advance: int = 0,
+ use_new_shuffle_tensor_slice: bool = False,
+ ):
+ self._sample_count = sample_count
+ self._seed = seed
+ self._shuffle = shuffle
+ self._start = distributed.get_global_rank() if start is None else start
+ self._step = distributed.get_global_size() if step is None else step
+ self._advance = advance
+ self._iter_count = 0
+ self._shuffle_tensor_slice_fn = (
+ _new_shuffle_tensor_slice if use_new_shuffle_tensor_slice else _shuffle_tensor_slice
+ )
+
+ def __iter__(self):
+ iter_count = self._advance // self._sample_count
+ if iter_count > 0:
+ self._advance -= iter_count * self._sample_count
+ self._iter_count += iter_count
+
+ if self._shuffle:
+ iterator = self._shuffled_iterator()
+ else:
+ iterator = self._iterator()
+
+ yield from itertools.islice(iterator, self._advance, None)
+
+ def _iterator(self):
+ assert not self._shuffle
+
+ while True:
+ iterable = range(self._sample_count)
+ yield from itertools.islice(iterable, self._start, None, self._step)
+
+ def _shuffled_iterator(self):
+ assert self._shuffle
+
+ # Instantiate a generator here (rather than in the ctor) to be keep the class
+ # picklable (requirement of mp.spawn)
+ generator = torch.Generator()
+
+ # Always shuffle everything first
+ generator.manual_seed(self._seed)
+ dtype = _get_torch_dtype(self._sample_count)
+ perm = torch.randperm(self._sample_count, dtype=dtype, generator=generator)
+
+ while True:
+ # Re-seed on each iteration to allow skipping whole permutations
+ seed = _make_seed(self._seed, self._start, self._iter_count)
+ generator.manual_seed(seed)
+
+ iterable = self._shuffle_tensor_slice_fn(
+ tensor=perm, start=self._start, step=self._step, generator=generator
+ )
+ yield from iterable
+ self._iter_count += 1
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/data/transforms.py b/torchhub/facebookresearch_dinov2_main/dinov2/data/transforms.py
new file mode 100644
index 0000000000000000000000000000000000000000..f1bc4cbd1a459a9f44314806cf9ccedea112ab14
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/data/transforms.py
@@ -0,0 +1,92 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Sequence
+
+import torch
+from torchvision import transforms
+
+
+class GaussianBlur(transforms.RandomApply):
+ """
+ Apply Gaussian Blur to the PIL image.
+ """
+
+ def __init__(self, *, p: float = 0.5, radius_min: float = 0.1, radius_max: float = 2.0):
+ # NOTE: torchvision is applying 1 - probability to return the original image
+ keep_p = 1 - p
+ transform = transforms.GaussianBlur(kernel_size=9, sigma=(radius_min, radius_max))
+ super().__init__(transforms=[transform], p=keep_p)
+
+
+class MaybeToTensor(transforms.ToTensor):
+ """
+ Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor, or keep as is if already a tensor.
+ """
+
+ def __call__(self, pic):
+ """
+ Args:
+ pic (PIL Image, numpy.ndarray or torch.tensor): Image to be converted to tensor.
+ Returns:
+ Tensor: Converted image.
+ """
+ if isinstance(pic, torch.Tensor):
+ return pic
+ return super().__call__(pic)
+
+
+# Use timm's names
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+
+def make_normalize_transform(
+ mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
+ std: Sequence[float] = IMAGENET_DEFAULT_STD,
+) -> transforms.Normalize:
+ return transforms.Normalize(mean=mean, std=std)
+
+
+# This roughly matches torchvision's preset for classification training:
+# https://github.com/pytorch/vision/blob/main/references/classification/presets.py#L6-L44
+def make_classification_train_transform(
+ *,
+ crop_size: int = 224,
+ interpolation=transforms.InterpolationMode.BICUBIC,
+ hflip_prob: float = 0.5,
+ mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
+ std: Sequence[float] = IMAGENET_DEFAULT_STD,
+):
+ transforms_list = [transforms.RandomResizedCrop(crop_size, interpolation=interpolation)]
+ if hflip_prob > 0.0:
+ transforms_list.append(transforms.RandomHorizontalFlip(hflip_prob))
+ transforms_list.extend(
+ [
+ MaybeToTensor(),
+ make_normalize_transform(mean=mean, std=std),
+ ]
+ )
+ return transforms.Compose(transforms_list)
+
+
+# This matches (roughly) torchvision's preset for classification evaluation:
+# https://github.com/pytorch/vision/blob/main/references/classification/presets.py#L47-L69
+def make_classification_eval_transform(
+ *,
+ resize_size: int = 256,
+ interpolation=transforms.InterpolationMode.BICUBIC,
+ crop_size: int = 224,
+ mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
+ std: Sequence[float] = IMAGENET_DEFAULT_STD,
+) -> transforms.Compose:
+ transforms_list = [
+ transforms.Resize(resize_size, interpolation=interpolation),
+ transforms.CenterCrop(crop_size),
+ MaybeToTensor(),
+ make_normalize_transform(mean=mean, std=std),
+ ]
+ return transforms.Compose(transforms_list)
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/distributed/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/distributed/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4ccd663f33d5a21ad1f9d25db7bd378ec52aeac2
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/distributed/__init__.py
@@ -0,0 +1,271 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import random
+import re
+import socket
+from typing import Dict, List
+
+import torch
+import torch.distributed as dist
+
+_LOCAL_RANK = -1
+_LOCAL_WORLD_SIZE = -1
+
+
+def is_enabled() -> bool:
+ """
+ Returns:
+ True if distributed training is enabled
+ """
+ return dist.is_available() and dist.is_initialized()
+
+
+def get_global_size() -> int:
+ """
+ Returns:
+ The number of processes in the process group
+ """
+ return dist.get_world_size() if is_enabled() else 1
+
+
+def get_global_rank() -> int:
+ """
+ Returns:
+ The rank of the current process within the global process group.
+ """
+ return dist.get_rank() if is_enabled() else 0
+
+
+def get_local_rank() -> int:
+ """
+ Returns:
+ The rank of the current process within the local (per-machine) process group.
+ """
+ if not is_enabled():
+ return 0
+ assert 0 <= _LOCAL_RANK < _LOCAL_WORLD_SIZE
+ return _LOCAL_RANK
+
+
+def get_local_size() -> int:
+ """
+ Returns:
+ The size of the per-machine process group,
+ i.e. the number of processes per machine.
+ """
+ if not is_enabled():
+ return 1
+ assert 0 <= _LOCAL_RANK < _LOCAL_WORLD_SIZE
+ return _LOCAL_WORLD_SIZE
+
+
+def is_main_process() -> bool:
+ """
+ Returns:
+ True if the current process is the main one.
+ """
+ return get_global_rank() == 0
+
+
+def _restrict_print_to_main_process() -> None:
+ """
+ This function disables printing when not in the main process
+ """
+ import builtins as __builtin__
+
+ builtin_print = __builtin__.print
+
+ def print(*args, **kwargs):
+ force = kwargs.pop("force", False)
+ if is_main_process() or force:
+ builtin_print(*args, **kwargs)
+
+ __builtin__.print = print
+
+
+def _get_master_port(seed: int = 0) -> int:
+ MIN_MASTER_PORT, MAX_MASTER_PORT = (20_000, 60_000)
+
+ master_port_str = os.environ.get("MASTER_PORT")
+ if master_port_str is None:
+ rng = random.Random(seed)
+ return rng.randint(MIN_MASTER_PORT, MAX_MASTER_PORT)
+
+ return int(master_port_str)
+
+
+def _get_available_port() -> int:
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+ # A "" host address means INADDR_ANY i.e. binding to all interfaces.
+ # Note this is not compatible with IPv6.
+ s.bind(("", 0))
+ port = s.getsockname()[1]
+ return port
+
+
+_TORCH_DISTRIBUTED_ENV_VARS = (
+ "MASTER_ADDR",
+ "MASTER_PORT",
+ "RANK",
+ "WORLD_SIZE",
+ "LOCAL_RANK",
+ "LOCAL_WORLD_SIZE",
+)
+
+
+def _collect_env_vars() -> Dict[str, str]:
+ return {env_var: os.environ[env_var] for env_var in _TORCH_DISTRIBUTED_ENV_VARS if env_var in os.environ}
+
+
+def _is_slurm_job_process() -> bool:
+ return "SLURM_JOB_ID" in os.environ
+
+
+def _parse_slurm_node_list(s: str) -> List[str]:
+ nodes = []
+ # Extract "hostname", "hostname[1-2,3,4-5]," substrings
+ p = re.compile(r"(([^\[]+)(?:\[([^\]]+)\])?),?")
+ for m in p.finditer(s):
+ prefix, suffixes = s[m.start(2) : m.end(2)], s[m.start(3) : m.end(3)]
+ for suffix in suffixes.split(","):
+ span = suffix.split("-")
+ if len(span) == 1:
+ nodes.append(prefix + suffix)
+ else:
+ width = len(span[0])
+ start, end = int(span[0]), int(span[1]) + 1
+ nodes.extend([prefix + f"{i:0{width}}" for i in range(start, end)])
+ return nodes
+
+
+def _check_env_variable(key: str, new_value: str):
+ # Only check for difference with preset environment variables
+ if key in os.environ and os.environ[key] != new_value:
+ raise RuntimeError(f"Cannot export environment variables as {key} is already set")
+
+
+class _TorchDistributedEnvironment:
+ def __init__(self):
+ self.master_addr = "127.0.0.1"
+ self.master_port = 0
+ self.rank = -1
+ self.world_size = -1
+ self.local_rank = -1
+ self.local_world_size = -1
+
+ if _is_slurm_job_process():
+ return self._set_from_slurm_env()
+
+ env_vars = _collect_env_vars()
+ if not env_vars:
+ # Environment is not set
+ pass
+ elif len(env_vars) == len(_TORCH_DISTRIBUTED_ENV_VARS):
+ # Environment is fully set
+ return self._set_from_preset_env()
+ else:
+ # Environment is partially set
+ collected_env_vars = ", ".join(env_vars.keys())
+ raise RuntimeError(f"Partially set environment: {collected_env_vars}")
+
+ if torch.cuda.device_count() > 0:
+ return self._set_from_local()
+
+ raise RuntimeError("Can't initialize PyTorch distributed environment")
+
+ # Slurm job created with sbatch, submitit, etc...
+ def _set_from_slurm_env(self):
+ # logger.info("Initialization from Slurm environment")
+ job_id = int(os.environ["SLURM_JOB_ID"])
+ node_count = int(os.environ["SLURM_JOB_NUM_NODES"])
+ nodes = _parse_slurm_node_list(os.environ["SLURM_JOB_NODELIST"])
+ assert len(nodes) == node_count
+
+ self.master_addr = nodes[0]
+ self.master_port = _get_master_port(seed=job_id)
+ self.rank = int(os.environ["SLURM_PROCID"])
+ self.world_size = int(os.environ["SLURM_NTASKS"])
+ assert self.rank < self.world_size
+ self.local_rank = int(os.environ["SLURM_LOCALID"])
+ self.local_world_size = self.world_size // node_count
+ assert self.local_rank < self.local_world_size
+
+ # Single node job with preset environment (i.e. torchrun)
+ def _set_from_preset_env(self):
+ # logger.info("Initialization from preset environment")
+ self.master_addr = os.environ["MASTER_ADDR"]
+ self.master_port = os.environ["MASTER_PORT"]
+ self.rank = int(os.environ["RANK"])
+ self.world_size = int(os.environ["WORLD_SIZE"])
+ assert self.rank < self.world_size
+ self.local_rank = int(os.environ["LOCAL_RANK"])
+ self.local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
+ assert self.local_rank < self.local_world_size
+
+ # Single node and GPU job (i.e. local script run)
+ def _set_from_local(self):
+ # logger.info("Initialization from local")
+ self.master_addr = "127.0.0.1"
+ self.master_port = _get_available_port()
+ self.rank = 0
+ self.world_size = 1
+ self.local_rank = 0
+ self.local_world_size = 1
+
+ def export(self, *, overwrite: bool) -> "_TorchDistributedEnvironment":
+ # See the "Environment variable initialization" section from
+ # https://pytorch.org/docs/stable/distributed.html for the complete list of
+ # environment variables required for the env:// initialization method.
+ env_vars = {
+ "MASTER_ADDR": self.master_addr,
+ "MASTER_PORT": str(self.master_port),
+ "RANK": str(self.rank),
+ "WORLD_SIZE": str(self.world_size),
+ "LOCAL_RANK": str(self.local_rank),
+ "LOCAL_WORLD_SIZE": str(self.local_world_size),
+ }
+ if not overwrite:
+ for k, v in env_vars.items():
+ _check_env_variable(k, v)
+
+ os.environ.update(env_vars)
+ return self
+
+
+def enable(*, set_cuda_current_device: bool = True, overwrite: bool = False, allow_nccl_timeout: bool = False):
+ """Enable distributed mode
+
+ Args:
+ set_cuda_current_device: If True, call torch.cuda.set_device() to set the
+ current PyTorch CUDA device to the one matching the local rank.
+ overwrite: If True, overwrites already set variables. Else fails.
+ """
+
+ global _LOCAL_RANK, _LOCAL_WORLD_SIZE
+ if _LOCAL_RANK >= 0 or _LOCAL_WORLD_SIZE >= 0:
+ raise RuntimeError("Distributed mode has already been enabled")
+ torch_env = _TorchDistributedEnvironment()
+ torch_env.export(overwrite=overwrite)
+
+ if set_cuda_current_device:
+ torch.cuda.set_device(torch_env.local_rank)
+
+ if allow_nccl_timeout:
+ # This allows to use torch distributed timeout in a NCCL backend
+ key, value = "NCCL_ASYNC_ERROR_HANDLING", "1"
+ if not overwrite:
+ _check_env_variable(key, value)
+ os.environ[key] = value
+
+ dist.init_process_group(backend="nccl")
+ dist.barrier()
+
+ # Finalize setup
+ _LOCAL_RANK = torch_env.local_rank
+ _LOCAL_WORLD_SIZE = torch_env.local_world_size
+ _restrict_print_to_main_process()
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/eval/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/eval/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0952fcc3f57e34b3747962e9ebd6fc57aeea63fa
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/eval/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/eval/knn.py b/torchhub/facebookresearch_dinov2_main/dinov2/eval/knn.py
new file mode 100644
index 0000000000000000000000000000000000000000..02ee261348e9871b10bfc40b7283b4f6205cba18
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/eval/knn.py
@@ -0,0 +1,405 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+from functools import partial
+import json
+import logging
+import os
+import sys
+from typing import List, Optional
+
+import torch
+from torch.nn.functional import one_hot, softmax
+
+import dinov2.distributed as distributed
+from dinov2.data import SamplerType, make_data_loader, make_dataset
+from dinov2.data.transforms import make_classification_eval_transform
+from dinov2.eval.metrics import AccuracyAveraging, build_topk_accuracy_metric
+from dinov2.eval.setup import get_args_parser as get_setup_args_parser
+from dinov2.eval.setup import setup_and_build_model
+from dinov2.eval.utils import ModelWithNormalize, evaluate, extract_features
+
+
+logger = logging.getLogger("dinov2")
+
+
+def get_args_parser(
+ description: Optional[str] = None,
+ parents: Optional[List[argparse.ArgumentParser]] = None,
+ add_help: bool = True,
+):
+ parents = parents or []
+ setup_args_parser = get_setup_args_parser(parents=parents, add_help=False)
+ parents = [setup_args_parser]
+ parser = argparse.ArgumentParser(
+ description=description,
+ parents=parents,
+ add_help=add_help,
+ )
+ parser.add_argument(
+ "--train-dataset",
+ dest="train_dataset_str",
+ type=str,
+ help="Training dataset",
+ )
+ parser.add_argument(
+ "--val-dataset",
+ dest="val_dataset_str",
+ type=str,
+ help="Validation dataset",
+ )
+ parser.add_argument(
+ "--nb_knn",
+ nargs="+",
+ type=int,
+ help="Number of NN to use. 20 is usually working the best.",
+ )
+ parser.add_argument(
+ "--temperature",
+ type=float,
+ help="Temperature used in the voting coefficient",
+ )
+ parser.add_argument(
+ "--gather-on-cpu",
+ action="store_true",
+ help="Whether to gather the train features on cpu, slower"
+ "but useful to avoid OOM for large datasets (e.g. ImageNet22k).",
+ )
+ parser.add_argument(
+ "--batch-size",
+ type=int,
+ help="Batch size.",
+ )
+ parser.add_argument(
+ "--n-per-class-list",
+ nargs="+",
+ type=int,
+ help="Number to take per class",
+ )
+ parser.add_argument(
+ "--n-tries",
+ type=int,
+ help="Number of tries",
+ )
+ parser.set_defaults(
+ train_dataset_str="ImageNet:split=TRAIN",
+ val_dataset_str="ImageNet:split=VAL",
+ nb_knn=[10, 20, 100, 200],
+ temperature=0.07,
+ batch_size=256,
+ n_per_class_list=[-1],
+ n_tries=1,
+ )
+ return parser
+
+
+class KnnModule(torch.nn.Module):
+ """
+ Gets knn of test features from all processes on a chunk of the train features
+
+ Each rank gets a chunk of the train features as well as a chunk of the test features.
+ In `compute_neighbors`, for each rank one after the other, its chunk of test features
+ is sent to all devices, partial knns are computed with each chunk of train features
+ then collated back on the original device.
+ """
+
+ def __init__(self, train_features, train_labels, nb_knn, T, device, num_classes=1000):
+ super().__init__()
+
+ self.global_rank = distributed.get_global_rank()
+ self.global_size = distributed.get_global_size()
+
+ self.device = device
+ self.train_features_rank_T = train_features.chunk(self.global_size)[self.global_rank].T.to(self.device)
+ self.candidates = train_labels.chunk(self.global_size)[self.global_rank].view(1, -1).to(self.device)
+
+ self.nb_knn = nb_knn
+ self.max_k = max(self.nb_knn)
+ self.T = T
+ self.num_classes = num_classes
+
+ def _get_knn_sims_and_labels(self, similarity, train_labels):
+ topk_sims, indices = similarity.topk(self.max_k, largest=True, sorted=True)
+ neighbors_labels = torch.gather(train_labels, 1, indices)
+ return topk_sims, neighbors_labels
+
+ def _similarity_for_rank(self, features_rank, source_rank):
+ # Send the features from `source_rank` to all ranks
+ broadcast_shape = torch.tensor(features_rank.shape).to(self.device)
+ torch.distributed.broadcast(broadcast_shape, source_rank)
+
+ broadcasted = features_rank
+ if self.global_rank != source_rank:
+ broadcasted = torch.zeros(*broadcast_shape, dtype=features_rank.dtype, device=self.device)
+ torch.distributed.broadcast(broadcasted, source_rank)
+
+ # Compute the neighbors for `source_rank` among `train_features_rank_T`
+ similarity_rank = torch.mm(broadcasted, self.train_features_rank_T)
+ candidate_labels = self.candidates.expand(len(similarity_rank), -1)
+ return self._get_knn_sims_and_labels(similarity_rank, candidate_labels)
+
+ def _gather_all_knn_for_rank(self, topk_sims, neighbors_labels, target_rank):
+ # Gather all neighbors for `target_rank`
+ topk_sims_rank = retrieved_rank = None
+ if self.global_rank == target_rank:
+ topk_sims_rank = [torch.zeros_like(topk_sims) for _ in range(self.global_size)]
+ retrieved_rank = [torch.zeros_like(neighbors_labels) for _ in range(self.global_size)]
+
+ torch.distributed.gather(topk_sims, topk_sims_rank, dst=target_rank)
+ torch.distributed.gather(neighbors_labels, retrieved_rank, dst=target_rank)
+
+ if self.global_rank == target_rank:
+ # Perform a second top-k on the k * global_size retrieved neighbors
+ topk_sims_rank = torch.cat(topk_sims_rank, dim=1)
+ retrieved_rank = torch.cat(retrieved_rank, dim=1)
+ results = self._get_knn_sims_and_labels(topk_sims_rank, retrieved_rank)
+ return results
+ return None
+
+ def compute_neighbors(self, features_rank):
+ for rank in range(self.global_size):
+ topk_sims, neighbors_labels = self._similarity_for_rank(features_rank, rank)
+ results = self._gather_all_knn_for_rank(topk_sims, neighbors_labels, rank)
+ if results is not None:
+ topk_sims_rank, neighbors_labels_rank = results
+ return topk_sims_rank, neighbors_labels_rank
+
+ def forward(self, features_rank):
+ """
+ Compute the results on all values of `self.nb_knn` neighbors from the full `self.max_k`
+ """
+ assert all(k <= self.max_k for k in self.nb_knn)
+
+ topk_sims, neighbors_labels = self.compute_neighbors(features_rank)
+ batch_size = neighbors_labels.shape[0]
+ topk_sims_transform = softmax(topk_sims / self.T, 1)
+ matmul = torch.mul(
+ one_hot(neighbors_labels, num_classes=self.num_classes),
+ topk_sims_transform.view(batch_size, -1, 1),
+ )
+ probas_for_k = {k: torch.sum(matmul[:, :k, :], 1) for k in self.nb_knn}
+ return probas_for_k
+
+
+class DictKeysModule(torch.nn.Module):
+ def __init__(self, keys):
+ super().__init__()
+ self.keys = keys
+
+ def forward(self, features_dict, targets):
+ for k in self.keys:
+ features_dict = features_dict[k]
+ return {"preds": features_dict, "target": targets}
+
+
+def create_module_dict(*, module, n_per_class_list, n_tries, nb_knn, train_features, train_labels):
+ modules = {}
+ mapping = create_class_indices_mapping(train_labels)
+ for npc in n_per_class_list:
+ if npc < 0: # Only one try needed when using the full data
+ full_module = module(
+ train_features=train_features,
+ train_labels=train_labels,
+ nb_knn=nb_knn,
+ )
+ modules["full"] = ModuleDictWithForward({"1": full_module})
+ continue
+ all_tries = {}
+ for t in range(n_tries):
+ final_indices = filter_train(mapping, npc, seed=t)
+ k_list = list(set(nb_knn + [npc]))
+ k_list = sorted([el for el in k_list if el <= npc])
+ all_tries[str(t)] = module(
+ train_features=train_features[final_indices],
+ train_labels=train_labels[final_indices],
+ nb_knn=k_list,
+ )
+ modules[f"{npc} per class"] = ModuleDictWithForward(all_tries)
+
+ return ModuleDictWithForward(modules)
+
+
+def filter_train(mapping, n_per_class, seed):
+ torch.manual_seed(seed)
+ final_indices = []
+ for k in mapping.keys():
+ index = torch.randperm(len(mapping[k]))[:n_per_class]
+ final_indices.append(mapping[k][index])
+ return torch.cat(final_indices).squeeze()
+
+
+def create_class_indices_mapping(labels):
+ unique_labels, inverse = torch.unique(labels, return_inverse=True)
+ mapping = {unique_labels[i]: (inverse == i).nonzero() for i in range(len(unique_labels))}
+ return mapping
+
+
+class ModuleDictWithForward(torch.nn.ModuleDict):
+ def forward(self, *args, **kwargs):
+ return {k: module(*args, **kwargs) for k, module in self._modules.items()}
+
+
+def eval_knn(
+ model,
+ train_dataset,
+ val_dataset,
+ accuracy_averaging,
+ nb_knn,
+ temperature,
+ batch_size,
+ num_workers,
+ gather_on_cpu,
+ n_per_class_list=[-1],
+ n_tries=1,
+):
+ model = ModelWithNormalize(model)
+
+ logger.info("Extracting features for train set...")
+ train_features, train_labels = extract_features(
+ model, train_dataset, batch_size, num_workers, gather_on_cpu=gather_on_cpu
+ )
+ logger.info(f"Train features created, shape {train_features.shape}.")
+
+ val_dataloader = make_data_loader(
+ dataset=val_dataset,
+ batch_size=batch_size,
+ num_workers=num_workers,
+ sampler_type=SamplerType.DISTRIBUTED,
+ drop_last=False,
+ shuffle=False,
+ persistent_workers=True,
+ )
+ num_classes = train_labels.max() + 1
+ metric_collection = build_topk_accuracy_metric(accuracy_averaging, num_classes=num_classes)
+
+ device = torch.cuda.current_device()
+ partial_module = partial(KnnModule, T=temperature, device=device, num_classes=num_classes)
+ knn_module_dict = create_module_dict(
+ module=partial_module,
+ n_per_class_list=n_per_class_list,
+ n_tries=n_tries,
+ nb_knn=nb_knn,
+ train_features=train_features,
+ train_labels=train_labels,
+ )
+ postprocessors, metrics = {}, {}
+ for n_per_class, knn_module in knn_module_dict.items():
+ for t, knn_try in knn_module.items():
+ postprocessors = {
+ **postprocessors,
+ **{(n_per_class, t, k): DictKeysModule([n_per_class, t, k]) for k in knn_try.nb_knn},
+ }
+ metrics = {**metrics, **{(n_per_class, t, k): metric_collection.clone() for k in knn_try.nb_knn}}
+ model_with_knn = torch.nn.Sequential(model, knn_module_dict)
+
+ # ============ evaluation ... ============
+ logger.info("Start the k-NN classification.")
+ _, results_dict = evaluate(model_with_knn, val_dataloader, postprocessors, metrics, device)
+
+ # Averaging the results over the n tries for each value of n_per_class
+ for n_per_class, knn_module in knn_module_dict.items():
+ first_try = list(knn_module.keys())[0]
+ k_list = knn_module[first_try].nb_knn
+ for k in k_list:
+ keys = results_dict[(n_per_class, first_try, k)].keys() # keys are e.g. `top-1` and `top-5`
+ results_dict[(n_per_class, k)] = {
+ key: torch.mean(torch.stack([results_dict[(n_per_class, t, k)][key] for t in knn_module.keys()]))
+ for key in keys
+ }
+ for t in knn_module.keys():
+ del results_dict[(n_per_class, t, k)]
+
+ return results_dict
+
+
+def eval_knn_with_model(
+ model,
+ output_dir,
+ train_dataset_str="ImageNet:split=TRAIN",
+ val_dataset_str="ImageNet:split=VAL",
+ nb_knn=(10, 20, 100, 200),
+ temperature=0.07,
+ autocast_dtype=torch.float,
+ accuracy_averaging=AccuracyAveraging.MEAN_ACCURACY,
+ transform=None,
+ gather_on_cpu=False,
+ batch_size=256,
+ num_workers=5,
+ n_per_class_list=[-1],
+ n_tries=1,
+):
+ transform = transform or make_classification_eval_transform()
+
+ train_dataset = make_dataset(
+ dataset_str=train_dataset_str,
+ transform=transform,
+ )
+ val_dataset = make_dataset(
+ dataset_str=val_dataset_str,
+ transform=transform,
+ )
+
+ with torch.cuda.amp.autocast(dtype=autocast_dtype):
+ results_dict_knn = eval_knn(
+ model=model,
+ train_dataset=train_dataset,
+ val_dataset=val_dataset,
+ accuracy_averaging=accuracy_averaging,
+ nb_knn=nb_knn,
+ temperature=temperature,
+ batch_size=batch_size,
+ num_workers=num_workers,
+ gather_on_cpu=gather_on_cpu,
+ n_per_class_list=n_per_class_list,
+ n_tries=n_tries,
+ )
+
+ results_dict = {}
+ if distributed.is_main_process():
+ for knn_ in results_dict_knn.keys():
+ top1 = results_dict_knn[knn_]["top-1"].item() * 100.0
+ top5 = results_dict_knn[knn_]["top-5"].item() * 100.0
+ results_dict[f"{knn_} Top 1"] = top1
+ results_dict[f"{knn_} Top 5"] = top5
+ logger.info(f"{knn_} classifier result: Top1: {top1:.2f} Top5: {top5:.2f}")
+
+ metrics_file_path = os.path.join(output_dir, "results_eval_knn.json")
+ with open(metrics_file_path, "a") as f:
+ for k, v in results_dict.items():
+ f.write(json.dumps({k: v}) + "\n")
+
+ if distributed.is_enabled():
+ torch.distributed.barrier()
+ return results_dict
+
+
+def main(args):
+ model, autocast_dtype = setup_and_build_model(args)
+ eval_knn_with_model(
+ model=model,
+ output_dir=args.output_dir,
+ train_dataset_str=args.train_dataset_str,
+ val_dataset_str=args.val_dataset_str,
+ nb_knn=args.nb_knn,
+ temperature=args.temperature,
+ autocast_dtype=autocast_dtype,
+ accuracy_averaging=AccuracyAveraging.MEAN_ACCURACY,
+ transform=None,
+ gather_on_cpu=args.gather_on_cpu,
+ batch_size=args.batch_size,
+ num_workers=5,
+ n_per_class_list=args.n_per_class_list,
+ n_tries=args.n_tries,
+ )
+ return 0
+
+
+if __name__ == "__main__":
+ description = "DINOv2 k-NN evaluation"
+ args_parser = get_args_parser(description=description)
+ args = args_parser.parse_args()
+ sys.exit(main(args))
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/eval/linear.py b/torchhub/facebookresearch_dinov2_main/dinov2/eval/linear.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d8202606999c0c01353904d8b02d2ff3509fef9
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/eval/linear.py
@@ -0,0 +1,626 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+from functools import partial
+import json
+import logging
+import os
+import sys
+from typing import List, Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+from fvcore.common.checkpoint import Checkpointer, PeriodicCheckpointer
+
+from dinov2.data import SamplerType, make_data_loader, make_dataset
+from dinov2.data.transforms import make_classification_eval_transform, make_classification_train_transform
+import dinov2.distributed as distributed
+from dinov2.eval.metrics import MetricType, build_metric
+from dinov2.eval.setup import get_args_parser as get_setup_args_parser
+from dinov2.eval.setup import setup_and_build_model
+from dinov2.eval.utils import ModelWithIntermediateLayers, evaluate
+from dinov2.logging import MetricLogger
+
+
+logger = logging.getLogger("dinov2")
+
+
+def get_args_parser(
+ description: Optional[str] = None,
+ parents: Optional[List[argparse.ArgumentParser]] = None,
+ add_help: bool = True,
+):
+ parents = parents or []
+ setup_args_parser = get_setup_args_parser(parents=parents, add_help=False)
+ parents = [setup_args_parser]
+ parser = argparse.ArgumentParser(
+ description=description,
+ parents=parents,
+ add_help=add_help,
+ )
+ parser.add_argument(
+ "--train-dataset",
+ dest="train_dataset_str",
+ type=str,
+ help="Training dataset",
+ )
+ parser.add_argument(
+ "--val-dataset",
+ dest="val_dataset_str",
+ type=str,
+ help="Validation dataset",
+ )
+ parser.add_argument(
+ "--test-datasets",
+ dest="test_dataset_strs",
+ type=str,
+ nargs="+",
+ help="Test datasets, none to reuse the validation dataset",
+ )
+ parser.add_argument(
+ "--epochs",
+ type=int,
+ help="Number of training epochs",
+ )
+ parser.add_argument(
+ "--batch-size",
+ type=int,
+ help="Batch Size (per GPU)",
+ )
+ parser.add_argument(
+ "--num-workers",
+ type=int,
+ help="Number de Workers",
+ )
+ parser.add_argument(
+ "--epoch-length",
+ type=int,
+ help="Length of an epoch in number of iterations",
+ )
+ parser.add_argument(
+ "--save-checkpoint-frequency",
+ type=int,
+ help="Number of epochs between two named checkpoint saves.",
+ )
+ parser.add_argument(
+ "--eval-period-iterations",
+ type=int,
+ help="Number of iterations between two evaluations.",
+ )
+ parser.add_argument(
+ "--learning-rates",
+ nargs="+",
+ type=float,
+ help="Learning rates to grid search.",
+ )
+ parser.add_argument(
+ "--no-resume",
+ action="store_true",
+ help="Whether to not resume from existing checkpoints",
+ )
+ parser.add_argument(
+ "--val-metric-type",
+ type=MetricType,
+ choices=list(MetricType),
+ help="Validation metric",
+ )
+ parser.add_argument(
+ "--test-metric-types",
+ type=MetricType,
+ choices=list(MetricType),
+ nargs="+",
+ help="Evaluation metric",
+ )
+ parser.add_argument(
+ "--classifier-fpath",
+ type=str,
+ help="Path to a file containing pretrained linear classifiers",
+ )
+ parser.add_argument(
+ "--val-class-mapping-fpath",
+ type=str,
+ help="Path to a file containing a mapping to adjust classifier outputs",
+ )
+ parser.add_argument(
+ "--test-class-mapping-fpaths",
+ nargs="+",
+ type=str,
+ help="Path to a file containing a mapping to adjust classifier outputs",
+ )
+ parser.set_defaults(
+ train_dataset_str="ImageNet:split=TRAIN",
+ val_dataset_str="ImageNet:split=VAL",
+ test_dataset_strs=None,
+ epochs=10,
+ batch_size=128,
+ num_workers=8,
+ epoch_length=1250,
+ save_checkpoint_frequency=20,
+ eval_period_iterations=1250,
+ learning_rates=[1e-5, 2e-5, 5e-5, 1e-4, 2e-4, 5e-4, 1e-3, 2e-3, 5e-3, 1e-2, 2e-2, 5e-2, 0.1],
+ val_metric_type=MetricType.MEAN_ACCURACY,
+ test_metric_types=None,
+ classifier_fpath=None,
+ val_class_mapping_fpath=None,
+ test_class_mapping_fpaths=[None],
+ )
+ return parser
+
+
+def has_ddp_wrapper(m: nn.Module) -> bool:
+ return isinstance(m, DistributedDataParallel)
+
+
+def remove_ddp_wrapper(m: nn.Module) -> nn.Module:
+ return m.module if has_ddp_wrapper(m) else m
+
+
+def _pad_and_collate(batch):
+ maxlen = max(len(targets) for image, targets in batch)
+ padded_batch = [
+ (image, np.pad(targets, (0, maxlen - len(targets)), constant_values=-1)) for image, targets in batch
+ ]
+ return torch.utils.data.default_collate(padded_batch)
+
+
+def create_linear_input(x_tokens_list, use_n_blocks, use_avgpool):
+ intermediate_output = x_tokens_list[-use_n_blocks:]
+ output = torch.cat([class_token for _, class_token in intermediate_output], dim=-1)
+ if use_avgpool:
+ output = torch.cat(
+ (
+ output,
+ torch.mean(intermediate_output[-1][0], dim=1), # patch tokens
+ ),
+ dim=-1,
+ )
+ output = output.reshape(output.shape[0], -1)
+ return output.float()
+
+
+class LinearClassifier(nn.Module):
+ """Linear layer to train on top of frozen features"""
+
+ def __init__(self, out_dim, use_n_blocks, use_avgpool, num_classes=1000):
+ super().__init__()
+ self.out_dim = out_dim
+ self.use_n_blocks = use_n_blocks
+ self.use_avgpool = use_avgpool
+ self.num_classes = num_classes
+ self.linear = nn.Linear(out_dim, num_classes)
+ self.linear.weight.data.normal_(mean=0.0, std=0.01)
+ self.linear.bias.data.zero_()
+
+ def forward(self, x_tokens_list):
+ output = create_linear_input(x_tokens_list, self.use_n_blocks, self.use_avgpool)
+ return self.linear(output)
+
+
+class AllClassifiers(nn.Module):
+ def __init__(self, classifiers_dict):
+ super().__init__()
+ self.classifiers_dict = nn.ModuleDict()
+ self.classifiers_dict.update(classifiers_dict)
+
+ def forward(self, inputs):
+ return {k: v.forward(inputs) for k, v in self.classifiers_dict.items()}
+
+ def __len__(self):
+ return len(self.classifiers_dict)
+
+
+class LinearPostprocessor(nn.Module):
+ def __init__(self, linear_classifier, class_mapping=None):
+ super().__init__()
+ self.linear_classifier = linear_classifier
+ self.register_buffer("class_mapping", None if class_mapping is None else torch.LongTensor(class_mapping))
+
+ def forward(self, samples, targets):
+ preds = self.linear_classifier(samples)
+ return {
+ "preds": preds[:, self.class_mapping] if self.class_mapping is not None else preds,
+ "target": targets,
+ }
+
+
+def scale_lr(learning_rates, batch_size):
+ return learning_rates * (batch_size * distributed.get_global_size()) / 256.0
+
+
+def setup_linear_classifiers(sample_output, n_last_blocks_list, learning_rates, batch_size, num_classes=1000):
+ linear_classifiers_dict = nn.ModuleDict()
+ optim_param_groups = []
+ for n in n_last_blocks_list:
+ for avgpool in [False, True]:
+ for _lr in learning_rates:
+ lr = scale_lr(_lr, batch_size)
+ out_dim = create_linear_input(sample_output, use_n_blocks=n, use_avgpool=avgpool).shape[1]
+ linear_classifier = LinearClassifier(
+ out_dim, use_n_blocks=n, use_avgpool=avgpool, num_classes=num_classes
+ )
+ linear_classifier = linear_classifier.cuda()
+ linear_classifiers_dict[
+ f"classifier_{n}_blocks_avgpool_{avgpool}_lr_{lr:.5f}".replace(".", "_")
+ ] = linear_classifier
+ optim_param_groups.append({"params": linear_classifier.parameters(), "lr": lr})
+
+ linear_classifiers = AllClassifiers(linear_classifiers_dict)
+ if distributed.is_enabled():
+ linear_classifiers = nn.parallel.DistributedDataParallel(linear_classifiers)
+
+ return linear_classifiers, optim_param_groups
+
+
+@torch.no_grad()
+def evaluate_linear_classifiers(
+ feature_model,
+ linear_classifiers,
+ data_loader,
+ metric_type,
+ metrics_file_path,
+ training_num_classes,
+ iteration,
+ prefixstring="",
+ class_mapping=None,
+ best_classifier_on_val=None,
+):
+ logger.info("running validation !")
+
+ num_classes = len(class_mapping) if class_mapping is not None else training_num_classes
+ metric = build_metric(metric_type, num_classes=num_classes)
+ postprocessors = {k: LinearPostprocessor(v, class_mapping) for k, v in linear_classifiers.classifiers_dict.items()}
+ metrics = {k: metric.clone() for k in linear_classifiers.classifiers_dict}
+
+ _, results_dict_temp = evaluate(
+ feature_model,
+ data_loader,
+ postprocessors,
+ metrics,
+ torch.cuda.current_device(),
+ )
+
+ logger.info("")
+ results_dict = {}
+ max_accuracy = 0
+ best_classifier = ""
+ for i, (classifier_string, metric) in enumerate(results_dict_temp.items()):
+ logger.info(f"{prefixstring} -- Classifier: {classifier_string} * {metric}")
+ if (
+ best_classifier_on_val is None and metric["top-1"].item() > max_accuracy
+ ) or classifier_string == best_classifier_on_val:
+ max_accuracy = metric["top-1"].item()
+ best_classifier = classifier_string
+
+ results_dict["best_classifier"] = {"name": best_classifier, "accuracy": max_accuracy}
+
+ logger.info(f"best classifier: {results_dict['best_classifier']}")
+
+ if distributed.is_main_process():
+ with open(metrics_file_path, "a") as f:
+ f.write(f"iter: {iteration}\n")
+ for k, v in results_dict.items():
+ f.write(json.dumps({k: v}) + "\n")
+ f.write("\n")
+
+ return results_dict
+
+
+def eval_linear(
+ *,
+ feature_model,
+ linear_classifiers,
+ train_data_loader,
+ val_data_loader,
+ metrics_file_path,
+ optimizer,
+ scheduler,
+ output_dir,
+ max_iter,
+ checkpoint_period, # In number of iter, creates a new file every period
+ running_checkpoint_period, # Period to update main checkpoint file
+ eval_period,
+ metric_type,
+ training_num_classes,
+ resume=True,
+ classifier_fpath=None,
+ val_class_mapping=None,
+):
+ checkpointer = Checkpointer(linear_classifiers, output_dir, optimizer=optimizer, scheduler=scheduler)
+ start_iter = checkpointer.resume_or_load(classifier_fpath or "", resume=resume).get("iteration", -1) + 1
+
+ periodic_checkpointer = PeriodicCheckpointer(checkpointer, checkpoint_period, max_iter=max_iter)
+ iteration = start_iter
+ logger.info("Starting training from iteration {}".format(start_iter))
+ metric_logger = MetricLogger(delimiter=" ")
+ header = "Training"
+
+ for data, labels in metric_logger.log_every(
+ train_data_loader,
+ 10,
+ header,
+ max_iter,
+ start_iter,
+ ):
+ data = data.cuda(non_blocking=True)
+ labels = labels.cuda(non_blocking=True)
+
+ features = feature_model(data)
+ outputs = linear_classifiers(features)
+
+ losses = {f"loss_{k}": nn.CrossEntropyLoss()(v, labels) for k, v in outputs.items()}
+ loss = sum(losses.values())
+
+ # compute the gradients
+ optimizer.zero_grad()
+ loss.backward()
+
+ # step
+ optimizer.step()
+ scheduler.step()
+
+ # log
+ if iteration % 10 == 0:
+ torch.cuda.synchronize()
+ metric_logger.update(loss=loss.item())
+ metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+ print("lr", optimizer.param_groups[0]["lr"])
+
+ if iteration - start_iter > 5:
+ if iteration % running_checkpoint_period == 0:
+ torch.cuda.synchronize()
+ if distributed.is_main_process():
+ logger.info("Checkpointing running_checkpoint")
+ periodic_checkpointer.save("running_checkpoint_linear_eval", iteration=iteration)
+ torch.cuda.synchronize()
+ periodic_checkpointer.step(iteration)
+
+ if eval_period > 0 and (iteration + 1) % eval_period == 0 and iteration != max_iter - 1:
+ _ = evaluate_linear_classifiers(
+ feature_model=feature_model,
+ linear_classifiers=remove_ddp_wrapper(linear_classifiers),
+ data_loader=val_data_loader,
+ metrics_file_path=metrics_file_path,
+ prefixstring=f"ITER: {iteration}",
+ metric_type=metric_type,
+ training_num_classes=training_num_classes,
+ iteration=iteration,
+ class_mapping=val_class_mapping,
+ )
+ torch.cuda.synchronize()
+
+ iteration = iteration + 1
+
+ val_results_dict = evaluate_linear_classifiers(
+ feature_model=feature_model,
+ linear_classifiers=remove_ddp_wrapper(linear_classifiers),
+ data_loader=val_data_loader,
+ metrics_file_path=metrics_file_path,
+ metric_type=metric_type,
+ training_num_classes=training_num_classes,
+ iteration=iteration,
+ class_mapping=val_class_mapping,
+ )
+ return val_results_dict, feature_model, linear_classifiers, iteration
+
+
+def make_eval_data_loader(test_dataset_str, batch_size, num_workers, metric_type):
+ test_dataset = make_dataset(
+ dataset_str=test_dataset_str,
+ transform=make_classification_eval_transform(),
+ )
+ test_data_loader = make_data_loader(
+ dataset=test_dataset,
+ batch_size=batch_size,
+ num_workers=num_workers,
+ sampler_type=SamplerType.DISTRIBUTED,
+ drop_last=False,
+ shuffle=False,
+ persistent_workers=False,
+ collate_fn=_pad_and_collate if metric_type == MetricType.IMAGENET_REAL_ACCURACY else None,
+ )
+ return test_data_loader
+
+
+def test_on_datasets(
+ feature_model,
+ linear_classifiers,
+ test_dataset_strs,
+ batch_size,
+ num_workers,
+ test_metric_types,
+ metrics_file_path,
+ training_num_classes,
+ iteration,
+ best_classifier_on_val,
+ prefixstring="",
+ test_class_mappings=[None],
+):
+ results_dict = {}
+ for test_dataset_str, class_mapping, metric_type in zip(test_dataset_strs, test_class_mappings, test_metric_types):
+ logger.info(f"Testing on {test_dataset_str}")
+ test_data_loader = make_eval_data_loader(test_dataset_str, batch_size, num_workers, metric_type)
+ dataset_results_dict = evaluate_linear_classifiers(
+ feature_model,
+ remove_ddp_wrapper(linear_classifiers),
+ test_data_loader,
+ metric_type,
+ metrics_file_path,
+ training_num_classes,
+ iteration,
+ prefixstring="",
+ class_mapping=class_mapping,
+ best_classifier_on_val=best_classifier_on_val,
+ )
+ results_dict[f"{test_dataset_str}_accuracy"] = 100.0 * dataset_results_dict["best_classifier"]["accuracy"]
+ return results_dict
+
+
+def run_eval_linear(
+ model,
+ output_dir,
+ train_dataset_str,
+ val_dataset_str,
+ batch_size,
+ epochs,
+ epoch_length,
+ num_workers,
+ save_checkpoint_frequency,
+ eval_period_iterations,
+ learning_rates,
+ autocast_dtype,
+ test_dataset_strs=None,
+ resume=True,
+ classifier_fpath=None,
+ val_class_mapping_fpath=None,
+ test_class_mapping_fpaths=[None],
+ val_metric_type=MetricType.MEAN_ACCURACY,
+ test_metric_types=None,
+):
+ seed = 0
+
+ if test_dataset_strs is None:
+ test_dataset_strs = [val_dataset_str]
+ if test_metric_types is None:
+ test_metric_types = [val_metric_type] * len(test_dataset_strs)
+ else:
+ assert len(test_metric_types) == len(test_dataset_strs)
+ assert len(test_dataset_strs) == len(test_class_mapping_fpaths)
+
+ train_transform = make_classification_train_transform()
+ train_dataset = make_dataset(
+ dataset_str=train_dataset_str,
+ transform=train_transform,
+ )
+ training_num_classes = len(torch.unique(torch.Tensor(train_dataset.get_targets().astype(int))))
+ sampler_type = SamplerType.SHARDED_INFINITE
+ # sampler_type = SamplerType.INFINITE
+
+ n_last_blocks_list = [1, 4]
+ n_last_blocks = max(n_last_blocks_list)
+ autocast_ctx = partial(torch.cuda.amp.autocast, enabled=True, dtype=autocast_dtype)
+ feature_model = ModelWithIntermediateLayers(model, n_last_blocks, autocast_ctx)
+ sample_output = feature_model(train_dataset[0][0].unsqueeze(0).cuda())
+
+ linear_classifiers, optim_param_groups = setup_linear_classifiers(
+ sample_output,
+ n_last_blocks_list,
+ learning_rates,
+ batch_size,
+ training_num_classes,
+ )
+
+ optimizer = torch.optim.SGD(optim_param_groups, momentum=0.9, weight_decay=0)
+ max_iter = epochs * epoch_length
+ scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, max_iter, eta_min=0)
+ checkpointer = Checkpointer(linear_classifiers, output_dir, optimizer=optimizer, scheduler=scheduler)
+ start_iter = checkpointer.resume_or_load(classifier_fpath or "", resume=resume).get("iteration", -1) + 1
+ train_data_loader = make_data_loader(
+ dataset=train_dataset,
+ batch_size=batch_size,
+ num_workers=num_workers,
+ shuffle=True,
+ seed=seed,
+ sampler_type=sampler_type,
+ sampler_advance=start_iter,
+ drop_last=True,
+ persistent_workers=True,
+ )
+ val_data_loader = make_eval_data_loader(val_dataset_str, batch_size, num_workers, val_metric_type)
+
+ checkpoint_period = save_checkpoint_frequency * epoch_length
+
+ if val_class_mapping_fpath is not None:
+ logger.info(f"Using class mapping from {val_class_mapping_fpath}")
+ val_class_mapping = np.load(val_class_mapping_fpath)
+ else:
+ val_class_mapping = None
+
+ test_class_mappings = []
+ for class_mapping_fpath in test_class_mapping_fpaths:
+ if class_mapping_fpath is not None and class_mapping_fpath != "None":
+ logger.info(f"Using class mapping from {class_mapping_fpath}")
+ class_mapping = np.load(class_mapping_fpath)
+ else:
+ class_mapping = None
+ test_class_mappings.append(class_mapping)
+
+ metrics_file_path = os.path.join(output_dir, "results_eval_linear.json")
+ val_results_dict, feature_model, linear_classifiers, iteration = eval_linear(
+ feature_model=feature_model,
+ linear_classifiers=linear_classifiers,
+ train_data_loader=train_data_loader,
+ val_data_loader=val_data_loader,
+ metrics_file_path=metrics_file_path,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ output_dir=output_dir,
+ max_iter=max_iter,
+ checkpoint_period=checkpoint_period,
+ running_checkpoint_period=epoch_length,
+ eval_period=eval_period_iterations,
+ metric_type=val_metric_type,
+ training_num_classes=training_num_classes,
+ resume=resume,
+ val_class_mapping=val_class_mapping,
+ classifier_fpath=classifier_fpath,
+ )
+ results_dict = {}
+ if len(test_dataset_strs) > 1 or test_dataset_strs[0] != val_dataset_str:
+ results_dict = test_on_datasets(
+ feature_model,
+ linear_classifiers,
+ test_dataset_strs,
+ batch_size,
+ 0, # num_workers,
+ test_metric_types,
+ metrics_file_path,
+ training_num_classes,
+ iteration,
+ val_results_dict["best_classifier"]["name"],
+ prefixstring="",
+ test_class_mappings=test_class_mappings,
+ )
+ results_dict["best_classifier"] = val_results_dict["best_classifier"]["name"]
+ results_dict[f"{val_dataset_str}_accuracy"] = 100.0 * val_results_dict["best_classifier"]["accuracy"]
+ logger.info("Test Results Dict " + str(results_dict))
+
+ return results_dict
+
+
+def main(args):
+ model, autocast_dtype = setup_and_build_model(args)
+ run_eval_linear(
+ model=model,
+ output_dir=args.output_dir,
+ train_dataset_str=args.train_dataset_str,
+ val_dataset_str=args.val_dataset_str,
+ test_dataset_strs=args.test_dataset_strs,
+ batch_size=args.batch_size,
+ epochs=args.epochs,
+ epoch_length=args.epoch_length,
+ num_workers=args.num_workers,
+ save_checkpoint_frequency=args.save_checkpoint_frequency,
+ eval_period_iterations=args.eval_period_iterations,
+ learning_rates=args.learning_rates,
+ autocast_dtype=autocast_dtype,
+ resume=not args.no_resume,
+ classifier_fpath=args.classifier_fpath,
+ val_metric_type=args.val_metric_type,
+ test_metric_types=args.test_metric_types,
+ val_class_mapping_fpath=args.val_class_mapping_fpath,
+ test_class_mapping_fpaths=args.test_class_mapping_fpaths,
+ )
+ return 0
+
+
+if __name__ == "__main__":
+ description = "DINOv2 linear evaluation"
+ args_parser = get_args_parser(description=description)
+ args = args_parser.parse_args()
+ sys.exit(main(args))
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/eval/log_regression.py b/torchhub/facebookresearch_dinov2_main/dinov2/eval/log_regression.py
new file mode 100644
index 0000000000000000000000000000000000000000..2e6ede2b616208cb49c7af67d58c8e6e4afb60e1
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/eval/log_regression.py
@@ -0,0 +1,445 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import gc
+import logging
+import sys
+import time
+from typing import List, Optional
+
+from cuml.linear_model import LogisticRegression
+import torch
+import torch.backends.cudnn as cudnn
+import torch.distributed
+from torch import nn
+from torch.utils.data import TensorDataset
+from torchmetrics import MetricTracker
+
+from dinov2.data import make_dataset
+from dinov2.data.transforms import make_classification_eval_transform
+from dinov2.distributed import get_global_rank, get_global_size
+from dinov2.eval.metrics import MetricType, build_metric
+from dinov2.eval.setup import get_args_parser as get_setup_args_parser
+from dinov2.eval.setup import setup_and_build_model
+from dinov2.eval.utils import evaluate, extract_features
+from dinov2.utils.dtype import as_torch_dtype
+
+
+logger = logging.getLogger("dinov2")
+
+DEFAULT_MAX_ITER = 1_000
+C_POWER_RANGE = torch.linspace(-6, 5, 45)
+_CPU_DEVICE = torch.device("cpu")
+
+
+def get_args_parser(
+ description: Optional[str] = None,
+ parents: Optional[List[argparse.ArgumentParser]] = None,
+ add_help: bool = True,
+):
+ parents = parents or []
+ setup_args_parser = get_setup_args_parser(parents=parents, add_help=False)
+ parents = [setup_args_parser]
+ parser = argparse.ArgumentParser(
+ description=description,
+ parents=parents,
+ add_help=add_help,
+ )
+ parser.add_argument(
+ "--train-dataset",
+ dest="train_dataset_str",
+ type=str,
+ help="Training dataset",
+ )
+ parser.add_argument(
+ "--val-dataset",
+ dest="val_dataset_str",
+ type=str,
+ help="Validation dataset",
+ )
+ parser.add_argument(
+ "--finetune-dataset-str",
+ dest="finetune_dataset_str",
+ type=str,
+ help="Fine-tuning dataset",
+ )
+ parser.add_argument(
+ "--finetune-on-val",
+ action="store_true",
+ help="If there is no finetune dataset, whether to choose the "
+ "hyperparameters on the val set instead of 10%% of the train dataset",
+ )
+ parser.add_argument(
+ "--metric-type",
+ type=MetricType,
+ choices=list(MetricType),
+ help="Metric type",
+ )
+ parser.add_argument(
+ "--train-features-device",
+ type=str,
+ help="Device to gather train features (cpu, cuda, cuda:0, etc.), default: %(default)s",
+ )
+ parser.add_argument(
+ "--train-dtype",
+ type=str,
+ help="Data type to convert the train features to (default: %(default)s)",
+ )
+ parser.add_argument(
+ "--max-train-iters",
+ type=int,
+ help="Maximum number of train iterations (default: %(default)s)",
+ )
+ parser.set_defaults(
+ train_dataset_str="ImageNet:split=TRAIN",
+ val_dataset_str="ImageNet:split=VAL",
+ finetune_dataset_str=None,
+ metric_type=MetricType.MEAN_ACCURACY,
+ train_features_device="cpu",
+ train_dtype="float64",
+ max_train_iters=DEFAULT_MAX_ITER,
+ finetune_on_val=False,
+ )
+ return parser
+
+
+class LogRegModule(nn.Module):
+ def __init__(
+ self,
+ C,
+ max_iter=DEFAULT_MAX_ITER,
+ dtype=torch.float64,
+ device=_CPU_DEVICE,
+ ):
+ super().__init__()
+ self.dtype = dtype
+ self.device = device
+ self.estimator = LogisticRegression(
+ penalty="l2",
+ C=C,
+ max_iter=max_iter,
+ output_type="numpy",
+ tol=1e-12,
+ linesearch_max_iter=50,
+ )
+
+ def forward(self, samples, targets):
+ samples_device = samples.device
+ samples = samples.to(dtype=self.dtype, device=self.device)
+ if self.device == _CPU_DEVICE:
+ samples = samples.numpy()
+ probas = self.estimator.predict_proba(samples)
+ return {"preds": torch.from_numpy(probas).to(samples_device), "target": targets}
+
+ def fit(self, train_features, train_labels):
+ train_features = train_features.to(dtype=self.dtype, device=self.device)
+ train_labels = train_labels.to(dtype=self.dtype, device=self.device)
+ if self.device == _CPU_DEVICE:
+ # both cuML and sklearn only work with numpy arrays on CPU
+ train_features = train_features.numpy()
+ train_labels = train_labels.numpy()
+ self.estimator.fit(train_features, train_labels)
+
+
+def evaluate_model(*, logreg_model, logreg_metric, test_data_loader, device):
+ postprocessors = {"metrics": logreg_model}
+ metrics = {"metrics": logreg_metric}
+ return evaluate(nn.Identity(), test_data_loader, postprocessors, metrics, device)
+
+
+def train_for_C(*, C, max_iter, train_features, train_labels, dtype=torch.float64, device=_CPU_DEVICE):
+ logreg_model = LogRegModule(C, max_iter=max_iter, dtype=dtype, device=device)
+ logreg_model.fit(train_features, train_labels)
+ return logreg_model
+
+
+def train_and_evaluate(
+ *,
+ C,
+ max_iter,
+ train_features,
+ train_labels,
+ logreg_metric,
+ test_data_loader,
+ train_dtype=torch.float64,
+ train_features_device,
+ eval_device,
+):
+ logreg_model = train_for_C(
+ C=C,
+ max_iter=max_iter,
+ train_features=train_features,
+ train_labels=train_labels,
+ dtype=train_dtype,
+ device=train_features_device,
+ )
+ return evaluate_model(
+ logreg_model=logreg_model,
+ logreg_metric=logreg_metric,
+ test_data_loader=test_data_loader,
+ device=eval_device,
+ )
+
+
+def sweep_C_values(
+ *,
+ train_features,
+ train_labels,
+ test_data_loader,
+ metric_type,
+ num_classes,
+ train_dtype=torch.float64,
+ train_features_device=_CPU_DEVICE,
+ max_train_iters=DEFAULT_MAX_ITER,
+):
+ if metric_type == MetricType.PER_CLASS_ACCURACY:
+ # If we want to output per-class accuracy, we select the hyperparameters with mean per class
+ metric_type = MetricType.MEAN_PER_CLASS_ACCURACY
+ logreg_metric = build_metric(metric_type, num_classes=num_classes)
+ metric_tracker = MetricTracker(logreg_metric, maximize=True)
+ ALL_C = 10**C_POWER_RANGE
+ logreg_models = {}
+
+ train_features = train_features.to(dtype=train_dtype, device=train_features_device)
+ train_labels = train_labels.to(device=train_features_device)
+
+ for i in range(get_global_rank(), len(ALL_C), get_global_size()):
+ C = ALL_C[i].item()
+ logger.info(
+ f"Training for C = {C:.5f}, dtype={train_dtype}, "
+ f"features: {train_features.shape}, {train_features.dtype}, "
+ f"labels: {train_labels.shape}, {train_labels.dtype}"
+ )
+ logreg_models[C] = train_for_C(
+ C=C,
+ max_iter=max_train_iters,
+ train_features=train_features,
+ train_labels=train_labels,
+ dtype=train_dtype,
+ device=train_features_device,
+ )
+
+ gather_list = [None for _ in range(get_global_size())]
+ torch.distributed.all_gather_object(gather_list, logreg_models)
+
+ logreg_models_gathered = {}
+ for logreg_dict in gather_list:
+ logreg_models_gathered.update(logreg_dict)
+
+ for i in range(len(ALL_C)):
+ metric_tracker.increment()
+ C = ALL_C[i].item()
+ evals = evaluate_model(
+ logreg_model=logreg_models_gathered[C],
+ logreg_metric=metric_tracker,
+ test_data_loader=test_data_loader,
+ device=torch.cuda.current_device(),
+ )
+ logger.info(f"Trained for C = {C:.5f}, accuracies = {evals}")
+
+ best_stats, which_epoch = metric_tracker.best_metric(return_step=True)
+ best_stats_100 = {k: 100.0 * v for k, v in best_stats.items()}
+ if which_epoch["top-1"] == i:
+ best_C = C
+ logger.info(f"Sweep best {best_stats_100}, best C = {best_C:.6f}")
+
+ return best_stats, best_C
+
+
+def eval_log_regression(
+ *,
+ model,
+ train_dataset,
+ val_dataset,
+ finetune_dataset,
+ metric_type,
+ batch_size,
+ num_workers,
+ finetune_on_val=False,
+ train_dtype=torch.float64,
+ train_features_device=_CPU_DEVICE,
+ max_train_iters=DEFAULT_MAX_ITER,
+):
+ """
+ Implements the "standard" process for log regression evaluation:
+ The value of C is chosen by training on train_dataset and evaluating on
+ finetune_dataset. Then, the final model is trained on a concatenation of
+ train_dataset and finetune_dataset, and is evaluated on val_dataset.
+ If there is no finetune_dataset, the value of C is the one that yields
+ the best results on a random 10% subset of the train dataset
+ """
+
+ start = time.time()
+
+ train_features, train_labels = extract_features(
+ model, train_dataset, batch_size, num_workers, gather_on_cpu=(train_features_device == _CPU_DEVICE)
+ )
+ val_features, val_labels = extract_features(
+ model, val_dataset, batch_size, num_workers, gather_on_cpu=(train_features_device == _CPU_DEVICE)
+ )
+ val_data_loader = torch.utils.data.DataLoader(
+ TensorDataset(val_features, val_labels),
+ batch_size=batch_size,
+ drop_last=False,
+ num_workers=0,
+ persistent_workers=False,
+ )
+
+ if finetune_dataset is None and finetune_on_val:
+ logger.info("Choosing hyperparameters on the val dataset")
+ finetune_features, finetune_labels = val_features, val_labels
+ elif finetune_dataset is None and not finetune_on_val:
+ logger.info("Choosing hyperparameters on 10% of the train dataset")
+ torch.manual_seed(0)
+ indices = torch.randperm(len(train_features), device=train_features.device)
+ finetune_index = indices[: len(train_features) // 10]
+ train_index = indices[len(train_features) // 10 :]
+ finetune_features, finetune_labels = train_features[finetune_index], train_labels[finetune_index]
+ train_features, train_labels = train_features[train_index], train_labels[train_index]
+ else:
+ logger.info("Choosing hyperparameters on the finetune dataset")
+ finetune_features, finetune_labels = extract_features(
+ model, finetune_dataset, batch_size, num_workers, gather_on_cpu=(train_features_device == _CPU_DEVICE)
+ )
+ # release the model - free GPU memory
+ del model
+ gc.collect()
+ torch.cuda.empty_cache()
+ finetune_data_loader = torch.utils.data.DataLoader(
+ TensorDataset(finetune_features, finetune_labels),
+ batch_size=batch_size,
+ drop_last=False,
+ )
+
+ if len(train_labels.shape) > 1:
+ num_classes = train_labels.shape[1]
+ else:
+ num_classes = train_labels.max() + 1
+
+ logger.info("Using cuML for logistic regression")
+
+ best_stats, best_C = sweep_C_values(
+ train_features=train_features,
+ train_labels=train_labels,
+ test_data_loader=finetune_data_loader,
+ metric_type=metric_type,
+ num_classes=num_classes,
+ train_dtype=train_dtype,
+ train_features_device=train_features_device,
+ max_train_iters=max_train_iters,
+ )
+
+ if not finetune_on_val:
+ logger.info("Best parameter found, concatenating features")
+ train_features = torch.cat((train_features, finetune_features))
+ train_labels = torch.cat((train_labels, finetune_labels))
+
+ logger.info("Training final model")
+ logreg_metric = build_metric(metric_type, num_classes=num_classes)
+ evals = train_and_evaluate(
+ C=best_C,
+ max_iter=max_train_iters,
+ train_features=train_features,
+ train_labels=train_labels,
+ logreg_metric=logreg_metric.clone(),
+ test_data_loader=val_data_loader,
+ eval_device=torch.cuda.current_device(),
+ train_dtype=train_dtype,
+ train_features_device=train_features_device,
+ )
+
+ best_stats = evals[1]["metrics"]
+
+ best_stats["best_C"] = best_C
+
+ logger.info(f"Log regression evaluation done in {int(time.time() - start)}s")
+ return best_stats
+
+
+def eval_log_regression_with_model(
+ model,
+ train_dataset_str="ImageNet:split=TRAIN",
+ val_dataset_str="ImageNet:split=VAL",
+ finetune_dataset_str=None,
+ autocast_dtype=torch.float,
+ finetune_on_val=False,
+ metric_type=MetricType.MEAN_ACCURACY,
+ train_dtype=torch.float64,
+ train_features_device=_CPU_DEVICE,
+ max_train_iters=DEFAULT_MAX_ITER,
+):
+ cudnn.benchmark = True
+
+ transform = make_classification_eval_transform(resize_size=224)
+ target_transform = None
+
+ train_dataset = make_dataset(dataset_str=train_dataset_str, transform=transform, target_transform=target_transform)
+ val_dataset = make_dataset(dataset_str=val_dataset_str, transform=transform, target_transform=target_transform)
+ if finetune_dataset_str is not None:
+ finetune_dataset = make_dataset(
+ dataset_str=finetune_dataset_str, transform=transform, target_transform=target_transform
+ )
+ else:
+ finetune_dataset = None
+
+ with torch.cuda.amp.autocast(dtype=autocast_dtype):
+ results_dict_logreg = eval_log_regression(
+ model=model,
+ train_dataset=train_dataset,
+ val_dataset=val_dataset,
+ finetune_dataset=finetune_dataset,
+ metric_type=metric_type,
+ batch_size=256,
+ num_workers=0, # 5,
+ finetune_on_val=finetune_on_val,
+ train_dtype=train_dtype,
+ train_features_device=train_features_device,
+ max_train_iters=max_train_iters,
+ )
+
+ results_dict = {
+ "top-1": results_dict_logreg["top-1"].cpu().numpy() * 100.0,
+ "top-5": results_dict_logreg.get("top-5", torch.tensor(0.0)).cpu().numpy() * 100.0,
+ "best_C": results_dict_logreg["best_C"],
+ }
+ logger.info(
+ "\n".join(
+ [
+ "Training of the supervised logistic regression on frozen features completed.\n"
+ "Top-1 test accuracy: {acc:.1f}".format(acc=results_dict["top-1"]),
+ "Top-5 test accuracy: {acc:.1f}".format(acc=results_dict["top-5"]),
+ "obtained for C = {c:.6f}".format(c=results_dict["best_C"]),
+ ]
+ )
+ )
+
+ torch.distributed.barrier()
+ return results_dict
+
+
+def main(args):
+ model, autocast_dtype = setup_and_build_model(args)
+ eval_log_regression_with_model(
+ model=model,
+ train_dataset_str=args.train_dataset_str,
+ val_dataset_str=args.val_dataset_str,
+ finetune_dataset_str=args.finetune_dataset_str,
+ autocast_dtype=autocast_dtype,
+ finetune_on_val=args.finetune_on_val,
+ metric_type=args.metric_type,
+ train_dtype=as_torch_dtype(args.train_dtype),
+ train_features_device=torch.device(args.train_features_device),
+ max_train_iters=args.max_train_iters,
+ )
+ return 0
+
+
+if __name__ == "__main__":
+ description = "DINOv2 logistic regression evaluation"
+ args_parser = get_args_parser(description=description)
+ args = args_parser.parse_args()
+ sys.exit(main(args))
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/eval/metrics.py b/torchhub/facebookresearch_dinov2_main/dinov2/eval/metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..80bf88da224e749dd6b3dd4b2bd90ec99eaec34e
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/eval/metrics.py
@@ -0,0 +1,114 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from enum import Enum
+import logging
+from typing import Any, Dict, Optional
+
+import torch
+from torch import Tensor
+from torchmetrics import Metric, MetricCollection
+from torchmetrics.classification import MulticlassAccuracy
+from torchmetrics.utilities.data import dim_zero_cat, select_topk
+
+
+logger = logging.getLogger("dinov2")
+
+
+class MetricType(Enum):
+ MEAN_ACCURACY = "mean_accuracy"
+ MEAN_PER_CLASS_ACCURACY = "mean_per_class_accuracy"
+ PER_CLASS_ACCURACY = "per_class_accuracy"
+ IMAGENET_REAL_ACCURACY = "imagenet_real_accuracy"
+
+ @property
+ def accuracy_averaging(self):
+ return getattr(AccuracyAveraging, self.name, None)
+
+ def __str__(self):
+ return self.value
+
+
+class AccuracyAveraging(Enum):
+ MEAN_ACCURACY = "micro"
+ MEAN_PER_CLASS_ACCURACY = "macro"
+ PER_CLASS_ACCURACY = "none"
+
+ def __str__(self):
+ return self.value
+
+
+def build_metric(metric_type: MetricType, *, num_classes: int, ks: Optional[tuple] = None):
+ if metric_type.accuracy_averaging is not None:
+ return build_topk_accuracy_metric(
+ average_type=metric_type.accuracy_averaging,
+ num_classes=num_classes,
+ ks=(1, 5) if ks is None else ks,
+ )
+ elif metric_type == MetricType.IMAGENET_REAL_ACCURACY:
+ return build_topk_imagenet_real_accuracy_metric(
+ num_classes=num_classes,
+ ks=(1, 5) if ks is None else ks,
+ )
+
+ raise ValueError(f"Unknown metric type {metric_type}")
+
+
+def build_topk_accuracy_metric(average_type: AccuracyAveraging, num_classes: int, ks: tuple = (1, 5)):
+ metrics: Dict[str, Metric] = {
+ f"top-{k}": MulticlassAccuracy(top_k=k, num_classes=int(num_classes), average=average_type.value) for k in ks
+ }
+ return MetricCollection(metrics)
+
+
+def build_topk_imagenet_real_accuracy_metric(num_classes: int, ks: tuple = (1, 5)):
+ metrics: Dict[str, Metric] = {f"top-{k}": ImageNetReaLAccuracy(top_k=k, num_classes=int(num_classes)) for k in ks}
+ return MetricCollection(metrics)
+
+
+class ImageNetReaLAccuracy(Metric):
+ is_differentiable: bool = False
+ higher_is_better: Optional[bool] = None
+ full_state_update: bool = False
+
+ def __init__(
+ self,
+ num_classes: int,
+ top_k: int = 1,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.num_classes = num_classes
+ self.top_k = top_k
+ self.add_state("tp", [], dist_reduce_fx="cat")
+
+ def update(self, preds: Tensor, target: Tensor) -> None: # type: ignore
+ # preds [B, D]
+ # target [B, A]
+ # preds_oh [B, D] with 0 and 1
+ # select top K highest probabilities, use one hot representation
+ preds_oh = select_topk(preds, self.top_k)
+ # target_oh [B, D + 1] with 0 and 1
+ target_oh = torch.zeros((preds_oh.shape[0], preds_oh.shape[1] + 1), device=target.device, dtype=torch.int32)
+ target = target.long()
+ # for undefined targets (-1) use a fake value `num_classes`
+ target[target == -1] = self.num_classes
+ # fill targets, use one hot representation
+ target_oh.scatter_(1, target, 1)
+ # target_oh [B, D] (remove the fake target at index `num_classes`)
+ target_oh = target_oh[:, :-1]
+ # tp [B] with 0 and 1
+ tp = (preds_oh * target_oh == 1).sum(dim=1)
+ # at least one match between prediction and target
+ tp.clip_(max=1)
+ # ignore instances where no targets are defined
+ mask = target_oh.sum(dim=1) > 0
+ tp = tp[mask]
+ self.tp.append(tp) # type: ignore
+
+ def compute(self) -> Tensor:
+ tp = dim_zero_cat(self.tp) # type: ignore
+ return tp.float().mean()
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/eval/setup.py b/torchhub/facebookresearch_dinov2_main/dinov2/eval/setup.py
new file mode 100644
index 0000000000000000000000000000000000000000..e7fadc2b63b994f569c8def82a43ed08ccd15b33
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/eval/setup.py
@@ -0,0 +1,76 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+from typing import Any, List, Optional, Tuple
+
+import torch
+import torch.backends.cudnn as cudnn
+
+from dinov2.models import build_model_from_cfg
+from dinov2.utils.config import setup
+import dinov2.utils.utils as dinov2_utils
+
+
+def get_args_parser(
+ description: Optional[str] = None,
+ parents: Optional[List[argparse.ArgumentParser]] = None,
+ add_help: bool = True,
+):
+ parser = argparse.ArgumentParser(
+ description=description,
+ parents=parents or [],
+ add_help=add_help,
+ )
+ parser.add_argument(
+ "--config-file",
+ type=str,
+ help="Model configuration file",
+ )
+ parser.add_argument(
+ "--pretrained-weights",
+ type=str,
+ help="Pretrained model weights",
+ )
+ parser.add_argument(
+ "--output-dir",
+ default="",
+ type=str,
+ help="Output directory to write results and logs",
+ )
+ parser.add_argument(
+ "--opts",
+ help="Extra configuration options",
+ default=[],
+ nargs="+",
+ )
+ return parser
+
+
+def get_autocast_dtype(config):
+ teacher_dtype_str = config.compute_precision.teacher.backbone.mixed_precision.param_dtype
+ if teacher_dtype_str == "fp16":
+ return torch.half
+ elif teacher_dtype_str == "bf16":
+ return torch.bfloat16
+ else:
+ return torch.float
+
+
+def build_model_for_eval(config, pretrained_weights):
+ model, _ = build_model_from_cfg(config, only_teacher=True)
+ dinov2_utils.load_pretrained_weights(model, pretrained_weights, "teacher")
+ model.eval()
+ model.cuda()
+ return model
+
+
+def setup_and_build_model(args) -> Tuple[Any, torch.dtype]:
+ cudnn.benchmark = True
+ config = setup(args)
+ model = build_model_for_eval(config, args.pretrained_weights)
+ autocast_dtype = get_autocast_dtype(config)
+ return model, autocast_dtype
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/eval/utils.py b/torchhub/facebookresearch_dinov2_main/dinov2/eval/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..b2f7e34f41ba6a0b911023e0c5375eef21f426fa
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/eval/utils.py
@@ -0,0 +1,147 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from typing import Dict, Optional
+
+import torch
+from torch import nn
+from torchmetrics import MetricCollection
+
+from dinov2.data import DatasetWithEnumeratedTargets, SamplerType, make_data_loader
+import dinov2.distributed as distributed
+from dinov2.logging import MetricLogger
+
+
+logger = logging.getLogger("dinov2")
+
+
+class ModelWithNormalize(torch.nn.Module):
+ def __init__(self, model):
+ super().__init__()
+ self.model = model
+
+ def forward(self, samples):
+ return nn.functional.normalize(self.model(samples), dim=1, p=2)
+
+
+class ModelWithIntermediateLayers(nn.Module):
+ def __init__(self, feature_model, n_last_blocks, autocast_ctx):
+ super().__init__()
+ self.feature_model = feature_model
+ self.feature_model.eval()
+ self.n_last_blocks = n_last_blocks
+ self.autocast_ctx = autocast_ctx
+
+ def forward(self, images):
+ with torch.inference_mode():
+ with self.autocast_ctx():
+ features = self.feature_model.get_intermediate_layers(
+ images, self.n_last_blocks, return_class_token=True
+ )
+ return features
+
+
+@torch.inference_mode()
+def evaluate(
+ model: nn.Module,
+ data_loader,
+ postprocessors: Dict[str, nn.Module],
+ metrics: Dict[str, MetricCollection],
+ device: torch.device,
+ criterion: Optional[nn.Module] = None,
+):
+ model.eval()
+ if criterion is not None:
+ criterion.eval()
+
+ for metric in metrics.values():
+ metric = metric.to(device)
+
+ metric_logger = MetricLogger(delimiter=" ")
+ header = "Test:"
+
+ for samples, targets, *_ in metric_logger.log_every(data_loader, 10, header):
+ outputs = model(samples.to(device))
+ targets = targets.to(device)
+
+ if criterion is not None:
+ loss = criterion(outputs, targets)
+ metric_logger.update(loss=loss.item())
+
+ for k, metric in metrics.items():
+ metric_inputs = postprocessors[k](outputs, targets)
+ metric.update(**metric_inputs)
+
+ metric_logger.synchronize_between_processes()
+ logger.info(f"Averaged stats: {metric_logger}")
+
+ stats = {k: metric.compute() for k, metric in metrics.items()}
+ metric_logger_stats = {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+ return metric_logger_stats, stats
+
+
+def all_gather_and_flatten(tensor_rank):
+ tensor_all_ranks = torch.empty(
+ distributed.get_global_size(),
+ *tensor_rank.shape,
+ dtype=tensor_rank.dtype,
+ device=tensor_rank.device,
+ )
+ tensor_list = list(tensor_all_ranks.unbind(0))
+ torch.distributed.all_gather(tensor_list, tensor_rank.contiguous())
+ return tensor_all_ranks.flatten(end_dim=1)
+
+
+def extract_features(model, dataset, batch_size, num_workers, gather_on_cpu=False):
+ dataset_with_enumerated_targets = DatasetWithEnumeratedTargets(dataset)
+ sample_count = len(dataset_with_enumerated_targets)
+ data_loader = make_data_loader(
+ dataset=dataset_with_enumerated_targets,
+ batch_size=batch_size,
+ num_workers=num_workers,
+ sampler_type=SamplerType.DISTRIBUTED,
+ drop_last=False,
+ shuffle=False,
+ )
+ return extract_features_with_dataloader(model, data_loader, sample_count, gather_on_cpu)
+
+
+@torch.inference_mode()
+def extract_features_with_dataloader(model, data_loader, sample_count, gather_on_cpu=False):
+ gather_device = torch.device("cpu") if gather_on_cpu else torch.device("cuda")
+ metric_logger = MetricLogger(delimiter=" ")
+ features, all_labels = None, None
+ for samples, (index, labels_rank) in metric_logger.log_every(data_loader, 10):
+ samples = samples.cuda(non_blocking=True)
+ labels_rank = labels_rank.cuda(non_blocking=True)
+ index = index.cuda(non_blocking=True)
+ features_rank = model(samples).float()
+
+ # init storage feature matrix
+ if features is None:
+ features = torch.zeros(sample_count, features_rank.shape[-1], device=gather_device)
+ labels_shape = list(labels_rank.shape)
+ labels_shape[0] = sample_count
+ all_labels = torch.full(labels_shape, fill_value=-1, device=gather_device)
+ logger.info(f"Storing features into tensor of shape {features.shape}")
+
+ # share indexes, features and labels between processes
+ index_all = all_gather_and_flatten(index).to(gather_device)
+ features_all_ranks = all_gather_and_flatten(features_rank).to(gather_device)
+ labels_all_ranks = all_gather_and_flatten(labels_rank).to(gather_device)
+
+ # update storage feature matrix
+ if len(index_all) > 0:
+ features.index_copy_(0, index_all, features_all_ranks)
+ all_labels.index_copy_(0, index_all, labels_all_ranks)
+
+ logger.info(f"Features shape: {tuple(features.shape)}")
+ logger.info(f"Labels shape: {tuple(all_labels.shape)}")
+
+ assert torch.all(all_labels > -1)
+
+ return features, all_labels
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/fsdp/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/fsdp/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..71d20397611619e6a02ea07f5305d650ffef2a51
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/fsdp/__init__.py
@@ -0,0 +1,158 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+from typing import Any
+
+import torch
+import dinov2.distributed as distributed
+from functools import partial
+from fvcore.common.checkpoint import Checkpointer
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import ShardingStrategy
+from torch.distributed.fsdp import MixedPrecision
+from torch.distributed.fsdp import StateDictType
+from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
+from torch.distributed.fsdp.wrap import ModuleWrapPolicy
+from torch.distributed.fsdp._runtime_utils import _reshard
+
+
+def get_fsdp_wrapper(model_cfg, modules_to_wrap=set()):
+ sharding_strategy_dict = {
+ "NO_SHARD": ShardingStrategy.NO_SHARD,
+ "SHARD_GRAD_OP": ShardingStrategy.SHARD_GRAD_OP,
+ "FULL_SHARD": ShardingStrategy.FULL_SHARD,
+ }
+
+ dtype_dict = {
+ "fp32": torch.float32,
+ "fp16": torch.float16,
+ "bf16": torch.bfloat16,
+ }
+
+ mixed_precision_config = MixedPrecision(
+ param_dtype=dtype_dict[model_cfg.mixed_precision.param_dtype],
+ reduce_dtype=dtype_dict[model_cfg.mixed_precision.reduce_dtype],
+ buffer_dtype=dtype_dict[model_cfg.mixed_precision.buffer_dtype],
+ )
+
+ sharding_strategy_config = sharding_strategy_dict[model_cfg.sharding_strategy]
+
+ local_rank = distributed.get_local_rank()
+
+ fsdp_wrapper = partial(
+ FSDP,
+ sharding_strategy=sharding_strategy_config,
+ mixed_precision=mixed_precision_config,
+ device_id=local_rank,
+ sync_module_states=True,
+ use_orig_params=True,
+ auto_wrap_policy=ModuleWrapPolicy(modules_to_wrap),
+ )
+ return fsdp_wrapper
+
+
+def is_fsdp(x):
+ return isinstance(x, FSDP)
+
+
+def is_sharded_fsdp(x):
+ return is_fsdp(x) and x.sharding_strategy is not ShardingStrategy.NO_SHARD
+
+
+def free_if_fsdp(x):
+ if is_sharded_fsdp(x):
+ handles = x._handles
+ true_list = [True for h in handles]
+ _reshard(x, handles, true_list)
+
+
+def get_fsdp_modules(x):
+ return FSDP.fsdp_modules(x)
+
+
+def reshard_fsdp_model(x):
+ for m in get_fsdp_modules(x):
+ free_if_fsdp(m)
+
+
+def rankstr():
+ return f"rank_{distributed.get_global_rank()}"
+
+
+class FSDPCheckpointer(Checkpointer):
+ def save(self, name: str, **kwargs: Any) -> None:
+ """
+ Dump model and checkpointables to a file.
+
+ Args:
+ name (str): name of the file.
+ kwargs (dict): extra arbitrary data to save.
+ """
+ if not self.save_dir or not self.save_to_disk:
+ return
+
+ data = {}
+ with FSDP.state_dict_type(self.model, StateDictType.LOCAL_STATE_DICT):
+ data["model"] = self.model.state_dict()
+
+ # data["model"] = self.model.state_dict()
+ for key, obj in self.checkpointables.items():
+ data[key] = obj.state_dict()
+ data.update(kwargs)
+
+ basename = f"{name}.{rankstr()}.pth"
+ save_file = os.path.join(self.save_dir, basename)
+ assert os.path.basename(save_file) == basename, basename
+ self.logger.info("Saving checkpoint to {}".format(save_file))
+ with self.path_manager.open(save_file, "wb") as f:
+ torch.save(data, f)
+ self.tag_last_checkpoint(basename)
+
+ def load(self, *args, **kwargs):
+ with FSDP.state_dict_type(self.model, StateDictType.LOCAL_STATE_DICT):
+ return super().load(*args, **kwargs)
+
+ def has_checkpoint(self) -> bool:
+ """
+ Returns:
+ bool: whether a checkpoint exists in the target directory.
+ """
+ save_file = os.path.join(self.save_dir, f"last_checkpoint.{rankstr()}")
+ return self.path_manager.exists(save_file)
+
+ def get_checkpoint_file(self) -> str:
+ """
+ Returns:
+ str: The latest checkpoint file in target directory.
+ """
+ save_file = os.path.join(self.save_dir, f"last_checkpoint.{rankstr()}")
+ try:
+ with self.path_manager.open(save_file, "r") as f:
+ last_saved = f.read().strip()
+ except IOError:
+ # if file doesn't exist, maybe because it has just been
+ # deleted by a separate process
+ return ""
+ # pyre-fixme[6]: For 2nd param expected `Union[PathLike[str], str]` but got
+ # `Union[bytes, str]`.
+ return os.path.join(self.save_dir, last_saved)
+
+ def tag_last_checkpoint(self, last_filename_basename: str) -> None:
+ """
+ Tag the last checkpoint.
+
+ Args:
+ last_filename_basename (str): the basename of the last filename.
+ """
+ if distributed.is_enabled():
+ torch.distributed.barrier()
+ save_file = os.path.join(self.save_dir, f"last_checkpoint.{rankstr()}")
+ with self.path_manager.open(save_file, "w") as f:
+ f.write(last_filename_basename) # pyre-ignore
+
+
+ShardedGradScaler = ShardedGradScaler
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/layers/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/layers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..31f196aacac5be8a7c537a3dfa8f97084671b466
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/layers/__init__.py
@@ -0,0 +1,12 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .dino_head import DINOHead
+from .mlp import Mlp
+from .patch_embed import PatchEmbed
+from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
+from .block import NestedTensorBlock
+from .attention import MemEffAttention
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/layers/attention.py b/torchhub/facebookresearch_dinov2_main/dinov2/layers/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f9b0c94b40967dfdff4f261c127cbd21328c905
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/layers/attention.py
@@ -0,0 +1,81 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+import logging
+
+from torch import Tensor
+from torch import nn
+
+
+logger = logging.getLogger("dinov2")
+
+
+try:
+ from xformers.ops import memory_efficient_attention, unbind, fmha
+
+ XFORMERS_AVAILABLE = True
+except ImportError:
+ logger.warning("xFormers not available")
+ XFORMERS_AVAILABLE = False
+
+
+class Attention(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int = 8,
+ qkv_bias: bool = False,
+ proj_bias: bool = True,
+ attn_drop: float = 0.0,
+ proj_drop: float = 0.0,
+ ) -> None:
+ super().__init__()
+ self.num_heads = num_heads
+ head_dim = dim // num_heads
+ self.scale = head_dim**-0.5
+
+ self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+ self.attn_drop = nn.Dropout(attn_drop)
+ self.proj = nn.Linear(dim, dim, bias=proj_bias)
+ self.proj_drop = nn.Dropout(proj_drop)
+
+ def forward(self, x: Tensor) -> Tensor:
+ B, N, C = x.shape
+ qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+ q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+ attn = q @ k.transpose(-2, -1)
+
+ attn = attn.softmax(dim=-1)
+ attn = self.attn_drop(attn)
+
+ x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ return x
+
+
+class MemEffAttention(Attention):
+ def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+ if not XFORMERS_AVAILABLE:
+ assert attn_bias is None, "xFormers is required for nested tensors usage"
+ return super().forward(x)
+
+ B, N, C = x.shape
+ qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+ q, k, v = unbind(qkv, 2)
+
+ x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+ x = x.reshape([B, N, C])
+
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ return x
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/layers/block.py b/torchhub/facebookresearch_dinov2_main/dinov2/layers/block.py
new file mode 100644
index 0000000000000000000000000000000000000000..25488f57cc0ad3c692f86b62555f6668e2a66db1
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/layers/block.py
@@ -0,0 +1,252 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+import logging
+from typing import Callable, List, Any, Tuple, Dict
+
+import torch
+from torch import nn, Tensor
+
+from .attention import Attention, MemEffAttention
+from .drop_path import DropPath
+from .layer_scale import LayerScale
+from .mlp import Mlp
+
+
+logger = logging.getLogger("dinov2")
+
+
+try:
+ from xformers.ops import fmha
+ from xformers.ops import scaled_index_add, index_select_cat
+
+ XFORMERS_AVAILABLE = True
+except ImportError:
+ logger.warning("xFormers not available")
+ XFORMERS_AVAILABLE = False
+
+
+class Block(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ qkv_bias: bool = False,
+ proj_bias: bool = True,
+ ffn_bias: bool = True,
+ drop: float = 0.0,
+ attn_drop: float = 0.0,
+ init_values=None,
+ drop_path: float = 0.0,
+ act_layer: Callable[..., nn.Module] = nn.GELU,
+ norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+ attn_class: Callable[..., nn.Module] = Attention,
+ ffn_layer: Callable[..., nn.Module] = Mlp,
+ ) -> None:
+ super().__init__()
+ # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
+ self.norm1 = norm_layer(dim)
+ self.attn = attn_class(
+ dim,
+ num_heads=num_heads,
+ qkv_bias=qkv_bias,
+ proj_bias=proj_bias,
+ attn_drop=attn_drop,
+ proj_drop=drop,
+ )
+ self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+ self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+ self.norm2 = norm_layer(dim)
+ mlp_hidden_dim = int(dim * mlp_ratio)
+ self.mlp = ffn_layer(
+ in_features=dim,
+ hidden_features=mlp_hidden_dim,
+ act_layer=act_layer,
+ drop=drop,
+ bias=ffn_bias,
+ )
+ self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+ self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+ self.sample_drop_ratio = drop_path
+
+ def forward(self, x: Tensor) -> Tensor:
+ def attn_residual_func(x: Tensor) -> Tensor:
+ return self.ls1(self.attn(self.norm1(x)))
+
+ def ffn_residual_func(x: Tensor) -> Tensor:
+ return self.ls2(self.mlp(self.norm2(x)))
+
+ if self.training and self.sample_drop_ratio > 0.1:
+ # the overhead is compensated only for a drop path rate larger than 0.1
+ x = drop_add_residual_stochastic_depth(
+ x,
+ residual_func=attn_residual_func,
+ sample_drop_ratio=self.sample_drop_ratio,
+ )
+ x = drop_add_residual_stochastic_depth(
+ x,
+ residual_func=ffn_residual_func,
+ sample_drop_ratio=self.sample_drop_ratio,
+ )
+ elif self.training and self.sample_drop_ratio > 0.0:
+ x = x + self.drop_path1(attn_residual_func(x))
+ x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2
+ else:
+ x = x + attn_residual_func(x)
+ x = x + ffn_residual_func(x)
+ return x
+
+
+def drop_add_residual_stochastic_depth(
+ x: Tensor,
+ residual_func: Callable[[Tensor], Tensor],
+ sample_drop_ratio: float = 0.0,
+) -> Tensor:
+ # 1) extract subset using permutation
+ b, n, d = x.shape
+ sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+ brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+ x_subset = x[brange]
+
+ # 2) apply residual_func to get residual
+ residual = residual_func(x_subset)
+
+ x_flat = x.flatten(1)
+ residual = residual.flatten(1)
+
+ residual_scale_factor = b / sample_subset_size
+
+ # 3) add the residual
+ x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+ return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+ b, n, d = x.shape
+ sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+ brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+ residual_scale_factor = b / sample_subset_size
+ return brange, residual_scale_factor
+
+
+def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
+ if scaling_vector is None:
+ x_flat = x.flatten(1)
+ residual = residual.flatten(1)
+ x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+ else:
+ x_plus_residual = scaled_index_add(
+ x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
+ )
+ return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+ """
+ this will perform the index select, cat the tensors, and provide the attn_bias from cache
+ """
+ batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
+ all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+ if all_shapes not in attn_bias_cache.keys():
+ seqlens = []
+ for b, x in zip(batch_sizes, x_list):
+ for _ in range(b):
+ seqlens.append(x.shape[1])
+ attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+ attn_bias._batch_sizes = batch_sizes
+ attn_bias_cache[all_shapes] = attn_bias
+
+ if branges is not None:
+ cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
+ else:
+ tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+ cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+ return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+ x_list: List[Tensor],
+ residual_func: Callable[[Tensor, Any], Tensor],
+ sample_drop_ratio: float = 0.0,
+ scaling_vector=None,
+) -> Tensor:
+ # 1) generate random set of indices for dropping samples in the batch
+ branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
+ branges = [s[0] for s in branges_scales]
+ residual_scale_factors = [s[1] for s in branges_scales]
+
+ # 2) get attention bias and index+concat the tensors
+ attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+ # 3) apply residual_func to get residual, and split the result
+ residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore
+
+ outputs = []
+ for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
+ outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
+ return outputs
+
+
+class NestedTensorBlock(Block):
+ def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
+ """
+ x_list contains a list of tensors to nest together and run
+ """
+ assert isinstance(self.attn, MemEffAttention)
+
+ if self.training and self.sample_drop_ratio > 0.0:
+
+ def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+ return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+ def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+ return self.mlp(self.norm2(x))
+
+ x_list = drop_add_residual_stochastic_depth_list(
+ x_list,
+ residual_func=attn_residual_func,
+ sample_drop_ratio=self.sample_drop_ratio,
+ scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
+ )
+ x_list = drop_add_residual_stochastic_depth_list(
+ x_list,
+ residual_func=ffn_residual_func,
+ sample_drop_ratio=self.sample_drop_ratio,
+ scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
+ )
+ return x_list
+ else:
+
+ def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+ return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+ def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+ return self.ls2(self.mlp(self.norm2(x)))
+
+ attn_bias, x = get_attn_bias_and_cat(x_list)
+ x = x + attn_residual_func(x, attn_bias=attn_bias)
+ x = x + ffn_residual_func(x)
+ return attn_bias.split(x)
+
+ def forward(self, x_or_x_list):
+ if isinstance(x_or_x_list, Tensor):
+ return super().forward(x_or_x_list)
+ elif isinstance(x_or_x_list, list):
+ assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage"
+ return self.forward_nested(x_or_x_list)
+ else:
+ raise AssertionError
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/layers/dino_head.py b/torchhub/facebookresearch_dinov2_main/dinov2/layers/dino_head.py
new file mode 100644
index 0000000000000000000000000000000000000000..7212db92a4fd8d4c7230e284e551a0234e9d8623
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/layers/dino_head.py
@@ -0,0 +1,59 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+from torch.nn.init import trunc_normal_
+from torch.nn.utils import weight_norm
+
+
+class DINOHead(nn.Module):
+ def __init__(
+ self,
+ in_dim,
+ out_dim,
+ use_bn=False,
+ nlayers=3,
+ hidden_dim=2048,
+ bottleneck_dim=256,
+ mlp_bias=True,
+ ):
+ super().__init__()
+ nlayers = max(nlayers, 1)
+ self.mlp = _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=hidden_dim, use_bn=use_bn, bias=mlp_bias)
+ self.apply(self._init_weights)
+ self.last_layer = weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
+ self.last_layer.weight_g.data.fill_(1)
+
+ def _init_weights(self, m):
+ if isinstance(m, nn.Linear):
+ trunc_normal_(m.weight, std=0.02)
+ if isinstance(m, nn.Linear) and m.bias is not None:
+ nn.init.constant_(m.bias, 0)
+
+ def forward(self, x):
+ x = self.mlp(x)
+ eps = 1e-6 if x.dtype == torch.float16 else 1e-12
+ x = nn.functional.normalize(x, dim=-1, p=2, eps=eps)
+ x = self.last_layer(x)
+ return x
+
+
+def _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=None, use_bn=False, bias=True):
+ if nlayers == 1:
+ return nn.Linear(in_dim, bottleneck_dim, bias=bias)
+ else:
+ layers = [nn.Linear(in_dim, hidden_dim, bias=bias)]
+ if use_bn:
+ layers.append(nn.BatchNorm1d(hidden_dim))
+ layers.append(nn.GELU())
+ for _ in range(nlayers - 2):
+ layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias))
+ if use_bn:
+ layers.append(nn.BatchNorm1d(hidden_dim))
+ layers.append(nn.GELU())
+ layers.append(nn.Linear(hidden_dim, bottleneck_dim, bias=bias))
+ return nn.Sequential(*layers)
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/layers/drop_path.py b/torchhub/facebookresearch_dinov2_main/dinov2/layers/drop_path.py
new file mode 100644
index 0000000000000000000000000000000000000000..af05625984dd14682cc96a63bf0c97bab1f123b1
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/layers/drop_path.py
@@ -0,0 +1,35 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
+
+
+from torch import nn
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+ if drop_prob == 0.0 or not training:
+ return x
+ keep_prob = 1 - drop_prob
+ shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
+ random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+ if keep_prob > 0.0:
+ random_tensor.div_(keep_prob)
+ output = x * random_tensor
+ return output
+
+
+class DropPath(nn.Module):
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+ def __init__(self, drop_prob=None):
+ super(DropPath, self).__init__()
+ self.drop_prob = drop_prob
+
+ def forward(self, x):
+ return drop_path(x, self.drop_prob, self.training)
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/layers/layer_scale.py b/torchhub/facebookresearch_dinov2_main/dinov2/layers/layer_scale.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca5daa52bd81d3581adeb2198ea5b7dba2a3aea1
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/layers/layer_scale.py
@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
+
+from typing import Union
+
+import torch
+from torch import Tensor
+from torch import nn
+
+
+class LayerScale(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ init_values: Union[float, Tensor] = 1e-5,
+ inplace: bool = False,
+ ) -> None:
+ super().__init__()
+ self.inplace = inplace
+ self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+ def forward(self, x: Tensor) -> Tensor:
+ return x.mul_(self.gamma) if self.inplace else x * self.gamma
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/layers/mlp.py b/torchhub/facebookresearch_dinov2_main/dinov2/layers/mlp.py
new file mode 100644
index 0000000000000000000000000000000000000000..5e4b315f972f9a9f54aef1e4ef4e81b52976f018
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/layers/mlp.py
@@ -0,0 +1,41 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
+
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+
+
+class Mlp(nn.Module):
+ def __init__(
+ self,
+ in_features: int,
+ hidden_features: Optional[int] = None,
+ out_features: Optional[int] = None,
+ act_layer: Callable[..., nn.Module] = nn.GELU,
+ drop: float = 0.0,
+ bias: bool = True,
+ ) -> None:
+ super().__init__()
+ out_features = out_features or in_features
+ hidden_features = hidden_features or in_features
+ self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+ self.act = act_layer()
+ self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+ self.drop = nn.Dropout(drop)
+
+ def forward(self, x: Tensor) -> Tensor:
+ x = self.fc1(x)
+ x = self.act(x)
+ x = self.drop(x)
+ x = self.fc2(x)
+ x = self.drop(x)
+ return x
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/layers/patch_embed.py b/torchhub/facebookresearch_dinov2_main/dinov2/layers/patch_embed.py
new file mode 100644
index 0000000000000000000000000000000000000000..574abe41175568d700a389b8b96d1ba554914779
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/layers/patch_embed.py
@@ -0,0 +1,89 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+from typing import Callable, Optional, Tuple, Union
+
+from torch import Tensor
+import torch.nn as nn
+
+
+def make_2tuple(x):
+ if isinstance(x, tuple):
+ assert len(x) == 2
+ return x
+
+ assert isinstance(x, int)
+ return (x, x)
+
+
+class PatchEmbed(nn.Module):
+ """
+ 2D image to patch embedding: (B,C,H,W) -> (B,N,D)
+
+ Args:
+ img_size: Image size.
+ patch_size: Patch token size.
+ in_chans: Number of input image channels.
+ embed_dim: Number of linear projection output channels.
+ norm_layer: Normalization layer.
+ """
+
+ def __init__(
+ self,
+ img_size: Union[int, Tuple[int, int]] = 224,
+ patch_size: Union[int, Tuple[int, int]] = 16,
+ in_chans: int = 3,
+ embed_dim: int = 768,
+ norm_layer: Optional[Callable] = None,
+ flatten_embedding: bool = True,
+ ) -> None:
+ super().__init__()
+
+ image_HW = make_2tuple(img_size)
+ patch_HW = make_2tuple(patch_size)
+ patch_grid_size = (
+ image_HW[0] // patch_HW[0],
+ image_HW[1] // patch_HW[1],
+ )
+
+ self.img_size = image_HW
+ self.patch_size = patch_HW
+ self.patches_resolution = patch_grid_size
+ self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+ self.in_chans = in_chans
+ self.embed_dim = embed_dim
+
+ self.flatten_embedding = flatten_embedding
+
+ self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
+ self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+ def forward(self, x: Tensor) -> Tensor:
+ _, _, H, W = x.shape
+ patch_H, patch_W = self.patch_size
+
+ assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
+ assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
+
+ x = self.proj(x) # B C H W
+ H, W = x.size(2), x.size(3)
+ x = x.flatten(2).transpose(1, 2) # B HW C
+ x = self.norm(x)
+ if not self.flatten_embedding:
+ x = x.reshape(-1, H, W, self.embed_dim) # B H W C
+ return x
+
+ def flops(self) -> float:
+ Ho, Wo = self.patches_resolution
+ flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+ if self.norm is not None:
+ flops += Ho * Wo * self.embed_dim
+ return flops
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/layers/swiglu_ffn.py b/torchhub/facebookresearch_dinov2_main/dinov2/layers/swiglu_ffn.py
new file mode 100644
index 0000000000000000000000000000000000000000..b3324b266fb0a50ccf8c3a0ede2ae10ac4dfa03e
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/layers/swiglu_ffn.py
@@ -0,0 +1,63 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+class SwiGLUFFN(nn.Module):
+ def __init__(
+ self,
+ in_features: int,
+ hidden_features: Optional[int] = None,
+ out_features: Optional[int] = None,
+ act_layer: Callable[..., nn.Module] = None,
+ drop: float = 0.0,
+ bias: bool = True,
+ ) -> None:
+ super().__init__()
+ out_features = out_features or in_features
+ hidden_features = hidden_features or in_features
+ self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+ self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+ def forward(self, x: Tensor) -> Tensor:
+ x12 = self.w12(x)
+ x1, x2 = x12.chunk(2, dim=-1)
+ hidden = F.silu(x1) * x2
+ return self.w3(hidden)
+
+
+try:
+ from xformers.ops import SwiGLU
+
+ XFORMERS_AVAILABLE = True
+except ImportError:
+ SwiGLU = SwiGLUFFN
+ XFORMERS_AVAILABLE = False
+
+
+class SwiGLUFFNFused(SwiGLU):
+ def __init__(
+ self,
+ in_features: int,
+ hidden_features: Optional[int] = None,
+ out_features: Optional[int] = None,
+ act_layer: Callable[..., nn.Module] = None,
+ drop: float = 0.0,
+ bias: bool = True,
+ ) -> None:
+ out_features = out_features or in_features
+ hidden_features = hidden_features or in_features
+ hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+ super().__init__(
+ in_features=in_features,
+ hidden_features=hidden_features,
+ out_features=out_features,
+ bias=bias,
+ )
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/logging/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/logging/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e80dadb2d57056e9f6f4989cd24a3c7e26fee23f
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/logging/__init__.py
@@ -0,0 +1,103 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import functools
+import logging
+import os
+import sys
+from typing import Optional
+
+import dinov2.distributed as distributed
+from .helpers import MetricLogger, SmoothedValue
+
+
+# So that calling _configure_logger multiple times won't add many handlers
+@functools.lru_cache()
+def _configure_logger(
+ name: Optional[str] = None,
+ *,
+ level: int = logging.DEBUG,
+ output: Optional[str] = None,
+):
+ """
+ Configure a logger.
+
+ Adapted from Detectron2.
+
+ Args:
+ name: The name of the logger to configure.
+ level: The logging level to use.
+ output: A file name or a directory to save log. If None, will not save log file.
+ If ends with ".txt" or ".log", assumed to be a file name.
+ Otherwise, logs will be saved to `output/log.txt`.
+
+ Returns:
+ The configured logger.
+ """
+
+ logger = logging.getLogger(name)
+ logger.setLevel(level)
+ logger.propagate = False
+
+ # Loosely match Google glog format:
+ # [IWEF]yyyymmdd hh:mm:ss.uuuuuu threadid file:line] msg
+ # but use a shorter timestamp and include the logger name:
+ # [IWEF]yyyymmdd hh:mm:ss logger threadid file:line] msg
+ fmt_prefix = "%(levelname).1s%(asctime)s %(process)s %(name)s %(filename)s:%(lineno)s] "
+ fmt_message = "%(message)s"
+ fmt = fmt_prefix + fmt_message
+ datefmt = "%Y%m%d %H:%M:%S"
+ formatter = logging.Formatter(fmt=fmt, datefmt=datefmt)
+
+ # stdout logging for main worker only
+ if distributed.is_main_process():
+ handler = logging.StreamHandler(stream=sys.stdout)
+ handler.setLevel(logging.DEBUG)
+ handler.setFormatter(formatter)
+ logger.addHandler(handler)
+
+ # file logging for all workers
+ if output:
+ if os.path.splitext(output)[-1] in (".txt", ".log"):
+ filename = output
+ else:
+ filename = os.path.join(output, "logs", "log.txt")
+
+ if not distributed.is_main_process():
+ global_rank = distributed.get_global_rank()
+ filename = filename + ".rank{}".format(global_rank)
+
+ os.makedirs(os.path.dirname(filename), exist_ok=True)
+
+ handler = logging.StreamHandler(open(filename, "a"))
+ handler.setLevel(logging.DEBUG)
+ handler.setFormatter(formatter)
+ logger.addHandler(handler)
+
+ return logger
+
+
+def setup_logging(
+ output: Optional[str] = None,
+ *,
+ name: Optional[str] = None,
+ level: int = logging.DEBUG,
+ capture_warnings: bool = True,
+) -> None:
+ """
+ Setup logging.
+
+ Args:
+ output: A file name or a directory to save log files. If None, log
+ files will not be saved. If output ends with ".txt" or ".log", it
+ is assumed to be a file name.
+ Otherwise, logs will be saved to `output/log.txt`.
+ name: The name of the logger to configure, by default the root logger.
+ level: The logging level to use.
+ capture_warnings: Whether warnings should be captured as logs.
+ """
+ logging.captureWarnings(capture_warnings)
+ _configure_logger(name, level=level, output=output)
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/logging/helpers.py b/torchhub/facebookresearch_dinov2_main/dinov2/logging/helpers.py
new file mode 100644
index 0000000000000000000000000000000000000000..16d643500d2ee10ffea5916aad07f9b9d7c0af6d
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/logging/helpers.py
@@ -0,0 +1,195 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections import defaultdict, deque
+import datetime
+import json
+import logging
+import time
+
+import torch
+
+import dinov2.distributed as distributed
+
+
+logger = logging.getLogger("dinov2")
+
+
+class MetricLogger(object):
+ def __init__(self, delimiter="\t", output_file=None):
+ self.meters = defaultdict(SmoothedValue)
+ self.delimiter = delimiter
+ self.output_file = output_file
+
+ def update(self, **kwargs):
+ for k, v in kwargs.items():
+ if isinstance(v, torch.Tensor):
+ v = v.item()
+ assert isinstance(v, (float, int))
+ self.meters[k].update(v)
+
+ def __getattr__(self, attr):
+ if attr in self.meters:
+ return self.meters[attr]
+ if attr in self.__dict__:
+ return self.__dict__[attr]
+ raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr))
+
+ def __str__(self):
+ loss_str = []
+ for name, meter in self.meters.items():
+ loss_str.append("{}: {}".format(name, str(meter)))
+ return self.delimiter.join(loss_str)
+
+ def synchronize_between_processes(self):
+ for meter in self.meters.values():
+ meter.synchronize_between_processes()
+
+ def add_meter(self, name, meter):
+ self.meters[name] = meter
+
+ def dump_in_output_file(self, iteration, iter_time, data_time):
+ if self.output_file is None or not distributed.is_main_process():
+ return
+ dict_to_dump = dict(
+ iteration=iteration,
+ iter_time=iter_time,
+ data_time=data_time,
+ )
+ dict_to_dump.update({k: v.median for k, v in self.meters.items()})
+ with open(self.output_file, "a") as f:
+ f.write(json.dumps(dict_to_dump) + "\n")
+ pass
+
+ def log_every(self, iterable, print_freq, header=None, n_iterations=None, start_iteration=0):
+ i = start_iteration
+ if not header:
+ header = ""
+ start_time = time.time()
+ end = time.time()
+ iter_time = SmoothedValue(fmt="{avg:.6f}")
+ data_time = SmoothedValue(fmt="{avg:.6f}")
+
+ if n_iterations is None:
+ n_iterations = len(iterable)
+
+ space_fmt = ":" + str(len(str(n_iterations))) + "d"
+
+ log_list = [
+ header,
+ "[{0" + space_fmt + "}/{1}]",
+ "eta: {eta}",
+ "{meters}",
+ "time: {time}",
+ "data: {data}",
+ ]
+ if torch.cuda.is_available():
+ log_list += ["max mem: {memory:.0f}"]
+
+ log_msg = self.delimiter.join(log_list)
+ MB = 1024.0 * 1024.0
+ for obj in iterable:
+ data_time.update(time.time() - end)
+ yield obj
+ iter_time.update(time.time() - end)
+ if i % print_freq == 0 or i == n_iterations - 1:
+ self.dump_in_output_file(iteration=i, iter_time=iter_time.avg, data_time=data_time.avg)
+ eta_seconds = iter_time.global_avg * (n_iterations - i)
+ eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+ if torch.cuda.is_available():
+ logger.info(
+ log_msg.format(
+ i,
+ n_iterations,
+ eta=eta_string,
+ meters=str(self),
+ time=str(iter_time),
+ data=str(data_time),
+ memory=torch.cuda.max_memory_allocated() / MB,
+ )
+ )
+ else:
+ logger.info(
+ log_msg.format(
+ i,
+ n_iterations,
+ eta=eta_string,
+ meters=str(self),
+ time=str(iter_time),
+ data=str(data_time),
+ )
+ )
+ i += 1
+ end = time.time()
+ if i >= n_iterations:
+ break
+ total_time = time.time() - start_time
+ total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+ logger.info("{} Total time: {} ({:.6f} s / it)".format(header, total_time_str, total_time / n_iterations))
+
+
+class SmoothedValue:
+ """Track a series of values and provide access to smoothed values over a
+ window or the global series average.
+ """
+
+ def __init__(self, window_size=20, fmt=None):
+ if fmt is None:
+ fmt = "{median:.4f} ({global_avg:.4f})"
+ self.deque = deque(maxlen=window_size)
+ self.total = 0.0
+ self.count = 0
+ self.fmt = fmt
+
+ def update(self, value, num=1):
+ self.deque.append(value)
+ self.count += num
+ self.total += value * num
+
+ def synchronize_between_processes(self):
+ """
+ Distributed synchronization of the metric
+ Warning: does not synchronize the deque!
+ """
+ if not distributed.is_enabled():
+ return
+ t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+ torch.distributed.barrier()
+ torch.distributed.all_reduce(t)
+ t = t.tolist()
+ self.count = int(t[0])
+ self.total = t[1]
+
+ @property
+ def median(self):
+ d = torch.tensor(list(self.deque))
+ return d.median().item()
+
+ @property
+ def avg(self):
+ d = torch.tensor(list(self.deque), dtype=torch.float32)
+ return d.mean().item()
+
+ @property
+ def global_avg(self):
+ return self.total / self.count
+
+ @property
+ def max(self):
+ return max(self.deque)
+
+ @property
+ def value(self):
+ return self.deque[-1]
+
+ def __str__(self):
+ return self.fmt.format(
+ median=self.median,
+ avg=self.avg,
+ global_avg=self.global_avg,
+ max=self.max,
+ value=self.value,
+ )
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/loss/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/loss/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..477b71b28259bf97b806df3f3d2f392dded866d6
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/loss/__init__.py
@@ -0,0 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .dino_clstoken_loss import DINOLoss
+from .ibot_patch_loss import iBOTPatchLoss
+from .koleo_loss import KoLeoLoss
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/loss/dino_clstoken_loss.py b/torchhub/facebookresearch_dinov2_main/dinov2/loss/dino_clstoken_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f33897efb1084e6c1c14ae00bc93ab332c61074
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/loss/dino_clstoken_loss.py
@@ -0,0 +1,100 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+
+
+class DINOLoss(nn.Module):
+ def __init__(
+ self,
+ out_dim,
+ student_temp=0.1,
+ center_momentum=0.9,
+ ):
+ super().__init__()
+ self.student_temp = student_temp
+ self.center_momentum = center_momentum
+ self.register_buffer("center", torch.zeros(1, out_dim))
+ self.updated = True
+ self.reduce_handle = None
+ self.len_teacher_output = None
+ self.async_batch_center = None
+
+ @torch.no_grad()
+ def softmax_center_teacher(self, teacher_output, teacher_temp):
+ self.apply_center_update()
+ # teacher centering and sharpening
+ return F.softmax((teacher_output - self.center) / teacher_temp, dim=-1)
+
+ @torch.no_grad()
+ def sinkhorn_knopp_teacher(self, teacher_output, teacher_temp, n_iterations=3):
+ teacher_output = teacher_output.float()
+ world_size = dist.get_world_size() if dist.is_initialized() else 1
+ Q = torch.exp(teacher_output / teacher_temp).t() # Q is K-by-B for consistency with notations from our paper
+ B = Q.shape[1] * world_size # number of samples to assign
+ K = Q.shape[0] # how many prototypes
+
+ # make the matrix sums to 1
+ sum_Q = torch.sum(Q)
+ if dist.is_initialized():
+ dist.all_reduce(sum_Q)
+ Q /= sum_Q
+
+ for it in range(n_iterations):
+ # normalize each row: total weight per prototype must be 1/K
+ sum_of_rows = torch.sum(Q, dim=1, keepdim=True)
+ if dist.is_initialized():
+ dist.all_reduce(sum_of_rows)
+ Q /= sum_of_rows
+ Q /= K
+
+ # normalize each column: total weight per sample must be 1/B
+ Q /= torch.sum(Q, dim=0, keepdim=True)
+ Q /= B
+
+ Q *= B # the columns must sum to 1 so that Q is an assignment
+ return Q.t()
+
+ def forward(self, student_output_list, teacher_out_softmaxed_centered_list):
+ """
+ Cross-entropy between softmax outputs of the teacher and student networks.
+ """
+ # TODO: Use cross_entropy_distribution here
+ total_loss = 0
+ for s in student_output_list:
+ lsm = F.log_softmax(s / self.student_temp, dim=-1)
+ for t in teacher_out_softmaxed_centered_list:
+ loss = torch.sum(t * lsm, dim=-1)
+ total_loss -= loss.mean()
+ return total_loss
+
+ @torch.no_grad()
+ def update_center(self, teacher_output):
+ self.reduce_center_update(teacher_output)
+
+ @torch.no_grad()
+ def reduce_center_update(self, teacher_output):
+ self.updated = False
+ self.len_teacher_output = len(teacher_output)
+ self.async_batch_center = torch.sum(teacher_output, dim=0, keepdim=True)
+ if dist.is_initialized():
+ self.reduce_handle = dist.all_reduce(self.async_batch_center, async_op=True)
+
+ @torch.no_grad()
+ def apply_center_update(self):
+ if self.updated is False:
+ world_size = dist.get_world_size() if dist.is_initialized() else 1
+
+ if self.reduce_handle is not None:
+ self.reduce_handle.wait()
+ _t = self.async_batch_center / (self.len_teacher_output * world_size)
+
+ self.center = self.center * self.center_momentum + _t * (1 - self.center_momentum)
+
+ self.updated = True
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/loss/ibot_patch_loss.py b/torchhub/facebookresearch_dinov2_main/dinov2/loss/ibot_patch_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..16bc5cf634d661f1fa337304273f60dcd43c79c3
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/loss/ibot_patch_loss.py
@@ -0,0 +1,152 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+
+import logging
+
+
+logger = logging.getLogger("dinov2")
+
+
+try:
+ from xformers.ops import cross_entropy
+
+ def lossfunc(t, s, temp):
+ s = s.float()
+ t = t.float()
+ if s.ndim == 2:
+ return -cross_entropy(s.unsqueeze(0), t.unsqueeze(0), temp, bw_inplace=True).squeeze(0)
+ elif s.ndim == 3:
+ return -cross_entropy(s, t, temp, bw_inplace=True)
+
+except ImportError:
+
+ def lossfunc(t, s, temp):
+ return torch.sum(t * F.log_softmax(s / temp, dim=-1), dim=-1)
+
+
+class iBOTPatchLoss(nn.Module):
+ def __init__(self, patch_out_dim, student_temp=0.1, center_momentum=0.9):
+ super().__init__()
+ self.student_temp = student_temp
+ self.center_momentum = center_momentum
+ self.register_buffer("center", torch.zeros(1, 1, patch_out_dim))
+ self.updated = True
+ self.reduce_handle = None
+ self.len_teacher_patch_tokens = None
+ self.async_batch_center = None
+
+ @torch.no_grad()
+ def softmax_center_teacher(self, teacher_patch_tokens, teacher_temp):
+ self.apply_center_update()
+ # teacher centering and sharpening
+ #
+ # WARNING:
+ # as self.center is a float32, everything gets casted to float32 afterwards
+ #
+ # teacher_patch_tokens = teacher_patch_tokens.float()
+ # return F.softmax((teacher_patch_tokens.sub_(self.center.to(teacher_patch_tokens.dtype))).mul_(1 / teacher_temp), dim=-1)
+
+ return F.softmax((teacher_patch_tokens - self.center) / teacher_temp, dim=-1)
+
+ # this is experimental, keep everything in float16 and let's see what happens:
+ # return F.softmax((teacher_patch_tokens.sub_(self.center)) / teacher_temp, dim=-1)
+
+ @torch.no_grad()
+ def sinkhorn_knopp_teacher(self, teacher_output, teacher_temp, n_masked_patches_tensor, n_iterations=3):
+ teacher_output = teacher_output.float()
+ # world_size = dist.get_world_size() if dist.is_initialized() else 1
+ Q = torch.exp(teacher_output / teacher_temp).t() # Q is K-by-B for consistency with notations from our paper
+ # B = Q.shape[1] * world_size # number of samples to assign
+ B = n_masked_patches_tensor
+ dist.all_reduce(B)
+ K = Q.shape[0] # how many prototypes
+
+ # make the matrix sums to 1
+ sum_Q = torch.sum(Q)
+ if dist.is_initialized():
+ dist.all_reduce(sum_Q)
+ Q /= sum_Q
+
+ for it in range(n_iterations):
+ # normalize each row: total weight per prototype must be 1/K
+ sum_of_rows = torch.sum(Q, dim=1, keepdim=True)
+ if dist.is_initialized():
+ dist.all_reduce(sum_of_rows)
+ Q /= sum_of_rows
+ Q /= K
+
+ # normalize each column: total weight per sample must be 1/B
+ Q /= torch.sum(Q, dim=0, keepdim=True)
+ Q /= B
+
+ Q *= B # the columns must sum to 1 so that Q is an assignment
+ return Q.t()
+
+ def forward(self, student_patch_tokens, teacher_patch_tokens, student_masks_flat):
+ """
+ Cross-entropy between softmax outputs of the teacher and student networks.
+ student_patch_tokens: (B, N, D) tensor
+ teacher_patch_tokens: (B, N, D) tensor
+ student_masks_flat: (B, N) tensor
+ """
+ t = teacher_patch_tokens
+ s = student_patch_tokens
+ loss = torch.sum(t * F.log_softmax(s / self.student_temp, dim=-1), dim=-1)
+ loss = torch.sum(loss * student_masks_flat.float(), dim=-1) / student_masks_flat.sum(dim=-1).clamp(min=1.0)
+ return -loss.mean()
+
+ def forward_masked(
+ self,
+ student_patch_tokens_masked,
+ teacher_patch_tokens_masked,
+ student_masks_flat,
+ n_masked_patches=None,
+ masks_weight=None,
+ ):
+ t = teacher_patch_tokens_masked
+ s = student_patch_tokens_masked
+ # loss = torch.sum(t * F.log_softmax(s / self.student_temp, dim=-1), dim=-1)
+ loss = lossfunc(t, s, self.student_temp)
+ if masks_weight is None:
+ masks_weight = (
+ (1 / student_masks_flat.sum(-1).clamp(min=1.0))
+ .unsqueeze(-1)
+ .expand_as(student_masks_flat)[student_masks_flat]
+ )
+ if n_masked_patches is not None:
+ loss = loss[:n_masked_patches]
+ loss = loss * masks_weight
+ return -loss.sum() / student_masks_flat.shape[0]
+
+ @torch.no_grad()
+ def update_center(self, teacher_patch_tokens):
+ self.reduce_center_update(teacher_patch_tokens)
+
+ @torch.no_grad()
+ def reduce_center_update(self, teacher_patch_tokens):
+ self.updated = False
+ self.len_teacher_patch_tokens = len(teacher_patch_tokens)
+ self.async_batch_center = torch.sum(teacher_patch_tokens.mean(1), dim=0, keepdim=True)
+ if dist.is_initialized():
+ self.reduce_handle = dist.all_reduce(self.async_batch_center, async_op=True)
+
+ @torch.no_grad()
+ def apply_center_update(self):
+ if self.updated is False:
+ world_size = dist.get_world_size() if dist.is_initialized() else 1
+
+ if self.reduce_handle is not None:
+ self.reduce_handle.wait()
+ _t = self.async_batch_center / (self.len_teacher_patch_tokens * world_size)
+
+ self.center = self.center * self.center_momentum + _t * (1 - self.center_momentum)
+
+ self.updated = True
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/loss/koleo_loss.py b/torchhub/facebookresearch_dinov2_main/dinov2/loss/koleo_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..e776d0426bb029cf48f25b0c94077720bc8421c4
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/loss/koleo_loss.py
@@ -0,0 +1,49 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# import torch.distributed as dist
+
+
+logger = logging.getLogger("dinov2")
+
+
+class KoLeoLoss(nn.Module):
+ """Kozachenko-Leonenko entropic loss regularizer from Sablayrolles et al. - 2018 - Spreading vectors for similarity search"""
+
+ def __init__(self):
+ super().__init__()
+ self.pdist = nn.PairwiseDistance(2, eps=1e-8)
+
+ def pairwise_NNs_inner(self, x):
+ """
+ Pairwise nearest neighbors for L2-normalized vectors.
+ Uses Torch rather than Faiss to remain on GPU.
+ """
+ # parwise dot products (= inverse distance)
+ dots = torch.mm(x, x.t())
+ n = x.shape[0]
+ dots.view(-1)[:: (n + 1)].fill_(-1) # Trick to fill diagonal with -1
+ # max inner prod -> min distance
+ _, I = torch.max(dots, dim=1) # noqa: E741
+ return I
+
+ def forward(self, student_output, eps=1e-8):
+ """
+ Args:
+ student_output (BxD): backbone output of student
+ """
+ with torch.cuda.amp.autocast(enabled=False):
+ student_output = F.normalize(student_output, eps=eps, p=2, dim=-1)
+ I = self.pairwise_NNs_inner(student_output) # noqa: E741
+ distances = self.pdist(student_output, student_output[I]) # BxD, BxD -> B
+ loss = -torch.log(distances + eps).mean()
+ return loss
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/models/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5e5a1f3832464f898752e57e865760e9864613cb
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/models/__init__.py
@@ -0,0 +1,41 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+
+from . import vision_transformer as vits
+
+
+logger = logging.getLogger("dinov2")
+
+
+def build_model(args, only_teacher=False, img_size=224):
+ args.arch = args.arch.removesuffix("_memeff")
+ if "vit" in args.arch:
+ vit_kwargs = dict(
+ img_size=img_size,
+ patch_size=args.patch_size,
+ init_values=args.layerscale,
+ ffn_layer=args.ffn_layer,
+ block_chunks=args.block_chunks,
+ qkv_bias=args.qkv_bias,
+ proj_bias=args.proj_bias,
+ ffn_bias=args.ffn_bias,
+ )
+ teacher = vits.__dict__[args.arch](**vit_kwargs)
+ if only_teacher:
+ return teacher, teacher.embed_dim
+ student = vits.__dict__[args.arch](
+ **vit_kwargs,
+ drop_path_rate=args.drop_path_rate,
+ drop_path_uniform=args.drop_path_uniform,
+ )
+ embed_dim = student.embed_dim
+ return student, teacher, embed_dim
+
+
+def build_model_from_cfg(cfg, only_teacher=False):
+ return build_model(cfg.student, only_teacher=only_teacher, img_size=cfg.crops.global_crops_size)
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/models/vision_transformer.py b/torchhub/facebookresearch_dinov2_main/dinov2/models/vision_transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..18e159a986336af813c8f0e505b946f42cd83e47
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/models/vision_transformer.py
@@ -0,0 +1,358 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+from functools import partial
+import math
+import logging
+from typing import Sequence, Tuple, Union, Callable
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from torch.nn.init import trunc_normal_
+
+from dinov2.layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block
+
+
+logger = logging.getLogger("dinov2")
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+ if not depth_first and include_root:
+ fn(module=module, name=name)
+ for child_name, child_module in module.named_children():
+ child_name = ".".join((name, child_name)) if name else child_name
+ named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+ if depth_first and include_root:
+ fn(module=module, name=name)
+ return module
+
+
+class BlockChunk(nn.ModuleList):
+ def forward(self, x):
+ for b in self:
+ x = b(x)
+ return x
+
+
+class DinoVisionTransformer(nn.Module):
+ def __init__(
+ self,
+ img_size=224,
+ patch_size=16,
+ in_chans=3,
+ embed_dim=768,
+ depth=12,
+ num_heads=12,
+ mlp_ratio=4.0,
+ qkv_bias=True,
+ ffn_bias=True,
+ proj_bias=True,
+ drop_path_rate=0.0,
+ drop_path_uniform=False,
+ init_values=None, # for layerscale: None or 0 => no layerscale
+ embed_layer=PatchEmbed,
+ act_layer=nn.GELU,
+ block_fn=Block,
+ ffn_layer="mlp",
+ block_chunks=1,
+ ):
+ """
+ Args:
+ img_size (int, tuple): input image size
+ patch_size (int, tuple): patch size
+ in_chans (int): number of input channels
+ embed_dim (int): embedding dimension
+ depth (int): depth of transformer
+ num_heads (int): number of attention heads
+ mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+ qkv_bias (bool): enable bias for qkv if True
+ proj_bias (bool): enable bias for proj in attn if True
+ ffn_bias (bool): enable bias for ffn if True
+ drop_path_rate (float): stochastic depth rate
+ drop_path_uniform (bool): apply uniform drop rate across blocks
+ weight_init (str): weight init scheme
+ init_values (float): layer-scale init values
+ embed_layer (nn.Module): patch embedding layer
+ act_layer (nn.Module): MLP activation layer
+ block_fn (nn.Module): transformer block class
+ ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+ block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
+ """
+ super().__init__()
+ norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+ self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
+ self.num_tokens = 1
+ self.n_blocks = depth
+ self.num_heads = num_heads
+ self.patch_size = patch_size
+
+ self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+ num_patches = self.patch_embed.num_patches
+
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+ self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+
+ if drop_path_uniform is True:
+ dpr = [drop_path_rate] * depth
+ else:
+ dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
+
+ if ffn_layer == "mlp":
+ logger.info("using MLP layer as FFN")
+ ffn_layer = Mlp
+ elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+ logger.info("using SwiGLU layer as FFN")
+ ffn_layer = SwiGLUFFNFused
+ elif ffn_layer == "identity":
+ logger.info("using Identity layer as FFN")
+
+ def f(*args, **kwargs):
+ return nn.Identity()
+
+ ffn_layer = f
+ else:
+ raise NotImplementedError
+
+ blocks_list = [
+ block_fn(
+ dim=embed_dim,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ proj_bias=proj_bias,
+ ffn_bias=ffn_bias,
+ drop_path=dpr[i],
+ norm_layer=norm_layer,
+ act_layer=act_layer,
+ ffn_layer=ffn_layer,
+ init_values=init_values,
+ )
+ for i in range(depth)
+ ]
+ if block_chunks > 0:
+ self.chunked_blocks = True
+ chunked_blocks = []
+ chunksize = depth // block_chunks
+ for i in range(0, depth, chunksize):
+ # this is to keep the block index consistent if we chunk the block list
+ chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
+ self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+ else:
+ self.chunked_blocks = False
+ self.blocks = nn.ModuleList(blocks_list)
+
+ self.norm = norm_layer(embed_dim)
+ self.head = nn.Identity()
+
+ self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+ self.init_weights()
+
+ def init_weights(self):
+ trunc_normal_(self.pos_embed, std=0.02)
+ nn.init.normal_(self.cls_token, std=1e-6)
+ named_apply(init_weights_vit_timm, self)
+
+ def interpolate_pos_encoding(self, x, w, h):
+ previous_dtype = x.dtype
+ npatch = x.shape[1] - 1
+ N = self.pos_embed.shape[1] - 1
+ if npatch == N and w == h:
+ return self.pos_embed
+ pos_embed = self.pos_embed.float()
+ class_pos_embed = pos_embed[:, 0]
+ patch_pos_embed = pos_embed[:, 1:]
+ dim = x.shape[-1]
+ w0 = w // self.patch_size
+ h0 = h // self.patch_size
+ # we add a small number to avoid floating point error in the interpolation
+ # see discussion at https://github.com/facebookresearch/dino/issues/8
+ w0, h0 = w0 + 0.1, h0 + 0.1
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+ scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+ mode="bicubic",
+ )
+
+ assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+ return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
+
+ def prepare_tokens_with_masks(self, x, masks=None):
+ B, nc, w, h = x.shape
+ x = self.patch_embed(x)
+ if masks is not None:
+ x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+ x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+ x = x + self.interpolate_pos_encoding(x, w, h)
+
+ return x
+
+ def forward_features_list(self, x_list, masks_list):
+ x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+ for blk in self.blocks:
+ x = blk(x)
+
+ all_x = x
+ output = []
+ for x, masks in zip(all_x, masks_list):
+ x_norm = self.norm(x)
+ output.append(
+ {
+ "x_norm_clstoken": x_norm[:, 0],
+ "x_norm_patchtokens": x_norm[:, 1:],
+ "x_prenorm": x,
+ "masks": masks,
+ }
+ )
+ return output
+
+ def forward_features(self, x, masks=None):
+ if isinstance(x, list):
+ return self.forward_features_list(x, masks)
+
+ x = self.prepare_tokens_with_masks(x, masks)
+
+ for blk in self.blocks:
+ x = blk(x)
+
+ x_norm = self.norm(x)
+ return {
+ "x_norm_clstoken": x_norm[:, 0],
+ "x_norm_patchtokens": x_norm[:, 1:],
+ "x_prenorm": x,
+ "masks": masks,
+ }
+
+ def _get_intermediate_layers_not_chunked(self, x, n=1):
+ x = self.prepare_tokens_with_masks(x)
+ # If n is an int, take the n last blocks. If it's a list, take them
+ output, total_block_len = [], len(self.blocks)
+ blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+ for i, blk in enumerate(self.blocks):
+ x = blk(x)
+ if i in blocks_to_take:
+ output.append(x)
+ assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+ return output
+
+ def _get_intermediate_layers_chunked(self, x, n=1):
+ x = self.prepare_tokens_with_masks(x)
+ output, i, total_block_len = [], 0, len(self.blocks[-1])
+ # If n is an int, take the n last blocks. If it's a list, take them
+ blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+ for block_chunk in self.blocks:
+ for blk in block_chunk[i:]: # Passing the nn.Identity()
+ x = blk(x)
+ if i in blocks_to_take:
+ output.append(x)
+ i += 1
+ assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+ return output
+
+ def get_intermediate_layers(
+ self,
+ x: torch.Tensor,
+ n: Union[int, Sequence] = 1, # Layers or n last layers to take
+ reshape: bool = False,
+ return_class_token: bool = False,
+ norm=True,
+ ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+ if self.chunked_blocks:
+ outputs = self._get_intermediate_layers_chunked(x, n)
+ else:
+ outputs = self._get_intermediate_layers_not_chunked(x, n)
+ if norm:
+ outputs = [self.norm(out) for out in outputs]
+ class_tokens = [out[:, 0] for out in outputs]
+ outputs = [out[:, 1:] for out in outputs]
+ if reshape:
+ B, _, w, h = x.shape
+ outputs = [
+ out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
+ for out in outputs
+ ]
+ if return_class_token:
+ return tuple(zip(outputs, class_tokens))
+ return tuple(outputs)
+
+ def forward(self, *args, is_training=False, **kwargs):
+ ret = self.forward_features(*args, **kwargs)
+ if is_training:
+ return ret
+ else:
+ return self.head(ret["x_norm_clstoken"])
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = ""):
+ """ViT weight initialization, original timm impl (for reproducibility)"""
+ if isinstance(module, nn.Linear):
+ trunc_normal_(module.weight, std=0.02)
+ if module.bias is not None:
+ nn.init.zeros_(module.bias)
+
+
+def vit_small(patch_size=16, **kwargs):
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=384,
+ depth=12,
+ num_heads=6,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ **kwargs,
+ )
+ return model
+
+
+def vit_base(patch_size=16, **kwargs):
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=768,
+ depth=12,
+ num_heads=12,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ **kwargs,
+ )
+ return model
+
+
+def vit_large(patch_size=16, **kwargs):
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=1024,
+ depth=24,
+ num_heads=16,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ **kwargs,
+ )
+ return model
+
+
+def vit_giant2(patch_size=16, **kwargs):
+ """
+ Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
+ """
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=1536,
+ depth=40,
+ num_heads=24,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ **kwargs,
+ )
+ return model
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/run/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/run/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0952fcc3f57e34b3747962e9ebd6fc57aeea63fa
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/run/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/run/eval/knn.py b/torchhub/facebookresearch_dinov2_main/dinov2/run/eval/knn.py
new file mode 100644
index 0000000000000000000000000000000000000000..15d674b78b0629aa0f041c2426c894925469a0e8
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/run/eval/knn.py
@@ -0,0 +1,60 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import os
+import sys
+
+from dinov2.eval.knn import get_args_parser as get_knn_args_parser
+from dinov2.logging import setup_logging
+from dinov2.run.submit import get_args_parser, submit_jobs
+
+
+logger = logging.getLogger("dinov2")
+
+
+class Evaluator:
+ def __init__(self, args):
+ self.args = args
+
+ def __call__(self):
+ from dinov2.eval.knn import main as knn_main
+
+ self._setup_args()
+ knn_main(self.args)
+
+ def checkpoint(self):
+ import submitit
+
+ logger.info(f"Requeuing {self.args}")
+ empty = type(self)(self.args)
+ return submitit.helpers.DelayedSubmission(empty)
+
+ def _setup_args(self):
+ import submitit
+
+ job_env = submitit.JobEnvironment()
+ self.args.output_dir = self.args.output_dir.replace("%j", str(job_env.job_id))
+ logger.info(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")
+ logger.info(f"Args: {self.args}")
+
+
+def main():
+ description = "Submitit launcher for DINOv2 k-NN evaluation"
+ knn_args_parser = get_knn_args_parser(add_help=False)
+ parents = [knn_args_parser]
+ args_parser = get_args_parser(description=description, parents=parents)
+ args = args_parser.parse_args()
+
+ setup_logging()
+
+ assert os.path.exists(args.config_file), "Configuration file does not exist!"
+ submit_jobs(Evaluator, args, name="dinov2:knn")
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/run/eval/linear.py b/torchhub/facebookresearch_dinov2_main/dinov2/run/eval/linear.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8c264762ac6bb82a3622c74e1e683ea5c6be437
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/run/eval/linear.py
@@ -0,0 +1,60 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import os
+import sys
+
+from dinov2.eval.linear import get_args_parser as get_linear_args_parser
+from dinov2.logging import setup_logging
+from dinov2.run.submit import get_args_parser, submit_jobs
+
+
+logger = logging.getLogger("dinov2")
+
+
+class Evaluator:
+ def __init__(self, args):
+ self.args = args
+
+ def __call__(self):
+ from dinov2.eval.linear import main as linear_main
+
+ self._setup_args()
+ linear_main(self.args)
+
+ def checkpoint(self):
+ import submitit
+
+ logger.info(f"Requeuing {self.args}")
+ empty = type(self)(self.args)
+ return submitit.helpers.DelayedSubmission(empty)
+
+ def _setup_args(self):
+ import submitit
+
+ job_env = submitit.JobEnvironment()
+ self.args.output_dir = self.args.output_dir.replace("%j", str(job_env.job_id))
+ logger.info(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")
+ logger.info(f"Args: {self.args}")
+
+
+def main():
+ description = "Submitit launcher for DINOv2 linear evaluation"
+ linear_args_parser = get_linear_args_parser(add_help=False)
+ parents = [linear_args_parser]
+ args_parser = get_args_parser(description=description, parents=parents)
+ args = args_parser.parse_args()
+
+ setup_logging()
+
+ assert os.path.exists(args.config_file), "Configuration file does not exist!"
+ submit_jobs(Evaluator, args, name="dinov2:linear")
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/run/eval/log_regression.py b/torchhub/facebookresearch_dinov2_main/dinov2/run/eval/log_regression.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d3d5a5742792fc8d4ca3b39c15c47e8aa349bc7
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/run/eval/log_regression.py
@@ -0,0 +1,60 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import os
+import sys
+
+from dinov2.eval.log_regression import get_args_parser as get_log_regression_args_parser
+from dinov2.logging import setup_logging
+from dinov2.run.submit import get_args_parser, submit_jobs
+
+
+logger = logging.getLogger("dinov2")
+
+
+class Evaluator:
+ def __init__(self, args):
+ self.args = args
+
+ def __call__(self):
+ from dinov2.eval.log_regression import main as log_regression_main
+
+ self._setup_args()
+ log_regression_main(self.args)
+
+ def checkpoint(self):
+ import submitit
+
+ logger.info(f"Requeuing {self.args}")
+ empty = type(self)(self.args)
+ return submitit.helpers.DelayedSubmission(empty)
+
+ def _setup_args(self):
+ import submitit
+
+ job_env = submitit.JobEnvironment()
+ self.args.output_dir = self.args.output_dir.replace("%j", str(job_env.job_id))
+ logger.info(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")
+ logger.info(f"Args: {self.args}")
+
+
+def main():
+ description = "Submitit launcher for DINOv2 logistic evaluation"
+ log_regression_args_parser = get_log_regression_args_parser(add_help=False)
+ parents = [log_regression_args_parser]
+ args_parser = get_args_parser(description=description, parents=parents)
+ args = args_parser.parse_args()
+
+ setup_logging()
+
+ assert os.path.exists(args.config_file), "Configuration file does not exist!"
+ submit_jobs(Evaluator, args, name="dinov2:logreg")
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/run/submit.py b/torchhub/facebookresearch_dinov2_main/dinov2/run/submit.py
new file mode 100644
index 0000000000000000000000000000000000000000..68140f3d6d93dc67ccd7c45fe712eb15483d1ad6
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/run/submit.py
@@ -0,0 +1,123 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import logging
+import os
+from pathlib import Path
+from typing import List, Optional
+
+import submitit
+
+from dinov2.utils.cluster import (
+ get_slurm_executor_parameters,
+ get_slurm_partition,
+ get_user_checkpoint_path,
+)
+
+
+logger = logging.getLogger("dinov2")
+
+
+def get_args_parser(
+ description: Optional[str] = None,
+ parents: Optional[List[argparse.ArgumentParser]] = None,
+ add_help: bool = True,
+) -> argparse.ArgumentParser:
+ parents = parents or []
+ slurm_partition = get_slurm_partition()
+ parser = argparse.ArgumentParser(
+ description=description,
+ parents=parents,
+ add_help=add_help,
+ )
+ parser.add_argument(
+ "--ngpus",
+ "--gpus",
+ "--gpus-per-node",
+ default=8,
+ type=int,
+ help="Number of GPUs to request on each node",
+ )
+ parser.add_argument(
+ "--nodes",
+ "--nnodes",
+ default=2,
+ type=int,
+ help="Number of nodes to request",
+ )
+ parser.add_argument(
+ "--timeout",
+ default=2800,
+ type=int,
+ help="Duration of the job",
+ )
+ parser.add_argument(
+ "--partition",
+ default=slurm_partition,
+ type=str,
+ help="Partition where to submit",
+ )
+ parser.add_argument(
+ "--use-volta32",
+ action="store_true",
+ help="Request V100-32GB GPUs",
+ )
+ parser.add_argument(
+ "--comment",
+ default="",
+ type=str,
+ help="Comment to pass to scheduler, e.g. priority message",
+ )
+ parser.add_argument(
+ "--exclude",
+ default="",
+ type=str,
+ help="Nodes to exclude",
+ )
+ return parser
+
+
+def get_shared_folder() -> Path:
+ user_checkpoint_path = get_user_checkpoint_path()
+ if user_checkpoint_path is None:
+ raise RuntimeError("Path to user checkpoint cannot be determined")
+ path = user_checkpoint_path / "experiments"
+ path.mkdir(exist_ok=True)
+ return path
+
+
+def submit_jobs(task_class, args, name: str):
+ if not args.output_dir:
+ args.output_dir = str(get_shared_folder() / "%j")
+
+ Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+ executor = submitit.AutoExecutor(folder=args.output_dir, slurm_max_num_timeout=30)
+
+ kwargs = {}
+ if args.use_volta32:
+ kwargs["slurm_constraint"] = "volta32gb"
+ if args.comment:
+ kwargs["slurm_comment"] = args.comment
+ if args.exclude:
+ kwargs["slurm_exclude"] = args.exclude
+
+ executor_params = get_slurm_executor_parameters(
+ nodes=args.nodes,
+ num_gpus_per_node=args.ngpus,
+ timeout_min=args.timeout, # max is 60 * 72
+ slurm_signal_delay_s=120,
+ slurm_partition=args.partition,
+ **kwargs,
+ )
+ executor.update_parameters(name=name, **executor_params)
+
+ task = task_class(args)
+ job = executor.submit(task)
+
+ logger.info(f"Submitted job_id: {job.job_id}")
+ str_output_dir = os.path.abspath(args.output_dir).replace("%j", str(job.job_id))
+ logger.info(f"Logs and checkpoints will be saved at: {str_output_dir}")
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/run/train/train.py b/torchhub/facebookresearch_dinov2_main/dinov2/run/train/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..24716f2a314820a4cc15289fe0cb13ad52cf343c
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/run/train/train.py
@@ -0,0 +1,60 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import os
+import sys
+
+from dinov2.logging import setup_logging
+from dinov2.train import get_args_parser as get_train_args_parser
+from dinov2.run.submit import get_args_parser, submit_jobs
+
+
+logger = logging.getLogger("dinov2")
+
+
+class Trainer(object):
+ def __init__(self, args):
+ self.args = args
+
+ def __call__(self):
+ from dinov2.train import main as train_main
+
+ self._setup_args()
+ train_main(self.args)
+
+ def checkpoint(self):
+ import submitit
+
+ logger.info(f"Requeuing {self.args}")
+ empty = type(self)(self.args)
+ return submitit.helpers.DelayedSubmission(empty)
+
+ def _setup_args(self):
+ import submitit
+
+ job_env = submitit.JobEnvironment()
+ self.args.output_dir = self.args.output_dir.replace("%j", str(job_env.job_id))
+ logger.info(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")
+ logger.info(f"Args: {self.args}")
+
+
+def main():
+ description = "Submitit launcher for DINOv2 training"
+ train_args_parser = get_train_args_parser(add_help=False)
+ parents = [train_args_parser]
+ args_parser = get_args_parser(description=description, parents=parents)
+ args = args_parser.parse_args()
+
+ setup_logging()
+
+ assert os.path.exists(args.config_file), "Configuration file does not exist!"
+ submit_jobs(Trainer, args, name="dinov2:train")
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/train/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/train/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0b66d17aa547ed5560e75a03f5c1587da2d4fd7
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/train/__init__.py
@@ -0,0 +1,8 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .train import get_args_parser, main
+from .ssl_meta_arch import SSLMetaArch
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/train/ssl_meta_arch.py b/torchhub/facebookresearch_dinov2_main/dinov2/train/ssl_meta_arch.py
new file mode 100644
index 0000000000000000000000000000000000000000..86d0c2413f9abc61953d0e12b43a5a843d97d244
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/train/ssl_meta_arch.py
@@ -0,0 +1,403 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from functools import partial
+import logging
+
+import torch
+from torch import nn
+
+from dinov2.loss import DINOLoss, iBOTPatchLoss, KoLeoLoss
+from dinov2.models import build_model_from_cfg
+from dinov2.layers import DINOHead
+from dinov2.utils.utils import has_batchnorms
+from dinov2.utils.param_groups import get_params_groups_with_decay, fuse_params_groups
+from dinov2.fsdp import get_fsdp_wrapper, ShardedGradScaler, get_fsdp_modules, reshard_fsdp_model
+
+from dinov2.models.vision_transformer import BlockChunk
+
+try:
+ from xformers.ops import fmha
+
+ XFORMERS_AVAILABLE = True
+except ImportError:
+ XFORMERS_AVAILABLE = False
+assert XFORMERS_AVAILABLE, "xFormers is required for DINOv2 training"
+
+
+logger = logging.getLogger("dinov2")
+
+
+class SSLMetaArch(nn.Module):
+ def __init__(self, cfg):
+ super().__init__()
+ self.cfg = cfg
+ self.fp16_scaler = ShardedGradScaler() if cfg.compute_precision.grad_scaler else None
+
+ student_model_dict = dict()
+ teacher_model_dict = dict()
+
+ student_backbone, teacher_backbone, embed_dim = build_model_from_cfg(cfg)
+ student_model_dict["backbone"] = student_backbone
+ teacher_model_dict["backbone"] = teacher_backbone
+ logger.info(f"OPTIONS -- architecture : embed_dim: {embed_dim}")
+
+ if cfg.student.pretrained_weights:
+ chkpt = torch.load(cfg.student.pretrained_weights)
+ logger.info(f"OPTIONS -- pretrained weights: loading from {cfg.student.pretrained_weights}")
+ student_backbone.load_state_dict(chkpt["model"], strict=False)
+
+ self.embed_dim = embed_dim
+ self.dino_out_dim = cfg.dino.head_n_prototypes
+
+ self.do_dino = cfg.dino.loss_weight > 0
+ self.do_koleo = cfg.dino.koleo_loss_weight > 0
+ self.do_ibot = cfg.ibot.loss_weight > 0
+ self.ibot_separate_head = cfg.ibot.separate_head
+
+ logger.info("OPTIONS -- DINO")
+ if self.do_dino:
+ logger.info(f"OPTIONS -- DINO -- loss_weight: {cfg.dino.loss_weight}")
+ logger.info(f"OPTIONS -- DINO -- head_n_prototypes: {cfg.dino.head_n_prototypes}")
+ logger.info(f"OPTIONS -- DINO -- head_bottleneck_dim: {cfg.dino.head_bottleneck_dim}")
+ logger.info(f"OPTIONS -- DINO -- head_hidden_dim: {cfg.dino.head_hidden_dim}")
+ self.dino_loss_weight = cfg.dino.loss_weight
+ dino_head = partial(
+ DINOHead,
+ in_dim=embed_dim,
+ out_dim=cfg.dino.head_n_prototypes,
+ hidden_dim=cfg.dino.head_hidden_dim,
+ bottleneck_dim=cfg.dino.head_bottleneck_dim,
+ nlayers=cfg.dino.head_nlayers,
+ )
+ self.dino_loss = DINOLoss(self.dino_out_dim)
+ if self.do_koleo:
+ logger.info("OPTIONS -- DINO -- applying KOLEO regularization")
+ self.koleo_loss = KoLeoLoss()
+
+ else:
+ logger.info("OPTIONS -- DINO -- not using DINO")
+
+ if self.do_dino or self.do_ibot:
+ student_model_dict["dino_head"] = dino_head()
+ teacher_model_dict["dino_head"] = dino_head()
+
+ logger.info("OPTIONS -- IBOT")
+ logger.info(f"OPTIONS -- IBOT -- loss_weight: {cfg.ibot.loss_weight}")
+ logger.info(f"OPTIONS -- IBOT masking -- ibot_mask_ratio_tuple: {cfg.ibot.mask_ratio_min_max}")
+ logger.info(f"OPTIONS -- IBOT masking -- ibot_mask_sample_probability: {cfg.ibot.mask_sample_probability}")
+ if self.do_ibot:
+ self.ibot_loss_weight = cfg.ibot.loss_weight
+ assert max(cfg.ibot.mask_ratio_min_max) > 0, "please provide a positive mask ratio tuple for ibot"
+ assert cfg.ibot.mask_sample_probability > 0, "please provide a positive mask probability for ibot"
+ self.ibot_out_dim = cfg.ibot.head_n_prototypes if self.ibot_separate_head else cfg.dino.head_n_prototypes
+ self.ibot_patch_loss = iBOTPatchLoss(self.ibot_out_dim)
+ if self.ibot_separate_head:
+ logger.info(f"OPTIONS -- IBOT -- loss_weight: {cfg.ibot.loss_weight}")
+ logger.info(f"OPTIONS -- IBOT -- head_n_prototypes: {cfg.ibot.head_n_prototypes}")
+ logger.info(f"OPTIONS -- IBOT -- head_bottleneck_dim: {cfg.ibot.head_bottleneck_dim}")
+ logger.info(f"OPTIONS -- IBOT -- head_hidden_dim: {cfg.ibot.head_hidden_dim}")
+ ibot_head = partial(
+ DINOHead,
+ in_dim=embed_dim,
+ out_dim=cfg.ibot.head_n_prototypes,
+ hidden_dim=cfg.ibot.head_hidden_dim,
+ bottleneck_dim=cfg.ibot.head_bottleneck_dim,
+ nlayers=cfg.ibot.head_nlayers,
+ )
+ student_model_dict["ibot_head"] = ibot_head()
+ teacher_model_dict["ibot_head"] = ibot_head()
+ else:
+ logger.info("OPTIONS -- IBOT -- head shared with DINO")
+
+ self.need_to_synchronize_fsdp_streams = True
+
+ self.student = nn.ModuleDict(student_model_dict)
+ self.teacher = nn.ModuleDict(teacher_model_dict)
+
+ # there is no backpropagation through the teacher, so no need for gradients
+ for p in self.teacher.parameters():
+ p.requires_grad = False
+ logger.info(f"Student and Teacher are built: they are both {cfg.student.arch} network.")
+
+ def forward(self, inputs):
+ raise NotImplementedError
+
+ def backprop_loss(self, loss):
+ if self.fp16_scaler is not None:
+ self.fp16_scaler.scale(loss).backward()
+ else:
+ loss.backward()
+
+ def forward_backward(self, images, teacher_temp):
+ n_global_crops = 2
+ assert n_global_crops == 2
+ n_local_crops = self.cfg.crops.local_crops_number
+
+ global_crops = images["collated_global_crops"].cuda(non_blocking=True)
+ local_crops = images["collated_local_crops"].cuda(non_blocking=True)
+
+ masks = images["collated_masks"].cuda(non_blocking=True)
+ mask_indices_list = images["mask_indices_list"].cuda(non_blocking=True)
+ n_masked_patches_tensor = images["n_masked_patches"].cuda(non_blocking=True)
+ n_masked_patches = mask_indices_list.shape[0]
+ upperbound = images["upperbound"]
+ masks_weight = images["masks_weight"].cuda(non_blocking=True)
+
+ n_local_crops_loss_terms = max(n_local_crops * n_global_crops, 1)
+ n_global_crops_loss_terms = (n_global_crops - 1) * n_global_crops
+
+ do_dino = self.do_dino
+ do_ibot = self.do_ibot
+
+ # loss scales
+ ibot_loss_scale = 1.0 / n_global_crops
+
+ # teacher output
+ @torch.no_grad()
+ def get_teacher_output():
+ x, n_global_crops_teacher = global_crops, n_global_crops
+ teacher_backbone_output_dict = self.teacher.backbone(x, is_training=True)
+ teacher_cls_tokens = teacher_backbone_output_dict["x_norm_clstoken"]
+ teacher_cls_tokens = teacher_cls_tokens.chunk(n_global_crops_teacher)
+ # watch out: these are chunked and cat'd in reverse so A is matched to B in the global crops dino loss
+ teacher_cls_tokens = torch.cat((teacher_cls_tokens[1], teacher_cls_tokens[0]))
+ ibot_teacher_patch_tokens = teacher_backbone_output_dict["x_norm_patchtokens"]
+ _dim = ibot_teacher_patch_tokens.shape[-1]
+ n_cls_tokens = teacher_cls_tokens.shape[0]
+
+ if do_ibot and not self.ibot_separate_head:
+ buffer_tensor_teacher = ibot_teacher_patch_tokens.new_zeros(upperbound + n_cls_tokens, _dim)
+ buffer_tensor_teacher[:n_cls_tokens].copy_(teacher_cls_tokens)
+ torch.index_select(
+ ibot_teacher_patch_tokens.flatten(0, 1),
+ dim=0,
+ index=mask_indices_list,
+ out=buffer_tensor_teacher[n_cls_tokens : n_cls_tokens + n_masked_patches],
+ )
+ tokens_after_head = self.teacher.dino_head(buffer_tensor_teacher)
+ teacher_cls_tokens_after_head = tokens_after_head[:n_cls_tokens]
+ masked_teacher_patch_tokens_after_head = tokens_after_head[
+ n_cls_tokens : n_cls_tokens + n_masked_patches
+ ]
+ elif do_ibot and self.ibot_separate_head:
+ buffer_tensor_teacher = ibot_teacher_patch_tokens.new_zeros(upperbound, _dim)
+ torch.index_select(
+ ibot_teacher_patch_tokens.flatten(0, 1),
+ dim=0,
+ index=mask_indices_list,
+ out=buffer_tensor_teacher[:n_masked_patches],
+ )
+ teacher_cls_tokens_after_head = self.teacher.dino_head(teacher_cls_tokens)
+ masked_teacher_patch_tokens_after_head = self.teacher.ibot_head(buffer_tensor_teacher)[
+ :n_masked_patches
+ ]
+ else:
+ teacher_cls_tokens_after_head = self.teacher.dino_head(teacher_cls_tokens)
+ masked_teacher_ibot_softmaxed_centered = None
+
+ if self.cfg.train.centering == "centering":
+ teacher_dino_softmaxed_centered_list = self.dino_loss.softmax_center_teacher(
+ teacher_cls_tokens_after_head, teacher_temp=teacher_temp
+ ).view(n_global_crops_teacher, -1, *teacher_cls_tokens_after_head.shape[1:])
+ self.dino_loss.update_center(teacher_cls_tokens_after_head)
+ if do_ibot:
+ masked_teacher_patch_tokens_after_head = masked_teacher_patch_tokens_after_head.unsqueeze(0)
+ masked_teacher_ibot_softmaxed_centered = self.ibot_patch_loss.softmax_center_teacher(
+ masked_teacher_patch_tokens_after_head[:, :n_masked_patches], teacher_temp=teacher_temp
+ )
+ masked_teacher_ibot_softmaxed_centered = masked_teacher_ibot_softmaxed_centered.squeeze(0)
+ self.ibot_patch_loss.update_center(masked_teacher_patch_tokens_after_head[:n_masked_patches])
+
+ elif self.cfg.train.centering == "sinkhorn_knopp":
+ teacher_dino_softmaxed_centered_list = self.dino_loss.sinkhorn_knopp_teacher(
+ teacher_cls_tokens_after_head, teacher_temp=teacher_temp
+ ).view(n_global_crops_teacher, -1, *teacher_cls_tokens_after_head.shape[1:])
+
+ if do_ibot:
+ masked_teacher_ibot_softmaxed_centered = self.ibot_patch_loss.sinkhorn_knopp_teacher(
+ masked_teacher_patch_tokens_after_head,
+ teacher_temp=teacher_temp,
+ n_masked_patches_tensor=n_masked_patches_tensor,
+ )
+
+ else:
+ raise NotImplementedError
+
+ return teacher_dino_softmaxed_centered_list, masked_teacher_ibot_softmaxed_centered
+
+ teacher_dino_softmaxed_centered_list, masked_teacher_ibot_softmaxed_centered = get_teacher_output()
+ reshard_fsdp_model(self.teacher)
+
+ loss_dict = {}
+
+ loss_accumulator = 0 # for backprop
+ student_global_backbone_output_dict, student_local_backbone_output_dict = self.student.backbone(
+ [global_crops, local_crops], masks=[masks, None], is_training=True
+ )
+
+ inputs_for_student_head_list = []
+
+ # 1a: local crops cls tokens
+ student_local_cls_tokens = student_local_backbone_output_dict["x_norm_clstoken"]
+ inputs_for_student_head_list.append(student_local_cls_tokens.unsqueeze(0))
+
+ # 1b: global crops cls tokens
+ student_global_cls_tokens = student_global_backbone_output_dict["x_norm_clstoken"]
+ inputs_for_student_head_list.append(student_global_cls_tokens.unsqueeze(0))
+
+ # 1c: global crops patch tokens
+ if do_ibot:
+ _dim = student_global_backbone_output_dict["x_norm_clstoken"].shape[-1]
+ ibot_student_patch_tokens = student_global_backbone_output_dict["x_norm_patchtokens"]
+ buffer_tensor_patch_tokens = ibot_student_patch_tokens.new_zeros(upperbound, _dim)
+ buffer_tensor_patch_tokens[:n_masked_patches].copy_(
+ torch.index_select(ibot_student_patch_tokens.flatten(0, 1), dim=0, index=mask_indices_list)
+ )
+ if not self.ibot_separate_head:
+ inputs_for_student_head_list.append(buffer_tensor_patch_tokens.unsqueeze(0))
+ else:
+ student_global_masked_patch_tokens_after_head = self.student.ibot_head(buffer_tensor_patch_tokens)[
+ :n_masked_patches
+ ]
+
+ # 2: run
+ _attn_bias, cat_inputs = fmha.BlockDiagonalMask.from_tensor_list(inputs_for_student_head_list)
+ outputs_list = _attn_bias.split(self.student.dino_head(cat_inputs))
+
+ # 3a: local crops cls tokens
+ student_local_cls_tokens_after_head = outputs_list.pop(0).squeeze(0)
+
+ # 3b: global crops cls tokens
+ student_global_cls_tokens_after_head = outputs_list.pop(0).squeeze(0)
+
+ # 3c: global crops patch tokens
+ if do_ibot and not self.ibot_separate_head:
+ student_global_masked_patch_tokens_after_head = outputs_list.pop(0).squeeze(0)[:n_masked_patches]
+
+ if n_local_crops > 0:
+ dino_local_crops_loss = self.dino_loss(
+ student_output_list=student_local_cls_tokens_after_head.chunk(n_local_crops),
+ teacher_out_softmaxed_centered_list=teacher_dino_softmaxed_centered_list,
+ ) / (n_global_crops_loss_terms + n_local_crops_loss_terms)
+
+ # store for display
+ loss_dict["dino_local_crops_loss"] = dino_local_crops_loss
+
+ # accumulate loss
+ loss_accumulator += self.dino_loss_weight * dino_local_crops_loss
+
+ # process global crops
+ loss_scales = 2 # this is here since we process global crops together
+
+ if do_dino:
+ # compute loss
+ dino_global_crops_loss = (
+ self.dino_loss(
+ student_output_list=[student_global_cls_tokens_after_head],
+ teacher_out_softmaxed_centered_list=[
+ teacher_dino_softmaxed_centered_list.flatten(0, 1)
+ ], # these were chunked and stacked in reverse so A is matched to B
+ )
+ * loss_scales
+ / (n_global_crops_loss_terms + n_local_crops_loss_terms)
+ )
+
+ loss_dict["dino_global_crops_loss"] = dino_global_crops_loss
+
+ # accumulate loss
+ loss_accumulator += self.dino_loss_weight * dino_global_crops_loss
+
+ student_cls_tokens = student_global_cls_tokens
+
+ if self.do_koleo:
+ koleo_loss = self.cfg.dino.koleo_loss_weight * sum(
+ self.koleo_loss(p) for p in student_cls_tokens.chunk(2)
+ ) # we don't apply koleo loss between cls tokens of a same image
+ loss_accumulator += koleo_loss
+ loss_dict["koleo_loss"] = (
+ koleo_loss / loss_scales
+ ) # this is to display the same losses as before but we can remove eventually
+
+ if do_ibot:
+ # compute loss
+ ibot_patch_loss = (
+ self.ibot_patch_loss.forward_masked(
+ student_global_masked_patch_tokens_after_head,
+ masked_teacher_ibot_softmaxed_centered,
+ student_masks_flat=masks,
+ n_masked_patches=n_masked_patches,
+ masks_weight=masks_weight,
+ )
+ * loss_scales
+ * ibot_loss_scale
+ )
+
+ # store for display
+ loss_dict["ibot_loss"] = ibot_patch_loss / 2
+
+ # accumulate loss
+ loss_accumulator += self.ibot_loss_weight * ibot_patch_loss
+
+ self.backprop_loss(loss_accumulator)
+
+ self.fsdp_synchronize_streams()
+
+ return loss_dict
+
+ def fsdp_synchronize_streams(self):
+ if self.need_to_synchronize_fsdp_streams:
+ torch.cuda.synchronize()
+ self.student.dino_head._streams = (
+ self.teacher.dino_head._streams
+ ) = self.student.backbone._streams = self.teacher.backbone._streams
+ self.need_to_synchronize_fsdp_streams = False
+
+ def update_teacher(self, m):
+ student_param_list = []
+ teacher_param_list = []
+ with torch.no_grad():
+ for k in self.student.keys():
+ for ms, mt in zip(get_fsdp_modules(self.student[k]), get_fsdp_modules(self.teacher[k])):
+ student_param_list += ms.params
+ teacher_param_list += mt.params
+ torch._foreach_mul_(teacher_param_list, m)
+ torch._foreach_add_(teacher_param_list, student_param_list, alpha=1 - m)
+
+ def train(self):
+ super().train()
+ self.teacher.eval()
+
+ def get_maybe_fused_params_for_submodel(self, m):
+ params_groups = get_params_groups_with_decay(
+ model=m,
+ lr_decay_rate=self.cfg.optim.layerwise_decay,
+ patch_embed_lr_mult=self.cfg.optim.patch_embed_lr_mult,
+ )
+ fused_params_groups = fuse_params_groups(params_groups)
+ logger.info("fusing param groups")
+
+ for g in fused_params_groups:
+ g["foreach"] = True
+ return fused_params_groups
+
+ def get_params_groups(self):
+ all_params_groups = []
+ for m in self.student.values():
+ all_params_groups += self.get_maybe_fused_params_for_submodel(m)
+ return all_params_groups
+
+ def prepare_for_distributed_training(self):
+ logger.info("DISTRIBUTED FSDP -- preparing model for distributed training")
+ if has_batchnorms(self.student):
+ raise NotImplementedError
+ # below will synchronize all student subnetworks across gpus:
+ for k, v in self.student.items():
+ self.teacher[k].load_state_dict(self.student[k].state_dict())
+ student_model_cfg = self.cfg.compute_precision.student[k]
+ self.student[k] = get_fsdp_wrapper(student_model_cfg, modules_to_wrap={BlockChunk})(self.student[k])
+ teacher_model_cfg = self.cfg.compute_precision.teacher[k]
+ self.teacher[k] = get_fsdp_wrapper(teacher_model_cfg, modules_to_wrap={BlockChunk})(self.teacher[k])
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/train/train.py b/torchhub/facebookresearch_dinov2_main/dinov2/train/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..5279b9c4317e56b5c0a9c39f7bf9bf56b04a1f8b
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/train/train.py
@@ -0,0 +1,319 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import logging
+import math
+import os
+from functools import partial
+
+from fvcore.common.checkpoint import PeriodicCheckpointer
+import torch
+
+from dinov2.data import SamplerType, make_data_loader, make_dataset
+from dinov2.data import collate_data_and_cast, DataAugmentationDINO, MaskingGenerator
+import dinov2.distributed as distributed
+from dinov2.fsdp import FSDPCheckpointer
+from dinov2.logging import MetricLogger
+from dinov2.utils.config import setup
+from dinov2.utils.utils import CosineScheduler
+
+from dinov2.train.ssl_meta_arch import SSLMetaArch
+
+
+torch.backends.cuda.matmul.allow_tf32 = True # PyTorch 1.12 sets this to False by default
+logger = logging.getLogger("dinov2")
+
+
+def get_args_parser(add_help: bool = True):
+ parser = argparse.ArgumentParser("DINOv2 training", add_help=add_help)
+ parser.add_argument("--config-file", default="", metavar="FILE", help="path to config file")
+ parser.add_argument(
+ "--no-resume",
+ action="store_true",
+ help="Whether to not attempt to resume from the checkpoint directory. ",
+ )
+ parser.add_argument("--eval-only", action="store_true", help="perform evaluation only")
+ parser.add_argument("--eval", type=str, default="", help="Eval type to perform")
+ parser.add_argument(
+ "opts",
+ help="""
+Modify config options at the end of the command. For Yacs configs, use
+space-separated "PATH.KEY VALUE" pairs.
+For python-based LazyConfig, use "path.key=value".
+ """.strip(),
+ default=None,
+ nargs=argparse.REMAINDER,
+ )
+ parser.add_argument(
+ "--output-dir",
+ "--output_dir",
+ default="",
+ type=str,
+ help="Output directory to save logs and checkpoints",
+ )
+
+ return parser
+
+
+def build_optimizer(cfg, params_groups):
+ return torch.optim.AdamW(params_groups, betas=(cfg.optim.adamw_beta1, cfg.optim.adamw_beta2))
+
+
+def build_schedulers(cfg):
+ OFFICIAL_EPOCH_LENGTH = cfg.train.OFFICIAL_EPOCH_LENGTH
+ lr = dict(
+ base_value=cfg.optim["lr"],
+ final_value=cfg.optim["min_lr"],
+ total_iters=cfg.optim["epochs"] * OFFICIAL_EPOCH_LENGTH,
+ warmup_iters=cfg.optim["warmup_epochs"] * OFFICIAL_EPOCH_LENGTH,
+ start_warmup_value=0,
+ )
+ wd = dict(
+ base_value=cfg.optim["weight_decay"],
+ final_value=cfg.optim["weight_decay_end"],
+ total_iters=cfg.optim["epochs"] * OFFICIAL_EPOCH_LENGTH,
+ )
+ momentum = dict(
+ base_value=cfg.teacher["momentum_teacher"],
+ final_value=cfg.teacher["final_momentum_teacher"],
+ total_iters=cfg.optim["epochs"] * OFFICIAL_EPOCH_LENGTH,
+ )
+ teacher_temp = dict(
+ base_value=cfg.teacher["teacher_temp"],
+ final_value=cfg.teacher["teacher_temp"],
+ total_iters=cfg.teacher["warmup_teacher_temp_epochs"] * OFFICIAL_EPOCH_LENGTH,
+ warmup_iters=cfg.teacher["warmup_teacher_temp_epochs"] * OFFICIAL_EPOCH_LENGTH,
+ start_warmup_value=cfg.teacher["warmup_teacher_temp"],
+ )
+
+ lr_schedule = CosineScheduler(**lr)
+ wd_schedule = CosineScheduler(**wd)
+ momentum_schedule = CosineScheduler(**momentum)
+ teacher_temp_schedule = CosineScheduler(**teacher_temp)
+ last_layer_lr_schedule = CosineScheduler(**lr)
+
+ last_layer_lr_schedule.schedule[
+ : cfg.optim["freeze_last_layer_epochs"] * OFFICIAL_EPOCH_LENGTH
+ ] = 0 # mimicking the original schedules
+
+ logger.info("Schedulers ready.")
+
+ return (
+ lr_schedule,
+ wd_schedule,
+ momentum_schedule,
+ teacher_temp_schedule,
+ last_layer_lr_schedule,
+ )
+
+
+def apply_optim_scheduler(optimizer, lr, wd, last_layer_lr):
+ for param_group in optimizer.param_groups:
+ is_last_layer = param_group["is_last_layer"]
+ lr_multiplier = param_group["lr_multiplier"]
+ wd_multiplier = param_group["wd_multiplier"]
+ param_group["weight_decay"] = wd * wd_multiplier
+ param_group["lr"] = (last_layer_lr if is_last_layer else lr) * lr_multiplier
+
+
+def do_test(cfg, model, iteration):
+ new_state_dict = model.teacher.state_dict()
+
+ if distributed.is_main_process():
+ iterstring = str(iteration)
+ eval_dir = os.path.join(cfg.train.output_dir, "eval", iterstring)
+ os.makedirs(eval_dir, exist_ok=True)
+ # save teacher checkpoint
+ teacher_ckp_path = os.path.join(eval_dir, "teacher_checkpoint.pth")
+ torch.save({"teacher": new_state_dict}, teacher_ckp_path)
+
+
+def do_train(cfg, model, resume=False):
+ model.train()
+ inputs_dtype = torch.half
+ fp16_scaler = model.fp16_scaler # for mixed precision training
+
+ # setup optimizer
+
+ optimizer = build_optimizer(cfg, model.get_params_groups())
+ (
+ lr_schedule,
+ wd_schedule,
+ momentum_schedule,
+ teacher_temp_schedule,
+ last_layer_lr_schedule,
+ ) = build_schedulers(cfg)
+
+ # checkpointer
+ checkpointer = FSDPCheckpointer(model, cfg.train.output_dir, optimizer=optimizer, save_to_disk=True)
+
+ start_iter = checkpointer.resume_or_load(cfg.MODEL.WEIGHTS, resume=resume).get("iteration", -1) + 1
+
+ OFFICIAL_EPOCH_LENGTH = cfg.train.OFFICIAL_EPOCH_LENGTH
+ max_iter = cfg.optim.epochs * OFFICIAL_EPOCH_LENGTH
+
+ periodic_checkpointer = PeriodicCheckpointer(
+ checkpointer,
+ period=3 * OFFICIAL_EPOCH_LENGTH,
+ max_iter=max_iter,
+ max_to_keep=3,
+ )
+
+ # setup data preprocessing
+
+ img_size = cfg.crops.global_crops_size
+ patch_size = cfg.student.patch_size
+ n_tokens = (img_size // patch_size) ** 2
+ mask_generator = MaskingGenerator(
+ input_size=(img_size // patch_size, img_size // patch_size),
+ max_num_patches=0.5 * img_size // patch_size * img_size // patch_size,
+ )
+
+ data_transform = DataAugmentationDINO(
+ cfg.crops.global_crops_scale,
+ cfg.crops.local_crops_scale,
+ cfg.crops.local_crops_number,
+ global_crops_size=cfg.crops.global_crops_size,
+ local_crops_size=cfg.crops.local_crops_size,
+ )
+
+ collate_fn = partial(
+ collate_data_and_cast,
+ mask_ratio_tuple=cfg.ibot.mask_ratio_min_max,
+ mask_probability=cfg.ibot.mask_sample_probability,
+ n_tokens=n_tokens,
+ mask_generator=mask_generator,
+ dtype=inputs_dtype,
+ )
+
+ # setup data loader
+
+ dataset = make_dataset(
+ dataset_str=cfg.train.dataset_path,
+ transform=data_transform,
+ target_transform=lambda _: (),
+ )
+ # sampler_type = SamplerType.INFINITE
+ sampler_type = SamplerType.SHARDED_INFINITE
+ data_loader = make_data_loader(
+ dataset=dataset,
+ batch_size=cfg.train.batch_size_per_gpu,
+ num_workers=cfg.train.num_workers,
+ shuffle=True,
+ seed=start_iter, # TODO: Fix this -- cfg.train.seed
+ sampler_type=sampler_type,
+ sampler_advance=0, # TODO(qas): fix this -- start_iter * cfg.train.batch_size_per_gpu,
+ drop_last=True,
+ collate_fn=collate_fn,
+ )
+
+ # training loop
+
+ iteration = start_iter
+
+ logger.info("Starting training from iteration {}".format(start_iter))
+ metrics_file = os.path.join(cfg.train.output_dir, "training_metrics.json")
+ metric_logger = MetricLogger(delimiter=" ", output_file=metrics_file)
+ header = "Training"
+
+ for data in metric_logger.log_every(
+ data_loader,
+ 10,
+ header,
+ max_iter,
+ start_iter,
+ ):
+ current_batch_size = data["collated_global_crops"].shape[0] / 2
+ if iteration > max_iter:
+ return
+
+ # apply schedules
+
+ lr = lr_schedule[iteration]
+ wd = wd_schedule[iteration]
+ mom = momentum_schedule[iteration]
+ teacher_temp = teacher_temp_schedule[iteration]
+ last_layer_lr = last_layer_lr_schedule[iteration]
+ apply_optim_scheduler(optimizer, lr, wd, last_layer_lr)
+
+ # compute losses
+
+ optimizer.zero_grad(set_to_none=True)
+ loss_dict = model.forward_backward(data, teacher_temp=teacher_temp)
+
+ # clip gradients
+
+ if fp16_scaler is not None:
+ if cfg.optim.clip_grad:
+ fp16_scaler.unscale_(optimizer)
+ for v in model.student.values():
+ v.clip_grad_norm_(cfg.optim.clip_grad)
+ fp16_scaler.step(optimizer)
+ fp16_scaler.update()
+ else:
+ if cfg.optim.clip_grad:
+ for v in model.student.values():
+ v.clip_grad_norm_(cfg.optim.clip_grad)
+ optimizer.step()
+
+ # perform teacher EMA update
+
+ model.update_teacher(mom)
+
+ # logging
+
+ if distributed.get_global_size() > 1:
+ for v in loss_dict.values():
+ torch.distributed.all_reduce(v)
+ loss_dict_reduced = {k: v.item() / distributed.get_global_size() for k, v in loss_dict.items()}
+
+ if math.isnan(sum(loss_dict_reduced.values())):
+ logger.info("NaN detected")
+ raise AssertionError
+ losses_reduced = sum(loss for loss in loss_dict_reduced.values())
+
+ metric_logger.update(lr=lr)
+ metric_logger.update(wd=wd)
+ metric_logger.update(mom=mom)
+ metric_logger.update(last_layer_lr=last_layer_lr)
+ metric_logger.update(current_batch_size=current_batch_size)
+ metric_logger.update(total_loss=losses_reduced, **loss_dict_reduced)
+
+ # checkpointing and testing
+
+ if cfg.evaluation.eval_period_iterations > 0 and (iteration + 1) % cfg.evaluation.eval_period_iterations == 0:
+ do_test(cfg, model, f"training_{iteration}")
+ torch.cuda.synchronize()
+ periodic_checkpointer.step(iteration)
+
+ iteration = iteration + 1
+ metric_logger.synchronize_between_processes()
+ return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+def main(args):
+ cfg = setup(args)
+
+ model = SSLMetaArch(cfg).to(torch.device("cuda"))
+ model.prepare_for_distributed_training()
+
+ logger.info("Model:\n{}".format(model))
+ if args.eval_only:
+ iteration = (
+ FSDPCheckpointer(model, save_dir=cfg.train.output_dir)
+ .resume_or_load(cfg.MODEL.WEIGHTS, resume=not args.no_resume)
+ .get("iteration", -1)
+ + 1
+ )
+ return do_test(cfg, model, f"manual_{iteration}")
+
+ do_train(cfg, model, resume=not args.no_resume)
+
+
+if __name__ == "__main__":
+ args = get_args_parser(add_help=True).parse_args()
+ main(args)
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/utils/__init__.py b/torchhub/facebookresearch_dinov2_main/dinov2/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0952fcc3f57e34b3747962e9ebd6fc57aeea63fa
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/utils/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/utils/cluster.py b/torchhub/facebookresearch_dinov2_main/dinov2/utils/cluster.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d98c05d68aa6e9dc165df3db06bd70d999b3fda
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/utils/cluster.py
@@ -0,0 +1,96 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from enum import Enum
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+
+class ClusterType(Enum):
+ AWS = "aws"
+ FAIR = "fair"
+ RSC = "rsc"
+
+
+def _guess_cluster_type() -> ClusterType:
+ uname = os.uname()
+ if uname.sysname == "Linux":
+ if uname.release.endswith("-aws"):
+ # Linux kernel versions on AWS instances are of the form "5.4.0-1051-aws"
+ return ClusterType.AWS
+ elif uname.nodename.startswith("rsc"):
+ # Linux kernel versions on RSC instances are standard ones but hostnames start with "rsc"
+ return ClusterType.RSC
+
+ return ClusterType.FAIR
+
+
+def get_cluster_type(cluster_type: Optional[ClusterType] = None) -> Optional[ClusterType]:
+ if cluster_type is None:
+ return _guess_cluster_type()
+
+ return cluster_type
+
+
+def get_checkpoint_path(cluster_type: Optional[ClusterType] = None) -> Optional[Path]:
+ cluster_type = get_cluster_type(cluster_type)
+ if cluster_type is None:
+ return None
+
+ CHECKPOINT_DIRNAMES = {
+ ClusterType.AWS: "checkpoints",
+ ClusterType.FAIR: "checkpoint",
+ ClusterType.RSC: "checkpoint/dino",
+ }
+ return Path("/") / CHECKPOINT_DIRNAMES[cluster_type]
+
+
+def get_user_checkpoint_path(cluster_type: Optional[ClusterType] = None) -> Optional[Path]:
+ checkpoint_path = get_checkpoint_path(cluster_type)
+ if checkpoint_path is None:
+ return None
+
+ username = os.environ.get("USER")
+ assert username is not None
+ return checkpoint_path / username
+
+
+def get_slurm_partition(cluster_type: Optional[ClusterType] = None) -> Optional[str]:
+ cluster_type = get_cluster_type(cluster_type)
+ if cluster_type is None:
+ return None
+
+ SLURM_PARTITIONS = {
+ ClusterType.AWS: "learnlab",
+ ClusterType.FAIR: "learnlab",
+ ClusterType.RSC: "learn",
+ }
+ return SLURM_PARTITIONS[cluster_type]
+
+
+def get_slurm_executor_parameters(
+ nodes: int, num_gpus_per_node: int, cluster_type: Optional[ClusterType] = None, **kwargs
+) -> Dict[str, Any]:
+ # create default parameters
+ params = {
+ "mem_gb": 0, # Requests all memory on a node, see https://slurm.schedmd.com/sbatch.html
+ "gpus_per_node": num_gpus_per_node,
+ "tasks_per_node": num_gpus_per_node, # one task per GPU
+ "cpus_per_task": 10,
+ "nodes": nodes,
+ "slurm_partition": get_slurm_partition(cluster_type),
+ }
+ # apply cluster-specific adjustments
+ cluster_type = get_cluster_type(cluster_type)
+ if cluster_type == ClusterType.AWS:
+ params["cpus_per_task"] = 12
+ del params["mem_gb"]
+ elif cluster_type == ClusterType.RSC:
+ params["cpus_per_task"] = 12
+ # set additional parameters / apply overrides
+ params.update(kwargs)
+ return params
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/utils/config.py b/torchhub/facebookresearch_dinov2_main/dinov2/utils/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..c3763a8b0808ad45cbbfc1dcb00d52b00113f9ad
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/utils/config.py
@@ -0,0 +1,73 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import logging
+import os
+
+from omegaconf import OmegaConf
+
+import dinov2.distributed as distributed
+from dinov2.logging import setup_logging
+from dinov2.utils import utils
+from dinov2.configs import dinov2_default_config
+
+
+logger = logging.getLogger("dinov2")
+
+
+def apply_scaling_rules_to_cfg(cfg): # to fix
+ if cfg.optim.scaling_rule == "sqrt_wrt_1024":
+ base_lr = cfg.optim.base_lr
+ cfg.optim.lr = base_lr
+ cfg.optim.lr *= math.sqrt(cfg.train.batch_size_per_gpu * distributed.get_global_size() / 1024.0)
+ logger.info(f"sqrt scaling learning rate; base: {base_lr}, new: {cfg.optim.lr}")
+ else:
+ raise NotImplementedError
+ return cfg
+
+
+def write_config(cfg, output_dir, name="config.yaml"):
+ logger.info(OmegaConf.to_yaml(cfg))
+ saved_cfg_path = os.path.join(output_dir, name)
+ with open(saved_cfg_path, "w") as f:
+ OmegaConf.save(config=cfg, f=f)
+ return saved_cfg_path
+
+
+def get_cfg_from_args(args):
+ args.output_dir = os.path.abspath(args.output_dir)
+ args.opts += [f"train.output_dir={args.output_dir}"]
+ default_cfg = OmegaConf.create(dinov2_default_config)
+ cfg = OmegaConf.load(args.config_file)
+ cfg = OmegaConf.merge(default_cfg, cfg, OmegaConf.from_cli(args.opts))
+ return cfg
+
+
+def default_setup(args):
+ distributed.enable(overwrite=True)
+ seed = getattr(args, "seed", 0)
+ rank = distributed.get_global_rank()
+
+ global logger
+ setup_logging(output=args.output_dir, level=logging.INFO)
+ logger = logging.getLogger("dinov2")
+
+ utils.fix_random_seeds(seed + rank)
+ logger.info("git:\n {}\n".format(utils.get_sha()))
+ logger.info("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
+
+
+def setup(args):
+ """
+ Create configs and perform basic setups.
+ """
+ cfg = get_cfg_from_args(args)
+ os.makedirs(args.output_dir, exist_ok=True)
+ default_setup(args)
+ apply_scaling_rules_to_cfg(cfg)
+ write_config(cfg, args.output_dir)
+ return cfg
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/utils/dtype.py b/torchhub/facebookresearch_dinov2_main/dinov2/utils/dtype.py
new file mode 100644
index 0000000000000000000000000000000000000000..cef122b25ff3533e004799a1d977f63eb213fee0
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/utils/dtype.py
@@ -0,0 +1,38 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+from typing import Dict, Union
+
+import numpy as np
+import torch
+
+
+TypeSpec = Union[str, np.dtype, torch.dtype]
+
+
+_NUMPY_TO_TORCH_DTYPE: Dict[np.dtype, torch.dtype] = {
+ np.dtype("bool"): torch.bool,
+ np.dtype("uint8"): torch.uint8,
+ np.dtype("int8"): torch.int8,
+ np.dtype("int16"): torch.int16,
+ np.dtype("int32"): torch.int32,
+ np.dtype("int64"): torch.int64,
+ np.dtype("float16"): torch.float16,
+ np.dtype("float32"): torch.float32,
+ np.dtype("float64"): torch.float64,
+ np.dtype("complex64"): torch.complex64,
+ np.dtype("complex128"): torch.complex128,
+}
+
+
+def as_torch_dtype(dtype: TypeSpec) -> torch.dtype:
+ if isinstance(dtype, torch.dtype):
+ return dtype
+ if isinstance(dtype, str):
+ dtype = np.dtype(dtype)
+ assert isinstance(dtype, np.dtype), f"Expected an instance of nunpy dtype, got {type(dtype)}"
+ return _NUMPY_TO_TORCH_DTYPE[dtype]
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/utils/param_groups.py b/torchhub/facebookresearch_dinov2_main/dinov2/utils/param_groups.py
new file mode 100644
index 0000000000000000000000000000000000000000..d707e70cc11591858d4166410d6ed80621cd49ff
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/utils/param_groups.py
@@ -0,0 +1,94 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections import defaultdict
+import logging
+
+
+logger = logging.getLogger("dinov2")
+
+
+def get_vit_lr_decay_rate(name, lr_decay_rate=1.0, num_layers=12, force_is_backbone=False, chunked_blocks=False):
+ """
+ Calculate lr decay rate for different ViT blocks.
+ Args:
+ name (string): parameter name.
+ lr_decay_rate (float): base lr decay rate.
+ num_layers (int): number of ViT blocks.
+ Returns:
+ lr decay rate for the given parameter.
+ """
+ layer_id = num_layers + 1
+ if name.startswith("backbone") or force_is_backbone:
+ if ".pos_embed" in name or ".patch_embed" in name or ".mask_token" in name or ".cls_token" in name:
+ layer_id = 0
+ elif force_is_backbone and (
+ "pos_embed" in name or "patch_embed" in name or "mask_token" in name or "cls_token" in name
+ ):
+ layer_id = 0
+ elif ".blocks." in name and ".residual." not in name:
+ layer_id = int(name[name.find(".blocks.") :].split(".")[2]) + 1
+ elif chunked_blocks and "blocks." in name and "residual." not in name:
+ layer_id = int(name[name.find("blocks.") :].split(".")[2]) + 1
+ elif "blocks." in name and "residual." not in name:
+ layer_id = int(name[name.find("blocks.") :].split(".")[1]) + 1
+
+ return lr_decay_rate ** (num_layers + 1 - layer_id)
+
+
+def get_params_groups_with_decay(model, lr_decay_rate=1.0, patch_embed_lr_mult=1.0):
+ chunked_blocks = False
+ if hasattr(model, "n_blocks"):
+ logger.info("chunked fsdp")
+ n_blocks = model.n_blocks
+ chunked_blocks = model.chunked_blocks
+ elif hasattr(model, "blocks"):
+ logger.info("first code branch")
+ n_blocks = len(model.blocks)
+ elif hasattr(model, "backbone"):
+ logger.info("second code branch")
+ n_blocks = len(model.backbone.blocks)
+ else:
+ logger.info("else code branch")
+ n_blocks = 0
+ all_param_groups = []
+
+ for name, param in model.named_parameters():
+ name = name.replace("_fsdp_wrapped_module.", "")
+ if not param.requires_grad:
+ continue
+ decay_rate = get_vit_lr_decay_rate(
+ name, lr_decay_rate, num_layers=n_blocks, force_is_backbone=n_blocks > 0, chunked_blocks=chunked_blocks
+ )
+ d = {"params": param, "is_last_layer": False, "lr_multiplier": decay_rate, "wd_multiplier": 1.0, "name": name}
+
+ if "last_layer" in name:
+ d.update({"is_last_layer": True})
+
+ if name.endswith(".bias") or "norm" in name or "gamma" in name:
+ d.update({"wd_multiplier": 0.0})
+
+ if "patch_embed" in name:
+ d.update({"lr_multiplier": d["lr_multiplier"] * patch_embed_lr_mult})
+
+ all_param_groups.append(d)
+ logger.info(f"""{name}: lr_multiplier: {d["lr_multiplier"]}, wd_multiplier: {d["wd_multiplier"]}""")
+
+ return all_param_groups
+
+
+def fuse_params_groups(all_params_groups, keys=("lr_multiplier", "wd_multiplier", "is_last_layer")):
+ fused_params_groups = defaultdict(lambda: {"params": []})
+ for d in all_params_groups:
+ identifier = ""
+ for k in keys:
+ identifier += k + str(d[k]) + "_"
+
+ for k in keys:
+ fused_params_groups[identifier][k] = d[k]
+ fused_params_groups[identifier]["params"].append(d["params"])
+
+ return fused_params_groups.values()
diff --git a/torchhub/facebookresearch_dinov2_main/dinov2/utils/utils.py b/torchhub/facebookresearch_dinov2_main/dinov2/utils/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..53e63eb427f6d5396c8dc153ab07e825c72b68b4
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/dinov2/utils/utils.py
@@ -0,0 +1,96 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import os
+import random
+import subprocess
+from urllib.parse import urlparse
+
+import numpy as np
+import torch
+from torch import nn
+
+
+logger = logging.getLogger("dinov2")
+
+
+def load_pretrained_weights(model, pretrained_weights, checkpoint_key):
+ if urlparse(pretrained_weights).scheme: # If it looks like an URL
+ state_dict = torch.hub.load_state_dict_from_url(pretrained_weights, map_location="cpu")
+ else:
+ state_dict = torch.load(pretrained_weights, map_location="cpu")
+ if checkpoint_key is not None and checkpoint_key in state_dict:
+ logger.info(f"Take key {checkpoint_key} in provided checkpoint dict")
+ state_dict = state_dict[checkpoint_key]
+ # remove `module.` prefix
+ state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+ # remove `backbone.` prefix induced by multicrop wrapper
+ state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
+ msg = model.load_state_dict(state_dict, strict=False)
+ logger.info("Pretrained weights found at {} and loaded with msg: {}".format(pretrained_weights, msg))
+
+
+def fix_random_seeds(seed=31):
+ """
+ Fix random seeds.
+ """
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+ np.random.seed(seed)
+ random.seed(seed)
+
+
+def get_sha():
+ cwd = os.path.dirname(os.path.abspath(__file__))
+
+ def _run(command):
+ return subprocess.check_output(command, cwd=cwd).decode("ascii").strip()
+
+ sha = "N/A"
+ diff = "clean"
+ branch = "N/A"
+ try:
+ sha = _run(["git", "rev-parse", "HEAD"])
+ subprocess.check_output(["git", "diff"], cwd=cwd)
+ diff = _run(["git", "diff-index", "HEAD"])
+ diff = "has uncommitted changes" if diff else "clean"
+ branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])
+ except Exception:
+ pass
+ message = f"sha: {sha}, status: {diff}, branch: {branch}"
+ return message
+
+
+class CosineScheduler(object):
+ def __init__(self, base_value, final_value, total_iters, warmup_iters=0, start_warmup_value=0, freeze_iters=0):
+ super().__init__()
+ self.final_value = final_value
+ self.total_iters = total_iters
+
+ freeze_schedule = np.zeros((freeze_iters))
+
+ warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)
+
+ iters = np.arange(total_iters - warmup_iters - freeze_iters)
+ schedule = final_value + 0.5 * (base_value - final_value) * (1 + np.cos(np.pi * iters / len(iters)))
+ self.schedule = np.concatenate((freeze_schedule, warmup_schedule, schedule))
+
+ assert len(self.schedule) == self.total_iters
+
+ def __getitem__(self, it):
+ if it >= self.total_iters:
+ return self.final_value
+ else:
+ return self.schedule[it]
+
+
+def has_batchnorms(model):
+ bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
+ for name, module in model.named_modules():
+ if isinstance(module, bn_types):
+ return True
+ return False
diff --git a/torchhub/facebookresearch_dinov2_main/hubconf.py b/torchhub/facebookresearch_dinov2_main/hubconf.py
new file mode 100644
index 0000000000000000000000000000000000000000..b36b42cd2136182ea956d8be785cf492418163d8
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/hubconf.py
@@ -0,0 +1,162 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+from enum import Enum
+from typing import Union
+
+import torch
+
+_DINOV2_BASE_URL = "https://dl.fbaipublicfiles.com/dinov2"
+
+
+def _make_dinov2_model_name(arch_name: str, patch_size: int, num_register_tokens: int = 0) -> str:
+ compact_arch_name = arch_name.replace("_", "")[:4]
+ registers_suffix = f"_reg{num_register_tokens}" if num_register_tokens else ""
+ return f"dinov2_{compact_arch_name}{patch_size}{registers_suffix}"
+
+
+class Weights(Enum):
+ LVD142M = "LVD142M"
+
+
+def _make_dinov2_model(
+ *,
+ arch_name: str = "vit_large",
+ img_size: int = 518,
+ patch_size: int = 14,
+ init_values: float = 1.0,
+ ffn_layer: str = "mlp",
+ block_chunks: int = 0,
+ num_register_tokens: int = 0,
+ interpolate_antialias: bool = False,
+ interpolate_offset: float = 0.1,
+ pretrained: bool = True,
+ weights: Union[Weights, str] = Weights.LVD142M,
+ **kwargs,
+):
+ import vision_transformer as vits
+
+ if isinstance(weights, str):
+ try:
+ weights = Weights[weights]
+ except KeyError:
+ raise AssertionError(f"Unsupported weights: {weights}")
+
+ model_base_name = _make_dinov2_model_name(arch_name, patch_size)
+ vit_kwargs = dict(
+ img_size=img_size,
+ patch_size=patch_size,
+ init_values=init_values,
+ ffn_layer=ffn_layer,
+ block_chunks=block_chunks,
+ num_register_tokens=num_register_tokens,
+ interpolate_antialias=interpolate_antialias,
+ interpolate_offset=interpolate_offset,
+ )
+ vit_kwargs.update(**kwargs)
+ model = vits.__dict__[arch_name](**vit_kwargs)
+
+ if pretrained:
+ model_full_name = _make_dinov2_model_name(arch_name, patch_size, num_register_tokens)
+ url = _DINOV2_BASE_URL + f"/{model_base_name}/{model_full_name}_pretrain.pth"
+ state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu")
+ model.load_state_dict(state_dict, strict=True)
+
+ return model
+
+
+def dinov2_vits14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+ """
+ DINOv2 ViT-S/14 model (optionally) pretrained on the LVD-142M dataset.
+ """
+ return _make_dinov2_model(arch_name="vit_small", pretrained=pretrained, weights=weights, **kwargs)
+
+
+def dinov2_vitb14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+ """
+ DINOv2 ViT-B/14 model (optionally) pretrained on the LVD-142M dataset.
+ """
+ return _make_dinov2_model(arch_name="vit_base", pretrained=pretrained, weights=weights, **kwargs)
+
+
+def dinov2_vitl14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+ """
+ DINOv2 ViT-L/14 model (optionally) pretrained on the LVD-142M dataset.
+ """
+ return _make_dinov2_model(arch_name="vit_large", pretrained=pretrained, weights=weights, **kwargs)
+
+
+def dinov2_vitg14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+ """
+ DINOv2 ViT-g/14 model (optionally) pretrained on the LVD-142M dataset.
+ """
+ return _make_dinov2_model(
+ arch_name="vit_giant2",
+ ffn_layer="swiglufused",
+ weights=weights,
+ pretrained=pretrained,
+ **kwargs,
+ )
+
+
+def dinov2_vits14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+ """
+ DINOv2 ViT-S/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+ """
+ return _make_dinov2_model(
+ arch_name="vit_small",
+ pretrained=pretrained,
+ weights=weights,
+ num_register_tokens=4,
+ interpolate_antialias=True,
+ interpolate_offset=0.0,
+ **kwargs,
+ )
+
+
+def dinov2_vitb14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+ """
+ DINOv2 ViT-B/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+ """
+ return _make_dinov2_model(
+ arch_name="vit_base",
+ pretrained=pretrained,
+ weights=weights,
+ num_register_tokens=4,
+ interpolate_antialias=True,
+ interpolate_offset=0.0,
+ **kwargs,
+ )
+
+
+def dinov2_vitl14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+ """
+ DINOv2 ViT-L/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+ """
+ return _make_dinov2_model(
+ arch_name="vit_large",
+ pretrained=pretrained,
+ weights=weights,
+ num_register_tokens=4,
+ interpolate_antialias=True,
+ interpolate_offset=0.0,
+ **kwargs,
+ )
+
+
+def dinov2_vitg14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+ """
+ DINOv2 ViT-g/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+ """
+ return _make_dinov2_model(
+ arch_name="vit_giant2",
+ ffn_layer="swiglufused",
+ weights=weights,
+ pretrained=pretrained,
+ num_register_tokens=4,
+ interpolate_antialias=True,
+ interpolate_offset=0.0,
+ **kwargs,
+ )
diff --git a/torchhub/facebookresearch_dinov2_main/pyproject.toml b/torchhub/facebookresearch_dinov2_main/pyproject.toml
new file mode 100644
index 0000000000000000000000000000000000000000..da67abd8ceabe6d427a96e5d9d4f04b25aebcd32
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/pyproject.toml
@@ -0,0 +1,29 @@
+[tool.black]
+line-length = 120
+
+[tool.pylint.master]
+persistent = false
+score = false
+
+[tool.pylint.messages_control]
+disable = "all"
+enable = [
+ "miscellaneous",
+ "similarities",
+]
+
+[tool.pylint.similarities]
+ignore-comments = true
+ignore-docstrings = true
+ignore-imports = true
+min-similarity-lines = 8
+
+[tool.pylint.reports]
+reports = false
+
+[tool.pylint.miscellaneous]
+notes = [
+ "FIXME",
+ "XXX",
+ "TODO",
+]
diff --git a/torchhub/facebookresearch_dinov2_main/requirements-dev.txt b/torchhub/facebookresearch_dinov2_main/requirements-dev.txt
new file mode 100644
index 0000000000000000000000000000000000000000..5cad34c34cde3a182b616d68b168588827eb9b7c
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/requirements-dev.txt
@@ -0,0 +1,3 @@
+black==22.6.0
+flake8==5.0.4
+pylint==2.15.0
diff --git a/torchhub/facebookresearch_dinov2_main/requirements.txt b/torchhub/facebookresearch_dinov2_main/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..04c159c443b89330ff3c84257c41b011f9791257
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/requirements.txt
@@ -0,0 +1,11 @@
+--extra-index-url https://download.pytorch.org/whl/cu117
+torch==2.0.0
+torchvision==0.15.0
+omegaconf
+torchmetrics==0.10.3
+fvcore
+iopath
+xformers==0.0.18
+submitit
+--extra-index-url https://pypi.nvidia.com
+cuml-cu11
diff --git a/torchhub/facebookresearch_dinov2_main/scripts/lint.sh b/torchhub/facebookresearch_dinov2_main/scripts/lint.sh
new file mode 100644
index 0000000000000000000000000000000000000000..b91acaf762c4be3a0c9d2a162210bfebfaacba08
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/scripts/lint.sh
@@ -0,0 +1,28 @@
+#!/bin/sh
+
+if [ -n "$1" ]; then
+ echo "linting \"$1\""
+fi
+
+echo "running black"
+if [ -n "$1" ]; then
+ black "$1"
+else
+ black dinov2
+fi
+
+echo "running flake8"
+if [ -n "$1" ]; then
+ flake8 "$1"
+else
+ flake8
+fi
+
+echo "running pylint"
+if [ -n "$1" ]; then
+ pylint "$1"
+else
+ pylint dinov2
+fi
+
+exit 0
diff --git a/torchhub/facebookresearch_dinov2_main/setup.cfg b/torchhub/facebookresearch_dinov2_main/setup.cfg
new file mode 100644
index 0000000000000000000000000000000000000000..3cac0c045434cde205eebe91fd5a2c35a1226b4b
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/setup.cfg
@@ -0,0 +1,7 @@
+[flake8]
+max-line-length = 120
+ignore = E203,E501,W503
+per-file-ignores =
+ __init__.py:F401
+exclude =
+ venv
diff --git a/torchhub/facebookresearch_dinov2_main/setup.py b/torchhub/facebookresearch_dinov2_main/setup.py
new file mode 100644
index 0000000000000000000000000000000000000000..001987cfeef6c5fe3469ea09cd4698352fa90939
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/setup.py
@@ -0,0 +1,87 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from pathlib import Path
+import re
+from typing import List, Tuple
+
+from setuptools import setup, find_packages
+
+
+NAME = "dinov2"
+DESCRIPTION = "PyTorch code and models for the DINOv2 self-supervised learning method."
+
+URL = "https://github.com/facebookresearch/dinov2"
+AUTHOR = "FAIR"
+REQUIRES_PYTHON = ">=3.9.0"
+HERE = Path(__file__).parent
+
+
+try:
+ with open(HERE / "README.md", encoding="utf-8") as f:
+ long_description = "\n" + f.read()
+except FileNotFoundError:
+ long_description = DESCRIPTION
+
+
+def get_requirements(path: str = HERE / "requirements.txt") -> Tuple[List[str], List[str]]:
+ requirements = []
+ extra_indices = []
+ with open(path) as f:
+ for line in f.readlines():
+ line = line.rstrip("\r\n")
+ if line.startswith("--extra-index-url "):
+ extra_indices.append(line[18:])
+ continue
+ requirements.append(line)
+ return requirements, extra_indices
+
+
+def get_package_version() -> str:
+ with open(HERE / "dinov2/__init__.py") as f:
+ result = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", f.read(), re.M)
+ if result:
+ return result.group(1)
+ raise RuntimeError("Can't get package version")
+
+
+requirements, extra_indices = get_requirements()
+version = get_package_version()
+dev_requirements, _ = get_requirements(HERE / "requirements-dev.txt")
+
+
+setup(
+ name=NAME,
+ version=version,
+ description=DESCRIPTION,
+ long_description=long_description,
+ long_description_content_type="text/markdown",
+ author=AUTHOR,
+ python_requires=REQUIRES_PYTHON,
+ url=URL,
+ packages=find_packages(),
+ package_data={
+ "": ["*.yaml"],
+ },
+ install_requires=requirements,
+ dependency_links=extra_indices,
+ extras_require={
+ "dev": dev_requirements,
+ },
+ install_package_data=True,
+ license="CC-BY-NC",
+ license_files=("LICENSE",),
+ classifiers=[
+ # Trove classifiers: https://github.com/pypa/trove-classifiers/blob/main/src/trove_classifiers/__init__.py
+ "Development Status :: 3 - Alpha",
+ "Intended Audience :: Developers",
+ "Intended Audience :: Science/Research",
+ "License :: Other/Proprietary License",
+ "Programming Language :: Python :: 3.9",
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
+ "Topic :: Software Development :: Libraries :: Python Modules",
+ ],
+)
diff --git a/torchhub/facebookresearch_dinov2_main/utils.py b/torchhub/facebookresearch_dinov2_main/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..9c6641404093652d5a2f19b4cf283d976ec39e64
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/utils.py
@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+import itertools
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+_DINOV2_BASE_URL = "https://dl.fbaipublicfiles.com/dinov2"
+
+
+def _make_dinov2_model_name(arch_name: str, patch_size: int, num_register_tokens: int = 0) -> str:
+ compact_arch_name = arch_name.replace("_", "")[:4]
+ registers_suffix = f"_reg{num_register_tokens}" if num_register_tokens else ""
+ return f"dinov2_{compact_arch_name}{patch_size}{registers_suffix}"
+
+
+class CenterPadding(nn.Module):
+ def __init__(self, multiple):
+ super().__init__()
+ self.multiple = multiple
+
+ def _get_pad(self, size):
+ new_size = math.ceil(size / self.multiple) * self.multiple
+ pad_size = new_size - size
+ pad_size_left = pad_size // 2
+ pad_size_right = pad_size - pad_size_left
+ return pad_size_left, pad_size_right
+
+ @torch.inference_mode()
+ def forward(self, x):
+ pads = list(itertools.chain.from_iterable(self._get_pad(m) for m in x.shape[:1:-1]))
+ output = F.pad(x, pads)
+ return output
diff --git a/torchhub/facebookresearch_dinov2_main/vision_transformer.py b/torchhub/facebookresearch_dinov2_main/vision_transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..121318f9c77a69a4467888cce44e49549e9954c0
--- /dev/null
+++ b/torchhub/facebookresearch_dinov2_main/vision_transformer.py
@@ -0,0 +1,395 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+from functools import partial
+import math
+import logging
+from typing import Sequence, Tuple, Union, Callable
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from torch.nn.init import trunc_normal_
+
+from dinov2.layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block
+
+
+logger = logging.getLogger("dinov2")
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+ if not depth_first and include_root:
+ fn(module=module, name=name)
+ for child_name, child_module in module.named_children():
+ child_name = ".".join((name, child_name)) if name else child_name
+ named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+ if depth_first and include_root:
+ fn(module=module, name=name)
+ return module
+
+
+class BlockChunk(nn.ModuleList):
+ def forward(self, x):
+ for b in self:
+ x = b(x)
+ return x
+
+
+class DinoVisionTransformer(nn.Module):
+ def __init__(
+ self,
+ img_size=224,
+ patch_size=16,
+ in_chans=3,
+ embed_dim=768,
+ depth=12,
+ num_heads=12,
+ mlp_ratio=4.0,
+ qkv_bias=True,
+ ffn_bias=True,
+ proj_bias=True,
+ drop_path_rate=0.0,
+ drop_path_uniform=False,
+ init_values=None, # for layerscale: None or 0 => no layerscale
+ embed_layer=PatchEmbed,
+ act_layer=nn.GELU,
+ block_fn=Block,
+ ffn_layer="mlp",
+ block_chunks=1,
+ num_register_tokens=0,
+ interpolate_antialias=False,
+ interpolate_offset=0.1,
+ ):
+ """
+ Args:
+ img_size (int, tuple): input image size
+ patch_size (int, tuple): patch size
+ in_chans (int): number of input channels
+ embed_dim (int): embedding dimension
+ depth (int): depth of transformer
+ num_heads (int): number of attention heads
+ mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+ qkv_bias (bool): enable bias for qkv if True
+ proj_bias (bool): enable bias for proj in attn if True
+ ffn_bias (bool): enable bias for ffn if True
+ drop_path_rate (float): stochastic depth rate
+ drop_path_uniform (bool): apply uniform drop rate across blocks
+ weight_init (str): weight init scheme
+ init_values (float): layer-scale init values
+ embed_layer (nn.Module): patch embedding layer
+ act_layer (nn.Module): MLP activation layer
+ block_fn (nn.Module): transformer block class
+ ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+ block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
+ num_register_tokens: (int) number of extra cls tokens (so-called "registers")
+ interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
+ interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
+ """
+ super().__init__()
+ norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+ self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
+ self.num_tokens = 1
+ self.n_blocks = depth
+ self.num_heads = num_heads
+ self.patch_size = patch_size
+ self.num_register_tokens = num_register_tokens
+ self.interpolate_antialias = interpolate_antialias
+ self.interpolate_offset = interpolate_offset
+
+ self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+ num_patches = self.patch_embed.num_patches
+
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+ self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+ assert num_register_tokens >= 0
+ self.register_tokens = (
+ nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
+ )
+
+ if drop_path_uniform is True:
+ dpr = [drop_path_rate] * depth
+ else:
+ dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
+
+ if ffn_layer == "mlp":
+ logger.info("using MLP layer as FFN")
+ ffn_layer = Mlp
+ elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+ logger.info("using SwiGLU layer as FFN")
+ ffn_layer = SwiGLUFFNFused
+ elif ffn_layer == "identity":
+ logger.info("using Identity layer as FFN")
+
+ def f(*args, **kwargs):
+ return nn.Identity()
+
+ ffn_layer = f
+ else:
+ raise NotImplementedError
+
+ blocks_list = [
+ block_fn(
+ dim=embed_dim,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ proj_bias=proj_bias,
+ ffn_bias=ffn_bias,
+ drop_path=dpr[i],
+ norm_layer=norm_layer,
+ act_layer=act_layer,
+ ffn_layer=ffn_layer,
+ init_values=init_values,
+ )
+ for i in range(depth)
+ ]
+ if block_chunks > 0:
+ self.chunked_blocks = True
+ chunked_blocks = []
+ chunksize = depth // block_chunks
+ for i in range(0, depth, chunksize):
+ # this is to keep the block index consistent if we chunk the block list
+ chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
+ self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+ else:
+ self.chunked_blocks = False
+ self.blocks = nn.ModuleList(blocks_list)
+
+ self.norm = norm_layer(embed_dim)
+ self.head = nn.Identity()
+
+ self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+ self.init_weights()
+
+ def init_weights(self):
+ trunc_normal_(self.pos_embed, std=0.02)
+ nn.init.normal_(self.cls_token, std=1e-6)
+ if self.register_tokens is not None:
+ nn.init.normal_(self.register_tokens, std=1e-6)
+ named_apply(init_weights_vit_timm, self)
+
+ def interpolate_pos_encoding(self, x, w, h):
+ previous_dtype = x.dtype
+ npatch = x.shape[1] - 1
+ N = self.pos_embed.shape[1] - 1
+ if npatch == N and w == h:
+ return self.pos_embed
+ pos_embed = self.pos_embed.float()
+ class_pos_embed = pos_embed[:, 0]
+ patch_pos_embed = pos_embed[:, 1:]
+ dim = x.shape[-1]
+ w0 = w // self.patch_size
+ h0 = h // self.patch_size
+ # we add a small number to avoid floating point error in the interpolation
+ # see discussion at https://github.com/facebookresearch/dino/issues/8
+ # DINOv2 with register modify the interpolate_offset from 0.1 to 0.0
+ w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset
+ # w0, h0 = w0 + 0.1, h0 + 0.1
+
+ sqrt_N = math.sqrt(N)
+ sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
+ scale_factor=(sx, sy),
+ # (int(w0), int(h0)), # to solve the upsampling shape issue
+ mode="bicubic",
+ antialias=self.interpolate_antialias
+ )
+
+ assert int(w0) == patch_pos_embed.shape[-2]
+ assert int(h0) == patch_pos_embed.shape[-1]
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+ return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
+
+ def prepare_tokens_with_masks(self, x, masks=None):
+ B, nc, w, h = x.shape
+ x = self.patch_embed(x)
+ if masks is not None:
+ x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+ x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+ x = x + self.interpolate_pos_encoding(x, w, h)
+
+ if self.register_tokens is not None:
+ x = torch.cat(
+ (
+ x[:, :1],
+ self.register_tokens.expand(x.shape[0], -1, -1),
+ x[:, 1:],
+ ),
+ dim=1,
+ )
+
+ return x
+
+ def forward_features_list(self, x_list, masks_list):
+ x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+ for blk in self.blocks:
+ x = blk(x)
+
+ all_x = x
+ output = []
+ for x, masks in zip(all_x, masks_list):
+ x_norm = self.norm(x)
+ output.append(
+ {
+ "x_norm_clstoken": x_norm[:, 0],
+ "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+ "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+ "x_prenorm": x,
+ "masks": masks,
+ }
+ )
+ return output
+
+ def forward_features(self, x, masks=None):
+ if isinstance(x, list):
+ return self.forward_features_list(x, masks)
+
+ x = self.prepare_tokens_with_masks(x, masks)
+
+ for blk in self.blocks:
+ x = blk(x)
+
+ x_norm = self.norm(x)
+ return {
+ "x_norm_clstoken": x_norm[:, 0],
+ "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+ "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+ "x_prenorm": x,
+ "masks": masks,
+ }
+
+ def _get_intermediate_layers_not_chunked(self, x, n=1):
+ x = self.prepare_tokens_with_masks(x)
+ # If n is an int, take the n last blocks. If it's a list, take them
+ output, total_block_len = [], len(self.blocks)
+ blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+ for i, blk in enumerate(self.blocks):
+ x = blk(x)
+ if i in blocks_to_take:
+ output.append(x)
+ assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+ return output
+
+ def _get_intermediate_layers_chunked(self, x, n=1):
+ x = self.prepare_tokens_with_masks(x)
+ output, i, total_block_len = [], 0, len(self.blocks[-1])
+ # If n is an int, take the n last blocks. If it's a list, take them
+ blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+ for block_chunk in self.blocks:
+ for blk in block_chunk[i:]: # Passing the nn.Identity()
+ x = blk(x)
+ if i in blocks_to_take:
+ output.append(x)
+ i += 1
+ assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+ return output
+
+ def get_intermediate_layers(
+ self,
+ x: torch.Tensor,
+ n: Union[int, Sequence] = 1, # Layers or n last layers to take
+ reshape: bool = False,
+ return_class_token: bool = False,
+ norm=True,
+ ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+ if self.chunked_blocks:
+ outputs = self._get_intermediate_layers_chunked(x, n)
+ else:
+ outputs = self._get_intermediate_layers_not_chunked(x, n)
+ if norm:
+ outputs = [self.norm(out) for out in outputs]
+ class_tokens = [out[:, 0] for out in outputs]
+ outputs = [out[:, 1 + self.num_register_tokens:] for out in outputs]
+ if reshape:
+ B, _, w, h = x.shape
+ outputs = [
+ out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
+ for out in outputs
+ ]
+ if return_class_token:
+ return tuple(zip(outputs, class_tokens))
+ return tuple(outputs)
+
+ def forward(self, *args, is_training=False, **kwargs):
+ ret = self.forward_features(*args, **kwargs)
+ if is_training:
+ return ret
+ else:
+ return self.head(ret["x_norm_clstoken"])
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = ""):
+ """ViT weight initialization, original timm impl (for reproducibility)"""
+ if isinstance(module, nn.Linear):
+ trunc_normal_(module.weight, std=0.02)
+ if module.bias is not None:
+ nn.init.zeros_(module.bias)
+
+
+def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=384,
+ depth=12,
+ num_heads=6,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ num_register_tokens=num_register_tokens,
+ **kwargs,
+ )
+ return model
+
+
+def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=768,
+ depth=12,
+ num_heads=12,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ num_register_tokens=num_register_tokens,
+ **kwargs,
+ )
+ return model
+
+
+def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=1024,
+ depth=24,
+ num_heads=16,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ num_register_tokens=num_register_tokens,
+ **kwargs,
+ )
+ return model
+
+
+def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
+ """
+ Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
+ """
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=1536,
+ depth=40,
+ num_heads=24,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ num_register_tokens=num_register_tokens,
+ **kwargs,
+ )
+ return model
diff --git a/ui_prediction.py b/ui_prediction.py
deleted file mode 100644
index 3fe8cda974245f38e5fac66e4eaf744a44ac1a41..0000000000000000000000000000000000000000
--- a/ui_prediction.py
+++ /dev/null
@@ -1,347 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Zhenyu Li
-
-import gradio as gr
-from PIL import Image
-import tempfile
-import torch
-import numpy as np
-
-from zoedepth.utils.arg_utils import parse_unknown
-import argparse
-from zoedepth.models.builder import build_model
-from zoedepth.utils.config import get_config_user
-import matplotlib
-import cv2
-
-from infer_user import regular_tile_param, random_tile_param
-from zoedepth.models.base_models.midas import Resize
-from torchvision.transforms import Compose
-from PIL import Image
-from torchvision import transforms
-import torch.nn.functional as F
-
-from zoedepth.models.base_models.midas import Resize
-from torchvision.transforms import Compose
-
-import gradio as gr
-import numpy as np
-import trimesh
-from zoedepth.utils.geometry import depth_to_points, create_triangles
-from functools import partial
-import tempfile
-
-def depth_edges_mask(depth, occ_filter_thr):
- """Returns a mask of edges in the depth map.
- Args:
- depth: 2D numpy array of shape (H, W) with dtype float32.
- Returns:
- mask: 2D numpy array of shape (H, W) with dtype bool.
- """
- # Compute the x and y gradients of the depth map.
- depth_dx, depth_dy = np.gradient(depth)
- # Compute the gradient magnitude.
- depth_grad = np.sqrt(depth_dx ** 2 + depth_dy ** 2)
- # Compute the edge mask.
- # mask = depth_grad > 0.05 # default in zoedepth
- mask = depth_grad > occ_filter_thr # preserve more edges (?)
- return mask
-
-def load_state_dict(model, state_dict):
- """Load state_dict into model, handling DataParallel and DistributedDataParallel. Also checks for "model" key in state_dict.
-
- DataParallel prefixes state_dict keys with 'module.' when saving.
- If the model is not a DataParallel model but the state_dict is, then prefixes are removed.
- If the model is a DataParallel model but the state_dict is not, then prefixes are added.
- """
- state_dict = state_dict.get('model', state_dict)
- # if model is a DataParallel model, then state_dict keys are prefixed with 'module.'
-
- do_prefix = isinstance(
- model, (torch.nn.DataParallel, torch.nn.parallel.DistributedDataParallel))
- state = {}
- for k, v in state_dict.items():
- if k.startswith('module.') and not do_prefix:
- k = k[7:]
-
- if not k.startswith('module.') and do_prefix:
- k = 'module.' + k
-
- state[k] = v
-
- model.load_state_dict(state, strict=True)
- print("Loaded successfully")
- return model
-
-def load_wts(model, checkpoint_path):
- ckpt = torch.load(checkpoint_path, map_location='cpu')
- return load_state_dict(model, ckpt)
-
-def load_ckpt(model, checkpoint):
- model = load_wts(model, checkpoint)
- print("Loaded weights from {0}".format(checkpoint))
- return model
-
-def colorize(value, cmap='magma_r', vmin=None, vmax=None):
- # normalize
- vmin = value.min() if vmin is None else vmin
- # vmax = value.max() if vmax is None else vmax
- vmax = np.percentile(value, 95) if vmax is None else vmax
-
- if vmin != vmax:
- value = (value - vmin) / (vmax - vmin) # vmin..vmax
- else:
- value = value * 0.
-
- cmapper = matplotlib.cm.get_cmap(cmap)
- value = cmapper(value, bytes=True) # ((1)xhxwx4)
-
- value = value[:, :, :3] # bgr -> rgb
- # rgb_value = value[..., ::-1]
- rgb_value = value
-
- return rgb_value
-
-def predict_depth(model, image, mode, pn, reso, ps, device=None):
-
- pil_image = image
- if device is not None:
- image = transforms.ToTensor()(pil_image).unsqueeze(0).to(device)
- else:
- image = transforms.ToTensor()(pil_image).unsqueeze(0).cuda()
-
- image_height, image_width = image.shape[-2], image.shape[-1]
-
- if reso != '':
- image_resolution = (int(reso.split('x')[0]), int(reso.split('x')[1]))
- else:
- image_resolution = (2160, 3840)
- image_hr = F.interpolate(image, image_resolution, mode='bicubic', align_corners=True)
- preprocess = Compose([Resize(512, 384, keep_aspect_ratio=False, ensure_multiple_of=32, resize_method="minimal")])
- image_lr = preprocess(image)
-
- if ps != '':
- patch_size = (int(ps.split('x')[0]), int(ps.split('x')[1]))
- else:
- patch_size = (int(image_resolution[0] // 4), int(image_resolution[1] // 4))
-
- avg_depth_map = regular_tile_param(
- model,
- image_hr,
- offset_x=0,
- offset_y=0,
- img_lr=image_lr,
- crop_size=patch_size,
- img_resolution=image_resolution,
- transform=preprocess,
- blr_mask=True)
-
- if mode== 'P16':
- pass
- elif mode== 'P49':
- regular_tile_param(
- model,
- image_hr,
- offset_x=patch_size[1]//2,
- offset_y=0,
- img_lr=image_lr,
- iter_pred=avg_depth_map.average_map,
- boundary=0,
- update=True,
- avg_depth_map=avg_depth_map,
- crop_size=patch_size,
- img_resolution=image_resolution,
- transform=preprocess,
- blr_mask=True)
- regular_tile_param(
- model,
- image_hr,
- offset_x=0,
- offset_y=patch_size[0]//2,
- img_lr=image_lr,
- iter_pred=avg_depth_map.average_map,
- boundary=0,
- update=True,
- avg_depth_map=avg_depth_map,
- crop_size=patch_size,
- img_resolution=image_resolution,
- transform=preprocess,
- blr_mask=True)
- regular_tile_param(
- model,
- image_hr,
- offset_x=patch_size[1]//2,
- offset_y=patch_size[0]//2,
- img_lr=image_lr,
- iter_pred=avg_depth_map.average_map,
- boundary=0,
- update=True,
- avg_depth_map=avg_depth_map,
- crop_size=patch_size,
- img_resolution=image_resolution,
- transform=preprocess,
- blr_mask=True)
- elif mode == 'R':
- regular_tile_param(
- model,
- image_hr,
- offset_x=patch_size[1]//2,
- offset_y=0,
- img_lr=image_lr,
- iter_pred=avg_depth_map.average_map,
- boundary=0,
- update=True,
- avg_depth_map=avg_depth_map,
- crop_size=patch_size,
- img_resolution=image_resolution,
- transform=preprocess,
- blr_mask=True)
- regular_tile_param(
- model,
- image_hr,
- offset_x=0,
- offset_y=patch_size[0]//2,
- img_lr=image_lr,
- iter_pred=avg_depth_map.average_map,
- boundary=0,
- update=True,
- avg_depth_map=avg_depth_map,
- crop_size=patch_size,
- img_resolution=image_resolution,
- transform=preprocess,
- blr_mask=True)
- regular_tile_param(
- model,
- image_hr,
- offset_x=patch_size[1]//2,
- offset_y=patch_size[0]//2,
- img_lr=image_lr,
- iter_pred=avg_depth_map.average_map,
- boundary=0,
- update=True,
- avg_depth_map=avg_depth_map,
- crop_size=patch_size,
- img_resolution=image_resolution,
- transform=preprocess,
- blr_mask=True)
-
- for i in range(int(pn)):
- random_tile_param(
- model,
- image_hr,
- img_lr=image_lr,
- iter_pred=avg_depth_map.average_map,
- boundary=0,
- update=True,
- avg_depth_map=avg_depth_map,
- crop_size=patch_size,
- img_resolution=image_resolution,
- transform=preprocess,
- blr_mask=True)
-
- depth = avg_depth_map.average_map.detach().cpu()
- depth = F.interpolate(depth.unsqueeze(dim=0).unsqueeze(dim=0), (image_height, image_width), mode='bicubic', align_corners=True).squeeze().numpy()
-
- return depth
-
-def create_demo(model):
- gr.Markdown("## Depth Prediction Demo")
-
- with gr.Accordion("Advanced options", open=False):
- mode = gr.Radio(["P49", "R"], label="Tiling mode", info="We recommand using P49 for fast evaluation and R with 1024 patches for best visualization results, respectively", elem_id='mode', value='R'),
- patch_number = gr.Slider(1, 1024, label="Please decide the number of random patches (Only useful in mode=R)", step=1, value=256)
- resolution = gr.Textbox(label="Proccessing resolution (Default 4K. Use 'x' to split height and width.)", elem_id='mode', value='2160x3840')
- patch_size = gr.Textbox(label="Patch size (Default 1/4 of image resolution. Use 'x' to split height and width.)", elem_id='mode', value='540x960')
-
- with gr.Row():
- input_image = gr.Image(label="Input Image", type='pil', elem_id='img-display-input')
- depth_image = gr.Image(label="Depth Map", elem_id='img-display-output')
- raw_file = gr.File(label="16-bit raw depth, multiplier:256")
- submit = gr.Button("Submit")
-
- def on_submit(image, mode, pn, reso, ps):
- depth = predict_depth(model, image, mode, pn, reso, ps)
- colored_depth = colorize(depth, cmap='gray_r')
- tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
- raw_depth = Image.fromarray((depth*256).astype('uint16'))
- raw_depth.save(tmp.name)
- return [colored_depth, tmp.name]
-
- submit.click(on_submit, inputs=[input_image, mode[0], patch_number, resolution, patch_size], outputs=[depth_image, raw_file])
- examples = gr.Examples(examples=["examples/example_1.jpeg", "examples/example_2.jpeg", "examples/example_3.jpeg"], inputs=[input_image])
-
-def get_mesh(model, image, mode, pn, reso, ps, keep_edges, occ_filter_thr, fov):
- depth = predict_depth(model, image, mode, pn, reso, ps)
-
- image.thumbnail((1024,1024)) # limit the size of the input image
- depth = F.interpolate(torch.from_numpy(depth).unsqueeze(dim=0).unsqueeze(dim=0), (image.height, image.width), mode='bicubic', align_corners=True).squeeze().numpy()
-
- pts3d = depth_to_points(depth[None], fov=float(fov))
- pts3d = pts3d.reshape(-1, 3)
-
- # Create a trimesh mesh from the points
- # Each pixel is connected to its 4 neighbors
- # colors are the RGB values of the image
-
- verts = pts3d.reshape(-1, 3)
- image = np.array(image)
- if keep_edges:
- triangles = create_triangles(image.shape[0], image.shape[1])
- else:
- triangles = create_triangles(image.shape[0], image.shape[1], mask=~depth_edges_mask(depth, occ_filter_thr=float(occ_filter_thr)))
- colors = image.reshape(-1, 3)
- mesh = trimesh.Trimesh(vertices=verts, faces=triangles, vertex_colors=colors)
-
- # Save as glb
- glb_file = tempfile.NamedTemporaryFile(suffix='.glb', delete=False)
- glb_path = glb_file.name
- mesh.export(glb_path)
- return glb_path
-
-def create_demo_3d(model):
-
- gr.Markdown("### Image to 3D Mesh")
- gr.Markdown("Convert a single 2D image to a 3D mesh")
-
- with gr.Accordion("Advanced options", open=False):
- mode = gr.Radio(["P49", "R"], label="Tiling mode", info="We recommand using P49 for fast evaluation and R with 1024 patches for best visualization results, respectively", elem_id='mode', value='R'),
- patch_number = gr.Slider(1, 1024, label="Please decide the number of random patches (Only useful in mode=R)", step=1, value=256)
- resolution = gr.Textbox(label="Proccessing resolution (Default 4K. Use 'x' to split height and width)", value='2160x3840')
- patch_size = gr.Textbox(label="Patch size (Default 1/4 of image resolution. Use 'x' to split height and width)", value='540x960')
-
- checkbox = gr.Checkbox(label="Keep occlusion edges", value=False)
- # occ_filter_thr = gr.Textbox(label="Occlusion filter threshold", info="Larger value will reserve more edges (Only useful when NOT keeping occlusion edges)", value='0.5')
- # fov = gr.Textbox(label="FOV for inv-projection", value='55')
-
- occ_filter_thr = gr.Slider(0.01, 5, label="Occlusion edge filter threshold", info="Larger value will reserve more occlusion edges (Only useful when NOT keeping occlusion edges)", step=0.01, value=0.2)
- fov = gr.Slider(5, 180, label="FOV for inv-projection", step=1, value=55)
-
-
- with gr.Row():
- input_image = gr.Image(label="Input Image", type='pil')
- result = gr.Model3D(label="3d mesh reconstruction", clear_color=[1.0, 1.0, 1.0, 1.0])
-
- submit = gr.Button("Submit")
- submit.click(partial(get_mesh, model), inputs=[input_image, mode[0], patch_number, resolution, patch_size, checkbox, occ_filter_thr, fov], outputs=[result])
- examples = gr.Examples(examples=["examples/example_1.jpeg", "examples/example_4.jpeg", "examples/example_3.jpeg"], inputs=[input_image])
diff --git a/ui_requirements.txt b/ui_requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..f96f1390dec919edaad836b5bf7cd3b04907b360
--- /dev/null
+++ b/ui_requirements.txt
@@ -0,0 +1,3 @@
+gradio==4.14.0
+gradio_imageslider
+trimesh==3.9.42
\ No newline at end of file
diff --git a/zoedepth/data/data_mono.py b/zoedepth/data/data_mono.py
index f07191ad15de554113a14332849734fede6ffd81..80a8486f239a35331df553f490e213f9bf71e735 100644
--- a/zoedepth/data/data_mono.py
+++ b/zoedepth/data/data_mono.py
@@ -24,8 +24,6 @@
# This file is partly inspired from BTS (https://github.com/cleinc/bts/blob/master/pytorch/bts_dataloader.py); author: Jin Han Lee
-# This file may include modifications from author Zhenyu Li
-
import itertools
import os
import random
@@ -51,13 +49,9 @@ from .ibims import get_ibims_loader
from .sun_rgbd_loader import get_sunrgbd_loader
from .vkitti import get_vkitti_loader
from .vkitti2 import get_vkitti2_loader
-from .u4k import get_u4k_loader
-from .middleburry import get_mid_loader
-from .gta import get_gta_loader
+
from .preprocess import CropParams, get_white_border, get_black_border
-import copy
-from zoedepth.utils.misc import get_boundaries
-from zoedepth.models.base_models.midas import Resize
+
def _is_pil_image(img):
return isinstance(img, Image.Image)
@@ -67,18 +61,12 @@ def _is_numpy_image(img):
return isinstance(img, np.ndarray) and (img.ndim in {2, 3})
-# def preprocessing_transforms(mode, **kwargs):
-# return transforms.Compose([
-# ToTensor(mode=mode, **kwargs)
-# ])
-
-def preprocessing_transforms(mode, sec_stage=False, **kwargs):
+def preprocessing_transforms(mode, **kwargs):
return transforms.Compose([
- ToTensor(mode=mode, sec_stage=sec_stage, **kwargs)
+ ToTensor(mode=mode, **kwargs)
])
-
class DepthDataLoader(object):
def __init__(self, config, mode, device='cpu', transform=None, **kwargs):
"""
@@ -136,26 +124,13 @@ class DepthDataLoader(object):
self.data = get_ddad_loader(config.ddad_root, resize_shape=(
352, 1216), batch_size=1, num_workers=1)
return
-
- # under construction
- if config.dataset == 'u4k':
- self.data = get_u4k_loader(config, mode, transform)
- return
-
- if config.dataset == 'mid':
- self.data = get_mid_loader(config, mode, transform)
- return
-
- if config.dataset == 'gta':
- self.data = get_gta_loader(config, mode, transform)
- return
-
+
img_size = self.config.get("img_size", None)
img_size = img_size if self.config.get(
"do_input_resize", False) else None
+
if transform is None:
- # transform = preprocessing_transforms(mode, size=img_size)
- transform = preprocessing_transforms(mode, size=img_size, sec_stage=config.get("sec_stage", False))
+ transform = preprocessing_transforms(mode, size=img_size)
if mode == 'train':
@@ -187,7 +162,7 @@ class DepthDataLoader(object):
else:
self.eval_sampler = None
self.data = DataLoader(self.testing_samples, 1,
- shuffle=False,
+ shuffle=kwargs.get("shuffle_test", False),
num_workers=1,
pin_memory=False,
sampler=self.eval_sampler)
@@ -302,16 +277,6 @@ class DataLoadPreprocess(Dataset):
with open(config.filenames_file, 'r') as f:
self.filenames = f.readlines()
- self.sec_stage = self.config.get("sec_stage", False)
- # self.crop_size = [120, 160] # 1/4
- self.crop_size = [120*2, 160*2] # 1/4
- self.overlap = self.config.get("overlap", False)
-
- self.consistency_training = self.config.get("consistency_training", False)
- self.overlap_length_h = self.config.get("overlap_length_h", int(60))
- self.overlap_length_w = self.config.get("overlap_length_w", int(80))
- print("current overlap_length_h and overlap_length_w are {} and {}".format(self.overlap_length_h, self.overlap_length_w))
-
self.mode = mode
self.transform = transform
self.to_tensor = ToTensor(mode)
@@ -324,39 +289,11 @@ class DataLoadPreprocess(Dataset):
def postprocess(self, sample):
return sample
- def get_crop_bbox(self, img):
- """Randomly get a crop bounding box."""
- margin_h = max(img.shape[0] - self.crop_size[0], 0)
- margin_w = max(img.shape[1] - self.crop_size[1], 0)
- offset_h = np.random.randint(0, margin_h + 1)
- offset_w = np.random.randint(0, margin_w + 1)
- crop_y1, crop_y2 = offset_h, offset_h + self.crop_size[0]
- crop_x1, crop_x2 = offset_w, offset_w + self.crop_size[1]
-
- return crop_y1, crop_y2, crop_x1, crop_x2
-
- def crop(self, img, crop_bbox, tmp=False):
- """Crop from ``img``"""
-
- crop_y1, crop_y2, crop_x1, crop_x2 = crop_bbox
- if tmp:
- templete = np.zeros((img.shape[0], img.shape[1], 1), dtype=np.float32)
- templete[crop_y1:crop_y2, crop_x1:crop_x2, :] = 1.0
- img = img[crop_y1:crop_y2, crop_x1:crop_x2, ...]
- return img, templete
-
- else:
- img = img[crop_y1:crop_y2, crop_x1:crop_x2, ...]
- return img
-
def __getitem__(self, idx):
sample_path = self.filenames[idx]
focal = float(sample_path.split()[2])
sample = {}
- height=480
- width=640
-
if self.mode == 'train':
if self.config.dataset == 'kitti' and self.config.use_right and random.random() > 0.5:
image_path = os.path.join(
@@ -404,7 +341,6 @@ class DataLoadPreprocess(Dataset):
if self.config.do_random_rotate and (self.config.aug):
- # NOTE: YES!
random_angle = (random.random() - 0.5) * 2 * self.config.degree
image = self.rotate_image(image, random_angle)
depth_gt = self.rotate_image(
@@ -413,91 +349,25 @@ class DataLoadPreprocess(Dataset):
image = np.asarray(image, dtype=np.float32) / 255.0
depth_gt = np.asarray(depth_gt, dtype=np.float32)
depth_gt = np.expand_dims(depth_gt, axis=2)
- disp_gt_copy = depth_gt[:, :, 0].copy()
if self.config.dataset == 'nyu':
depth_gt = depth_gt / 1000.0
else:
depth_gt = depth_gt / 256.0
- # if self.config.aug and (self.config.random_crop):
- # image, depth_gt = self.random_crop(
- # image, depth_gt, self.config.input_height, self.config.input_width)
+ if self.config.aug and (self.config.random_crop):
+ image, depth_gt = self.random_crop(
+ image, depth_gt, self.config.input_height, self.config.input_width)
- image, depth_gt = self.train_preprocess(image, depth_gt)
- img_temp = copy.deepcopy(image)
- depth_gt_temp = copy.deepcopy(depth_gt)
- if self.random_crop: # use in sec_stage
- if self.consistency_training:
- crop_y1, crop_y2, crop_x1, crop_x2 = self.get_crop_bbox(image) # ensure the prob of crop is the same
- while True:
- # shift_x = random.randint(self.overlap_length//3, self.overlap_length)
- # shift_y = random.randint(self.overlap_length//3, self.overlap_length)
- shift_x = self.overlap_length_w
- shift_y = self.overlap_length_h
- if random.random() > 0.5:
- shift_x = shift_x * -1
- if random.random() > 0.5:
- shift_y = shift_y * -1
- crop_y1_shift, crop_y2_shift, crop_x1_shift, crop_x2_shift = crop_y1 + shift_y, crop_y2 + shift_y, crop_x1 + shift_x, crop_x2 + shift_x
- if crop_y1_shift > 0 and crop_x1_shift > 0 and crop_y2_shift < image.shape[0] and crop_x2_shift < image.shape[1]:
- break
- bbox_ori = (crop_y1, crop_y2, crop_x1, crop_x2)
- bbox_shift = (crop_y1_shift, crop_y2_shift, crop_x1_shift, crop_x2_shift)
- image_ori, crop_area_ori = self.crop(image, bbox_ori, tmp=True)
- image_shift, crop_area_shift = self.crop(image, bbox_shift, tmp=True)
- depth_gt_ori = self.crop(depth_gt, bbox_ori)
- depth_gt_shift = self.crop(depth_gt, bbox_shift)
- disp_gt_copy_ori = self.crop(disp_gt_copy, bbox_ori)
- disp_gt_copy_shift = self.crop(disp_gt_copy, bbox_shift)
-
- bboxs_ori = torch.tensor([crop_x1 / width * 160 * 2, crop_y1 / height * 120 * 2, crop_x2 / width * 160 * 2, crop_y2 / height * 120 * 2])
- bboxs_shift = torch.tensor([crop_x1_shift / width * 160 * 2, crop_y1_shift / height * 120 * 2, crop_x2_shift / width * 160 * 2, crop_y2_shift / height * 120 * 2])
- bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2])
- bboxs_raw_shift = torch.tensor([crop_x1_shift, crop_y1_shift, crop_x2_shift, crop_y2_shift])
-
- else:
- bbox = self.get_crop_bbox(image)
- image, crop_area = self.crop(image, bbox, tmp=True)
- depth_gt = self.crop(depth_gt, bbox)
- disp_gt_copy = self.crop(disp_gt_copy, bbox)
-
- crop_y1, crop_y2, crop_x1, crop_x2 = bbox
- bboxs_res = torch.tensor([crop_x1 / width * 160 * 2, crop_y1 / height * 120 * 2, crop_x2 / width * 160 * 2, crop_y2 / height * 120 * 2]) # coord in 384, 512
- bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2])
+ if self.config.aug and self.config.random_translate:
+ # print("Random Translation!")
+ image, depth_gt = self.random_translate(image, depth_gt, self.config.max_translation)
+ image, depth_gt = self.train_preprocess(image, depth_gt)
mask = np.logical_and(depth_gt > self.config.min_depth,
depth_gt < self.config.max_depth).squeeze()[None, ...]
- mask_raw = np.logical_and(depth_gt_temp > self.config.min_depth, depth_gt_temp < self.config.max_depth).squeeze()[None, ...]
- sample = {'image': image, 'depth': depth_gt, 'focal': focal, 'mask': mask, 'image_raw': image.copy(), 'mask_raw': mask_raw}
-
- if self.random_crop:
- if self.consistency_training:
- image = np.concatenate([image_ori, image_shift], axis=-1)
- depth_gt = np.concatenate([depth_gt_ori, depth_gt_shift], axis=-1)
- crop_area = np.concatenate([crop_area_ori, crop_area_shift], axis=-1)
- bboxs_res = torch.cat([bboxs_ori, bboxs_shift], dim=-1)
- bboxes_raw_res = torch.cat([bboxs_raw, bboxs_raw_shift], dim=-1)
- mask = np.logical_and(depth_gt > self.config.min_depth,
- depth_gt < self.config.max_depth)
-
- # hack the sample dict
- sample['image'] = image
- sample['depth'] = depth_gt
- sample['crop_area'] = crop_area
- sample['bbox'] = bboxs_res
- sample['bbox_raw'] = bboxes_raw_res
- sample['shift'] = torch.tensor([shift_y, shift_x]) # h direction, then w direction
- sample['mask'] = mask
-
- else:
- if bboxs_res is not None:
- sample['bbox'] = bboxs_res
- sample['bbox_raw'] = bboxs_raw
- sample['crop_area'] = crop_area
-
- if self.config.aug and self.config.random_translate:
- image, depth_gt = self.random_translate(image, depth_gt, self.config.max_translation)
+ sample = {'image': image, 'depth': depth_gt, 'focal': focal,
+ 'mask': mask, **sample}
else:
if self.mode == 'online_eval':
@@ -525,7 +395,6 @@ class DataLoadPreprocess(Dataset):
if has_valid_depth:
depth_gt = np.asarray(depth_gt, dtype=np.float32)
depth_gt = np.expand_dims(depth_gt, axis=2)
- disp_gt_copy = depth_gt[:, :, 0].copy()
if self.config.dataset == 'nyu':
depth_gt = depth_gt / 1000.0
else:
@@ -547,46 +416,10 @@ class DataLoadPreprocess(Dataset):
depth_gt = depth_gt[top_margin:top_margin +
352, left_margin:left_margin + 1216, :]
- # NOTE: start insert something new for sec_stage training
- if self.sec_stage:
- img_temp = copy.deepcopy(image)
- depth_gt_temp = copy.deepcopy(depth_gt)
-
- x_start, y_start = [0, 240], [0, 320]
- # x_start, y_start = [0 + 3 * self.overlap / 2, 120 + self.overlap / 2, 240 - self.overlap / 2, 360 - 3 * self.overlap / 2], \
- # [0 + 3 * self.overlap / 2, 160 + self.overlap / 2, 320 - self.overlap / 2, 480 - 3 * self.overlap / 2]
- img_crops = []
- bboxs_roi = []
- crop_areas = []
- bboxs_raw_list = []
- for x in x_start:
- for y in y_start:
- bbox = (int(x), int(x+240), int(y), int(y+320))
- img_crop, crop_area = self.crop(image, bbox, tmp=True)
- img_crops.append(img_crop)
- crop_areas.append(crop_area)
- crop_y1, crop_y2, crop_x1, crop_x2 = bbox
- bbox_roi = torch.tensor([crop_x1 / width * 160 * 2, crop_y1 / height * 120 * 2, crop_x2 / width * 160 * 2, crop_y2 / height * 120 * 2])
- bboxs_roi.append(bbox_roi)
- bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2])
- bboxs_raw_list.append(bboxs_raw)
-
- image = img_crops
- bboxs_roi = torch.stack(bboxs_roi, dim=0)
- bboxs_raw = torch.stack(bboxs_raw_list, dim=0)
- disp_gt_edges = get_boundaries(disp_gt_copy, th=1, dilation=0)
-
-
if self.mode == 'online_eval':
sample = {'image': image, 'depth': depth_gt, 'focal': focal, 'has_valid_depth': has_valid_depth,
'image_path': sample_path.split()[0], 'depth_path': sample_path.split()[1],
- 'mask': mask, 'image_raw': image.copy(), 'disp_gt_edges': disp_gt_edges}
- if bboxs_roi is not None:
- sample['bbox'] = bboxs_roi
- sample['bbox_raw'] = bboxs_raw
- if crop_areas is not None:
- sample['crop_area'] = crop_areas
-
+ 'mask': mask}
else:
sample = {'image': image, 'focal': focal}
@@ -596,9 +429,6 @@ class DataLoadPreprocess(Dataset):
sample['mask'] = mask
if self.transform:
- # sample = self.transform(sample)
- sample['img_temp'] = img_temp
- sample['depth_gt_temp'] = depth_gt_temp
sample = self.transform(sample)
sample = self.postprocess(sample)
@@ -680,177 +510,64 @@ class DataLoadPreprocess(Dataset):
return len(self.filenames)
-# class ToTensor(object):
-# def __init__(self, mode, do_normalize=False, size=None):
-# self.mode = mode
-# self.normalize = transforms.Normalize(
-# mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if do_normalize else nn.Identity()
-# self.size = size
-# if size is not None:
-# self.resize = transforms.Resize(size=size)
-# else:
-# self.resize = nn.Identity()
-
-# def __call__(self, sample):
-# image, focal = sample['image'], sample['focal']
-# image = self.to_tensor(image)
-# image = self.normalize(image)
-# image = self.resize(image)
-
-# if self.mode == 'test':
-# return {'image': image, 'focal': focal}
-
-# depth = sample['depth']
-# if self.mode == 'train':
-# depth = self.to_tensor(depth)
-# return {**sample, 'image': image, 'depth': depth, 'focal': focal}
-# else:
-# has_valid_depth = sample['has_valid_depth']
-# image = self.resize(image)
-# return {**sample, 'image': image, 'depth': depth, 'focal': focal, 'has_valid_depth': has_valid_depth,
-# 'image_path': sample['image_path'], 'depth_path': sample['depth_path']}
-
-# def to_tensor(self, pic):
-# if not (_is_pil_image(pic) or _is_numpy_image(pic)):
-# raise TypeError(
-# 'pic should be PIL Image or ndarray. Got {}'.format(type(pic)))
-
-# if isinstance(pic, np.ndarray):
-# img = torch.from_numpy(pic.transpose((2, 0, 1)))
-# return img
-
-# # handle PIL Image
-# if pic.mode == 'I':
-# img = torch.from_numpy(np.array(pic, np.int32, copy=False))
-# elif pic.mode == 'I;16':
-# img = torch.from_numpy(np.array(pic, np.int16, copy=False))
-# else:
-# img = torch.ByteTensor(
-# torch.ByteStorage.from_buffer(pic.tobytes()))
-# # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK
-# if pic.mode == 'YCbCr':
-# nchannel = 3
-# elif pic.mode == 'I;16':
-# nchannel = 1
-# else:
-# nchannel = len(pic.mode)
-# img = img.view(pic.size[1], pic.size[0], nchannel)
-
-# img = img.transpose(0, 1).transpose(0, 2).contiguous()
-# if isinstance(img, torch.ByteTensor):
-# return img.float()
-# else:
-# return img
-
-
class ToTensor(object):
- def __init__(self, mode, do_normalize=False, size=None, sec_stage=False):
+ def __init__(self, mode, do_normalize=False, size=None):
self.mode = mode
- # don't do normalization as default
self.normalize = transforms.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if do_normalize else nn.Identity()
self.size = size
if size is not None:
- # self.resize = transforms.Resize(size=size)
- net_h, net_w = size
- self.resize = Resize(net_w, net_h, keep_aspect_ratio=False, ensure_multiple_of=32, resize_method="minimal")
+ self.resize = transforms.Resize(size=size)
else:
self.resize = nn.Identity()
- self.sec_stage = sec_stage
def __call__(self, sample):
image, focal = sample['image'], sample['focal']
- crop_areas = sample.get('crop_area', None)
-
- if isinstance(image, list):
- # there must be crop_areas
- # only infer on eval sec_stage
- imgs_process = []
- crp_process = []
- for img, crp in zip(image, crop_areas):
- img = self.to_tensor(img)
- img = self.normalize(img)
- img = img.unsqueeze(dim=0)
- img = self.resize(img)
- img = img.squeeze(dim=0)
- imgs_process.append(img)
-
- crp = self.to_tensor(crp)
- crp = crp.unsqueeze(dim=0)
- crp = self.resize(crp)
- crp = crp.squeeze(dim=0)
- crp_process.append(crp)
-
- image = torch.cat(imgs_process, dim=0)
- crop_areas = torch.cat(crp_process, dim=0)
-
- img_temp = sample['img_temp']
- img_temp = self.to_tensor(img_temp)
- img_temp = self.normalize(img_temp)
- img_temp = img_temp.unsqueeze(dim=0)
- img_temp = self.resize(img_temp) #NOTE: hack
- img_temp = img_temp.squeeze(dim=0)
- image_raw = copy.deepcopy(img_temp)
-
-
- else:
- image = self.to_tensor(image)
- image = self.normalize(image)
-
- if crop_areas is not None:
- crop_areas = self.to_tensor(crop_areas)
- crop_areas = crop_areas.unsqueeze(dim=0)
- crop_areas = self.resize(crop_areas)
- crop_areas = crop_areas.squeeze(dim=0)
-
- if self.sec_stage:
- img_temp = sample['img_temp']
- img_temp = self.to_tensor(img_temp)
- img_temp = self.normalize(img_temp)
-
- img_temp = img_temp.unsqueeze(dim=0)
- img_temp = self.resize(img_temp)
- image_raw = img_temp.squeeze(dim=0)
-
- image = image.unsqueeze(dim=0)
- image = self.resize(image)
- image = image.squeeze(dim=0)
-
- else:
- # in the first stage, this hr info is reserved
- image_raw = copy.deepcopy(image)
- image = image.unsqueeze(dim=0)
- image = self.resize(image)
- image = image.squeeze(dim=0)
+ image = self.to_tensor(image)
+ image = self.normalize(image)
+ image = self.resize(image)
if self.mode == 'test':
- return_dict = {'image': image, 'focal': focal}
- if crop_areas is not None:
- return_dict['crop_area'] = crop_areas
- return return_dict
-
+ return {'image': image, 'focal': focal}
depth = sample['depth']
- depth = self.to_tensor(depth)
- depth_gt_temp = sample['depth_gt_temp']
- depth_gt_raw = self.to_tensor(depth_gt_temp)
-
if self.mode == 'train':
- return_dict = {**sample, 'image': image, 'depth': depth, 'focal': focal, 'image_raw': image_raw, 'depth_raw': depth_gt_raw}
- if crop_areas is not None:
- return_dict['crop_area'] = crop_areas
- return return_dict
+ depth = self.to_tensor(depth)
+ return {**sample, 'image': image, 'depth': depth, 'focal': focal}
else:
has_valid_depth = sample['has_valid_depth']
- # image = self.resize(image)
- return_dict = {**sample, 'image': image, 'depth': depth, 'focal': focal, 'image_raw': image_raw,
- 'has_valid_depth': has_valid_depth, 'image_path': sample['image_path'], 'depth_path': sample['depth_path'],
- 'depth_raw': depth_gt_raw}
- if crop_areas is not None:
- return_dict['crop_area'] = crop_areas
- return return_dict
+ image = self.resize(image)
+ return {**sample, 'image': image, 'depth': depth, 'focal': focal, 'has_valid_depth': has_valid_depth,
+ 'image_path': sample['image_path'], 'depth_path': sample['depth_path']}
def to_tensor(self, pic):
+ if not (_is_pil_image(pic) or _is_numpy_image(pic)):
+ raise TypeError(
+ 'pic should be PIL Image or ndarray. Got {}'.format(type(pic)))
+
if isinstance(pic, np.ndarray):
- img = torch.from_numpy(pic.transpose((2, 0, 1))) # img here
+ img = torch.from_numpy(pic.transpose((2, 0, 1)))
+ return img
+
+ # handle PIL Image
+ if pic.mode == 'I':
+ img = torch.from_numpy(np.array(pic, np.int32, copy=False))
+ elif pic.mode == 'I;16':
+ img = torch.from_numpy(np.array(pic, np.int16, copy=False))
+ else:
+ img = torch.ByteTensor(
+ torch.ByteStorage.from_buffer(pic.tobytes()))
+ # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK
+ if pic.mode == 'YCbCr':
+ nchannel = 3
+ elif pic.mode == 'I;16':
+ nchannel = 1
+ else:
+ nchannel = len(pic.mode)
+ img = img.view(pic.size[1], pic.size[0], nchannel)
+
+ img = img.transpose(0, 1).transpose(0, 2).contiguous()
+ if isinstance(img, torch.ByteTensor):
+ return img.float()
+ else:
return img
diff --git a/zoedepth/data/gta.py b/zoedepth/data/gta.py
deleted file mode 100644
index 8f7eb133e158546a01eeb188c8ce2fc77984a915..0000000000000000000000000000000000000000
--- a/zoedepth/data/gta.py
+++ /dev/null
@@ -1,320 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Zhenyu Li
-
-# This file is partly inspired from ZoeDepth (https://github.com/isl-org/ZoeDepth/blob/main/zoedepth/data/data_mono.py); author: Shariq Farooq Bhat
-
-import os
-import numpy as np
-import torch
-from PIL import Image
-from torch.utils.data import DataLoader, Dataset
-from torchvision import transforms
-import os.path as osp
-import random
-import torch.nn as nn
-import cv2
-import copy
-from zoedepth.utils.misc import get_boundaries
-from zoedepth.models.base_models.midas import Resize
-
-
-from .u4k import U4KDataset, remove_leading_slash
-
-import re
-import numpy as np
-import sys
-import matplotlib.pyplot as plt
-import imageio
-
-class GTA(U4KDataset):
- def __init__(self, config, mode, data_root, split):
- super().__init__(config, mode, data_root, split)
- self.crop_size = [270, 480] # 1/4
-
- def load_data_list(self):
- # os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
-
- """Load annotation from directory.
- Args:
- data_root (str): Data root for img_dir/ann_dir.
- split (str|None): Split txt file. If split is specified, only file
- with suffix in the splits will be loaded. Otherwise, all images
- in img_dir/ann_dir will be loaded. Default: None
- Returns:
- list[dict]: All image info of dataset.
- """
- data_root = self.data_root
- split = self.split
-
- self.invalid_depth_num = 0
- img_infos = []
- if split is not None:
- with open(split) as f:
- for line in f:
- img_info_l = dict()
-
- img_l, depth_map_l = line.strip().split(" ")
-
- img_info_l['depth_map_path'] = osp.join(data_root, remove_leading_slash(depth_map_l))
- img_info_l['img_path'] = osp.join(data_root, remove_leading_slash(img_l))
- img_info_l['depth_fields'] = []
-
- img_infos.append(img_info_l)
- else:
- raise NotImplementedError
-
- # github issue:: make sure the same order
- img_infos = sorted(img_infos, key=lambda x: x['img_path'])
- return img_infos
-
-
- def __getitem__(self, idx):
-
- img_file_path = self.data_infos[idx]['img_path']
- disp_path = self.data_infos[idx]['depth_map_path']
-
- height=1080
- width=1920
-
- image = Image.open(img_file_path).convert("RGB")
- image = np.asarray(image, dtype=np.uint8) / 255.0
- image = image.astype(np.float32)
-
- # disp_gt, scale = readPFM(disp_path)
-
- # ref_depth_full = cv2.imread(disp_path, cv2.IMREAD_UNCHANGED)
- ref_depth_full = imageio.imread(disp_path)
- ref_depth_full = np.array(ref_depth_full).astype(np.float32) / 256
- # invalid_mask_full = np.logical_or(ref_depth_full < self.config.min_depth, ref_depth_full > self.config.max_depth)
- # ref_depth_full[invalid_mask_full] = 0
- # ref_depth_full[ref_depth_full < self.config.min_depth] = self.config.min_depth
- # ref_depth_full[ref_depth_full > self.config.max_depth] = self.config.max_depth
-
- depth_gt = ref_depth_full
- disp_gt_copy = depth_gt
- depth_gt = depth_gt[..., np.newaxis]
- depth_gt[depth_gt > self.config.max_depth] = self.config.max_depth # for vis
-
- # disp_gt_edges = get_boundaries(disp_gt_copy, th=1/20, dilation=0)
-
- bbox = None
- bboxs_res = None
- crop_areas = None
- bboxs_roi = None # hack for infer
-
- if self.mode == 'train':
- image, depth_gt = self.train_preprocess(image, depth_gt)
- img_temp = copy.deepcopy(image)
- depth_gt_temp = copy.deepcopy(depth_gt)
-
- if self.random_crop: # use in sec_stage
- if self.consistency_training:
- crop_y1, crop_y2, crop_x1, crop_x2 = self.get_crop_bbox(image) # ensure the prob of crop is the same
- while True:
- # shift_x = random.randint(self.overlap_length//3, self.overlap_length)
- # shift_y = random.randint(self.overlap_length//3, self.overlap_length)
- shift_x = self.overlap_length_w
- shift_y = self.overlap_length_h
- if random.random() > 0.5:
- shift_x = shift_x * -1
- if random.random() > 0.5:
- shift_y = shift_y * -1
- crop_y1_shift, crop_y2_shift, crop_x1_shift, crop_x2_shift = crop_y1 + shift_y, crop_y2 + shift_y, crop_x1 + shift_x, crop_x2 + shift_x
- if crop_y1_shift > 0 and crop_x1_shift > 0 and crop_y2_shift < image.shape[0] and crop_x2_shift < image.shape[1]:
- break
- bbox_ori = (crop_y1, crop_y2, crop_x1, crop_x2)
- bbox_shift = (crop_y1_shift, crop_y2_shift, crop_x1_shift, crop_x2_shift)
- image_ori, crop_area_ori = self.crop(image, bbox_ori, tmp=True)
- image_shift, crop_area_shift = self.crop(image, bbox_shift, tmp=True)
- depth_gt_ori = self.crop(depth_gt, bbox_ori)
- depth_gt_shift = self.crop(depth_gt, bbox_shift)
- disp_gt_copy_ori = self.crop(disp_gt_copy, bbox_ori)
- disp_gt_copy_shift = self.crop(disp_gt_copy, bbox_shift)
-
- bboxs_ori = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384])
- bboxs_shift = torch.tensor([crop_x1_shift / width * 512, crop_y1_shift / height * 384, crop_x2_shift / width * 512, crop_y2_shift / height * 384])
- bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2])
- bboxs_raw_shift = torch.tensor([crop_x1_shift, crop_y1_shift, crop_x2_shift, crop_y2_shift])
-
- else:
- bbox = self.get_crop_bbox(image)
- image, crop_area = self.crop(image, bbox, tmp=True)
- depth_gt = self.crop(depth_gt, bbox)
- disp_gt_copy = self.crop(disp_gt_copy, bbox)
-
- crop_y1, crop_y2, crop_x1, crop_x2 = bbox
- bboxs_res = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384]) # coord in 384, 512
- bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2])
-
-
- mask = np.logical_and(depth_gt > self.config.min_depth,
- depth_gt < self.config.max_depth).squeeze()[None, ...]
- mask_raw = np.logical_and(depth_gt_temp > self.config.min_depth, depth_gt_temp < self.config.max_depth).squeeze()[None, ...]
- sample = {'image': image, 'depth': depth_gt, 'focal': 0, 'mask': mask, 'image_raw': image.copy(), 'mask_raw': mask_raw}
-
-
- if self.random_crop:
- if self.consistency_training:
- image = np.concatenate([image_ori, image_shift], axis=-1)
- depth_gt = np.concatenate([depth_gt_ori, depth_gt_shift], axis=-1)
- crop_area = np.concatenate([crop_area_ori, crop_area_shift], axis=-1)
- bboxs_res = torch.cat([bboxs_ori, bboxs_shift], dim=-1)
- bboxes_raw_res = torch.cat([bboxs_raw, bboxs_raw_shift], dim=-1)
- mask = np.logical_and(depth_gt > self.config.min_depth,
- depth_gt < self.config.max_depth)
-
- # hack the sample dict
- sample['image'] = image
- sample['depth'] = depth_gt
- sample['crop_area'] = crop_area
- sample['bbox'] = bboxs_res
- sample['bbox_raw'] = bboxes_raw_res
- sample['shift'] = torch.tensor([shift_y, shift_x]) # h direction, then w direction
- sample['mask'] = mask
-
- else:
- if bboxs_res is not None:
- sample['bbox'] = bboxs_res
- sample['bbox_raw'] = bboxs_raw
- sample['crop_area'] = crop_area
-
- if self.sampled_training:
- self.data_sampler(sample, disp_gt_copy)
-
- # update mask
- sample_points = sample['sample_points']
- sample_mask = np.logical_and(sample_points[:, -1] > self.config.min_depth,
- sample_points[:, -1] < self.config.max_depth).squeeze()[None, ...]
- sample['sample_mask'] = sample_mask
-
-
- else:
- # nothing needs to be changed for consistency training.
-
- img_temp = copy.deepcopy(image)
- depth_gt_temp = copy.deepcopy(depth_gt)
- if self.sec_stage:
- # x_start, y_start = [0, 540, 1080, 1620], [0, 960, 1920, 2880]
- x_start, y_start = [0 + 3 * self.overlap / 2, 270 + self.overlap / 2, 540 - self.overlap / 2, 810 - 3 * self.overlap / 2], \
- [0 + 3 * self.overlap / 2, 480 + self.overlap / 2, 960 - self.overlap / 2, 1440 - 3 * self.overlap / 2]
- img_crops = []
- bboxs_roi = []
- crop_areas = []
- bboxs_raw_list = []
- for x in x_start:
- for y in y_start:
- bbox = (int(x), int(x+270), int(y), int(y+480))
- img_crop, crop_area = self.crop(image, bbox, tmp=True)
- img_crops.append(img_crop)
- crop_areas.append(crop_area)
- crop_y1, crop_y2, crop_x1, crop_x2 = bbox
- bbox_roi = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384])
- bboxs_roi.append(bbox_roi)
- bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2])
- bboxs_raw_list.append(bboxs_raw)
-
- image = img_crops
- bboxs_roi = torch.stack(bboxs_roi, dim=0)
- bboxs_raw = torch.stack(bboxs_raw_list, dim=0)
-
- # bbox = (820, 1360 ,1440, 2400) # a hack version for quick evaluation
- # image = self.crop(image, bbox)
- # depth_gt = self.crop(depth_gt, bbox)
- # disp_gt_copy = self.crop(disp_gt_copy, bbox)
-
- mask = np.logical_and(depth_gt > self.config.min_depth,
- depth_gt < self.config.max_depth).squeeze()[None, ...]
-
- disp_gt_edges = get_boundaries(disp_gt_copy, th=1, dilation=0)
- if self.mode == 'online_eval':
- sample = {'image': image, 'depth': depth_gt, 'focal': 0, 'has_valid_depth': True,
- 'image_path': img_file_path, 'depth_path': disp_path, 'depth_factor_path': 0,
- 'mask': mask, 'image_raw': image.copy(), 'disp_gt_edges': disp_gt_edges}
- if bboxs_roi is not None:
- sample['bbox'] = bboxs_roi
- sample['bbox_raw'] = bboxs_raw
- if crop_areas is not None:
- sample['crop_area'] = crop_areas
-
-
- else:
- sample = {'image': image, 'focal': 0, 'image_raw': image.copy(), 'disp_gt_edges': disp_gt_edges, 'image_path': img_file_path}
- if bboxs_roi is not None:
- sample['bbox'] = bboxs_roi
- sample['bbox_raw'] = bboxs_raw
- if crop_areas is not None:
- sample['crop_area'] = crop_areas
-
- if self.transform:
- sample['img_temp'] = img_temp
- sample['depth_gt_temp'] = depth_gt_temp
- sample = self.transform(sample)
-
- sample['dataset'] = self.config.dataset
- return sample
-
- def __len__(self):
- return len(self.data_infos)
-
-def get_gta_loader(config, mode, transform):
- if mode == 'train':
- log = 0
- dataset = GTA(config, mode, config.data_path, config.filenames_train)
- dataset[0]
- if config.distributed:
- train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
- else:
- train_sampler = None
-
- dataloader = DataLoader(dataset,
- batch_size=config.batch_size,
- shuffle=(train_sampler is None),
- num_workers=config.workers,
- pin_memory=True,
- persistent_workers=True,
- sampler=train_sampler)
-
- elif mode == 'online_eval':
- dataset = GTA(config, mode, config.data_path, config.filenames_val)
- # dataset = U4KDataset(config, mode, config.data_path, config.filenames_train)
-
-
- if config.distributed: # redundant. here only for readability and to be more explicit
- # Give whole test set to all processes (and report evaluation only on one) regardless
- eval_sampler = None
- else:
- eval_sampler = None
-
- dataloader = DataLoader(dataset, 1,
- shuffle=False,
- num_workers=1,
- pin_memory=False,
- sampler=eval_sampler)
-
- else:
- dataset = GTA(config, mode, config.data_path, config.filenames_test)
- dataloader = DataLoader(dataset, 1, shuffle=False, num_workers=1)
-
- return dataloader
-
diff --git a/zoedepth/data/middleburry.py b/zoedepth/data/middleburry.py
deleted file mode 100644
index cf3b82fabb09d986d9901dbd456c8f6cb671ffda..0000000000000000000000000000000000000000
--- a/zoedepth/data/middleburry.py
+++ /dev/null
@@ -1,373 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Zhenyu Li
-
-# This file is partly inspired from ZoeDepth (https://github.com/isl-org/ZoeDepth/blob/main/zoedepth/data/data_mono.py); author: Shariq Farooq Bhat
-
-import os
-import numpy as np
-import torch
-from PIL import Image
-from torch.utils.data import DataLoader, Dataset
-from torchvision import transforms
-import os.path as osp
-import random
-import torch.nn as nn
-import cv2
-import copy
-from zoedepth.utils.misc import get_boundaries
-from zoedepth.models.base_models.midas import Resize
-
-
-from .u4k import U4KDataset, remove_leading_slash
-
-import re
-import numpy as np
-import sys
-import matplotlib.pyplot as plt
-
-def readPFM(file):
- file = open(file, 'rb')
-
- color = None
- width = None
- height = None
- scale = None
- endian = None
-
- header = file.readline().rstrip()
- if (sys.version[0]) == '3':
- header = header.decode('utf-8')
- if header == 'PF':
- color = True
- elif header == 'Pf':
- color = False
- else:
- raise Exception('Not a PFM file.')
-
- if (sys.version[0]) == '3':
- dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode('utf-8'))
- else:
- dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline())
- if dim_match:
- width, height = map(int, dim_match.groups())
- else:
- raise Exception('Malformed PFM header.')
-
- if (sys.version[0]) == '3':
- scale = float(file.readline().rstrip().decode('utf-8'))
- else:
- scale = float(file.readline().rstrip())
-
- if scale < 0: # little-endian
- endian = '<'
- scale = -scale
- else:
- endian = '>' # big-endian
-
- data = np.fromfile(file, endian + 'f')
- shape = (height, width, 3) if color else (height, width)
-
- data = np.reshape(data, shape)
- data = np.flipud(data)
- return data, scale
-
-class MiddleBurry(U4KDataset):
- def load_data_list(self):
- """Load annotation from directory.
- Args:
- data_root (str): Data root for img_dir/ann_dir.
- split (str|None): Split txt file. If split is specified, only file
- with suffix in the splits will be loaded. Otherwise, all images
- in img_dir/ann_dir will be loaded. Default: None
- Returns:
- list[dict]: All image info of dataset.
- """
- data_root = self.data_root
- split = self.split
-
- self.invalid_depth_num = 0
- img_infos = []
- if split is not None:
- with open(split) as f:
- for line in f:
- img_info_l = dict()
-
- img_l, depth_map_l = line.strip().split(" ")
-
- img_info_l['depth_map_path'] = osp.join(data_root, remove_leading_slash(depth_map_l))
- img_info_l['img_path'] = osp.join(data_root, remove_leading_slash(img_l))
- img_info_l['depth_fields'] = []
-
- filename = img_info_l['depth_map_path']
- ext_name_l = filename.replace('disp0.pfm', 'calib.txt')
- with open(ext_name_l, 'r') as f:
- ext_l = f.readlines()
-
- cam_info = ext_l[0].strip()
- cam_info_f = float(cam_info.split(' ')[0].split('[')[1])
- base = float(ext_l[3].strip().split('=')[1])
- doffs = float(ext_l[2].strip().split('=')[1])
-
- f = cam_info_f
- img_info_l['focal'] = f
- base = base
- img_info_l['depth_factor'] = base * f
- img_info_l['doffs'] = doffs
-
- img_infos.append(img_info_l)
- else:
- raise NotImplementedError
-
- # github issue:: make sure the same order
- img_infos = sorted(img_infos, key=lambda x: x['img_path'])
- if self.mode == 'train':
- img_infos = img_infos * 100
- return img_infos
-
-
- def __getitem__(self, idx):
-
- img_file_path = self.data_infos[idx]['img_path']
- disp_path = self.data_infos[idx]['depth_map_path']
- depth_factor = self.data_infos[idx]['depth_factor']
-
- height=2160
- width=3840
- height = 1840
- width = 2300
-
- image = Image.open(img_file_path).convert("RGB")
- image = np.asarray(image, dtype=np.uint8) / 255.0
- image = image.astype(np.float32)
-
- disp_gt, scale = readPFM(disp_path)
- disp_gt = disp_gt.astype(np.float32)
-
- h, w, _ = image.shape
- h_start = int(h / 2 - height / 2)
- h_end = h_start + height
- w_start = int(w / 2 - width / 2)
- w_end = w_start + width
- image = image[h_start:h_end, w_start:w_end, :]
- disp_gt = disp_gt[h_start:h_end, w_start:w_end]
-
- disp_gt_copy = disp_gt.copy()
- disp_gt = disp_gt[..., np.newaxis]
- invalid_mask = disp_gt == np.inf
-
- depth_gt = depth_factor / (disp_gt + self.data_infos[idx]['doffs'])
- depth_gt = depth_gt / 1000
- depth_gt[invalid_mask] = 0 # set to a invalid number
- disp_gt_copy[invalid_mask[:, :, 0]] = 0
- focal = self.data_infos[idx]['focal']
-
- bbox = None
- bboxs_res = None
- crop_areas = None
- bboxs_roi = None # hack for infer
-
- if self.mode == 'train':
- image, depth_gt = self.train_preprocess(image, depth_gt)
- img_temp = copy.deepcopy(image)
- depth_gt_temp = copy.deepcopy(depth_gt)
-
- if self.random_crop: # use in sec_stage
- if self.consistency_training:
- crop_y1, crop_y2, crop_x1, crop_x2 = self.get_crop_bbox(image) # ensure the prob of crop is the same
- while True:
- # shift_x = random.randint(self.overlap_length//3, self.overlap_length)
- # shift_y = random.randint(self.overlap_length//3, self.overlap_length)
- shift_x = self.overlap_length_w
- shift_y = self.overlap_length_h
- if random.random() > 0.5:
- shift_x = shift_x * -1
- if random.random() > 0.5:
- shift_y = shift_y * -1
- crop_y1_shift, crop_y2_shift, crop_x1_shift, crop_x2_shift = crop_y1 + shift_y, crop_y2 + shift_y, crop_x1 + shift_x, crop_x2 + shift_x
- if crop_y1_shift > 0 and crop_x1_shift > 0 and crop_y2_shift < image.shape[0] and crop_x2_shift < image.shape[1]:
- break
- bbox_ori = (crop_y1, crop_y2, crop_x1, crop_x2)
- bbox_shift = (crop_y1_shift, crop_y2_shift, crop_x1_shift, crop_x2_shift)
- image_ori, crop_area_ori = self.crop(image, bbox_ori, tmp=True)
- image_shift, crop_area_shift = self.crop(image, bbox_shift, tmp=True)
- depth_gt_ori = self.crop(depth_gt, bbox_ori)
- depth_gt_shift = self.crop(depth_gt, bbox_shift)
- disp_gt_copy_ori = self.crop(disp_gt_copy, bbox_ori)
- disp_gt_copy_shift = self.crop(disp_gt_copy, bbox_shift)
-
- bboxs_ori = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384])
- bboxs_shift = torch.tensor([crop_x1_shift / width * 512, crop_y1_shift / height * 384, crop_x2_shift / width * 512, crop_y2_shift / height * 384])
-
- else:
- bbox = self.get_crop_bbox(image)
- image, crop_area = self.crop(image, bbox, tmp=True)
- depth_gt = self.crop(depth_gt, bbox)
- disp_gt_copy = self.crop(disp_gt_copy, bbox)
-
- crop_y1, crop_y2, crop_x1, crop_x2 = bbox
- bboxs_res = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384]) # coord in 384, 512
-
-
- mask = np.logical_and(depth_gt > self.config.min_depth,
- depth_gt < self.config.max_depth).squeeze()[None, ...]
- mask_raw = np.logical_and(depth_gt_temp > self.config.min_depth, depth_gt_temp < self.config.max_depth).squeeze()[None, ...]
- sample = {'image': image, 'depth': depth_gt, 'focal': focal, 'mask': mask, 'image_raw': image.copy(), 'mask_raw': mask_raw}
-
-
- if self.random_crop:
- if self.consistency_training:
- image = np.concatenate([image_ori, image_shift], axis=-1)
- depth_gt = np.concatenate([depth_gt_ori, depth_gt_shift], axis=-1)
- crop_area = np.concatenate([crop_area_ori, crop_area_shift], axis=-1)
- bboxs_res = torch.cat([bboxs_ori, bboxs_shift], dim=-1)
- mask = np.logical_and(depth_gt > self.config.min_depth,
- depth_gt < self.config.max_depth)
-
- # hack the sample dict
- sample['image'] = image
- sample['depth'] = depth_gt
- sample['crop_area'] = crop_area
- sample['bbox'] = bboxs_res
- sample['shift'] = torch.tensor([shift_y, shift_x]) # h direction, then w direction
- sample['mask'] = mask
-
- else:
- if bboxs_res is not None:
- sample['bbox'] = bboxs_res
- sample['crop_area'] = crop_area
-
- if self.sampled_training:
- self.data_sampler(sample, disp_gt_copy)
-
- # update mask
- sample_points = sample['sample_points']
- sample_mask = np.logical_and(sample_points[:, -1] > self.config.min_depth,
- sample_points[:, -1] < self.config.max_depth).squeeze()[None, ...]
- sample['sample_mask'] = sample_mask
-
-
- else:
- # nothing needs to be changed for consistency training.
-
- img_temp = copy.deepcopy(image)
- depth_gt_temp = copy.deepcopy(depth_gt)
-
- if self.sec_stage:
- # x_start, y_start = [0, 540, 1080, 1620], [0, 960, 1920, 2880]
- x_start, y_start = [0 + 3 * self.overlap / 2, 540 + self.overlap / 2, 1080 - self.overlap / 2, 1620 - 3 * self.overlap / 2], \
- [0 + 3 * self.overlap / 2, 960 + self.overlap / 2, 1920 - self.overlap / 2, 2880 - 3 * self.overlap / 2]
- img_crops = []
- bboxs_roi = []
- crop_areas = []
- for x in x_start:
- for y in y_start:
- bbox = (int(x), int(x+540), int(y), int(y+960))
- img_crop, crop_area = self.crop(image, bbox, tmp=True)
- img_crops.append(img_crop)
- crop_areas.append(crop_area)
- crop_y1, crop_y2, crop_x1, crop_x2 = bbox
- bbox_roi = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384])
- bboxs_roi.append(bbox_roi)
-
- image = img_crops
- bboxs_roi = torch.stack(bboxs_roi, dim=0)
-
- # bbox = (820, 1360 ,1440, 2400) # a hack version for quick evaluation
- # image = self.crop(image, bbox)
- # depth_gt = self.crop(depth_gt, bbox)
- # disp_gt_copy = self.crop(disp_gt_copy, bbox)
-
- mask = np.logical_and(depth_gt > self.config.min_depth,
- depth_gt < self.config.max_depth).squeeze()[None, ...]
-
- disp_gt_edges = get_boundaries(disp_gt_copy, th=1, dilation=0)
-
- if self.mode == 'online_eval':
- sample = {'image': image, 'depth': depth_gt, 'focal': focal, 'has_valid_depth': True,
- 'image_path': img_file_path, 'depth_path': disp_path, 'depth_factor_path': depth_factor,
- 'mask': mask, 'image_raw': image.copy(), 'disp_gt_edges': disp_gt_edges}
- if bboxs_roi is not None:
- sample['bbox'] = bboxs_roi
- if crop_areas is not None:
- sample['crop_area'] = crop_areas
-
-
- else:
- sample = {'image': image, 'focal': focal, 'image_raw': image.copy(), 'disp_gt_edges': disp_gt_edges, 'image_path': img_file_path}
- if bboxs_roi is not None:
- sample['bbox'] = bboxs_roi
- if crop_areas is not None:
- sample['crop_area'] = crop_areas
-
- if self.transform:
- sample['img_temp'] = img_temp
- sample['depth_gt_temp'] = depth_gt_temp
- sample = self.transform(sample)
-
- sample['dataset'] = self.config.dataset
- return sample
-
- def __len__(self):
- return len(self.data_infos)
-
-def get_mid_loader(config, mode, transform):
- if mode == 'train':
- log = 0
- dataset = MiddleBurry(config, mode, config.data_path, config.filenames_train)
- if config.distributed:
- train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
- else:
- train_sampler = None
-
- dataloader = DataLoader(dataset,
- batch_size=config.batch_size,
- shuffle=(train_sampler is None),
- num_workers=config.workers,
- pin_memory=True,
- persistent_workers=True,
- sampler=train_sampler)
-
- elif mode == 'online_eval':
- dataset = MiddleBurry(config, mode, config.data_path, config.filenames_val)
- # dataset = U4KDataset(config, mode, config.data_path, config.filenames_train)
-
-
- if config.distributed: # redundant. here only for readability and to be more explicit
- # Give whole test set to all processes (and report evaluation only on one) regardless
- eval_sampler = None
- else:
- eval_sampler = None
-
- dataloader = DataLoader(dataset, 1,
- shuffle=False,
- num_workers=1,
- pin_memory=False,
- sampler=eval_sampler)
-
- else:
- dataset = MiddleBurry(config, mode, config.data_path, config.filenames_test)
- dataloader = DataLoader(dataset, 1, shuffle=False, num_workers=1)
-
- return dataloader
-
diff --git a/zoedepth/data/u4k.py b/zoedepth/data/u4k.py
deleted file mode 100644
index 0249c239486ee7e78c6ce261c2e85b9a2e9213c1..0000000000000000000000000000000000000000
--- a/zoedepth/data/u4k.py
+++ /dev/null
@@ -1,668 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Zhenyu Li
-
-# This file is partly inspired from ZoeDepth (https://github.com/isl-org/ZoeDepth/blob/main/zoedepth/data/data_mono.py); author: Shariq Farooq Bhat
-
-import os
-import numpy as np
-import torch
-from PIL import Image
-from torch.utils.data import DataLoader, Dataset
-from torchvision import transforms
-import os.path as osp
-import random
-import torch.nn as nn
-import cv2
-import copy
-from zoedepth.utils.misc import get_boundaries
-from zoedepth.models.base_models.midas import Resize
-
-class SampleDataPairs(object):
-
- def __init__(self,
- num_sample_inout=50000,
- sampling_strategy='random', # or 'dda'
- dilation_factor=10,
- crop_size=(2160, 3840),
- ):
-
- self.num_sample_inout = num_sample_inout
- self.sampling_strategy = sampling_strategy
- self.dilation_factor = dilation_factor
- self.crop_height, self.crop_width = crop_size[0], crop_size[1]
- self.__init_grid()
-
- def __init_grid(self):
- nu = np.linspace(0, self.crop_width - 1, self.crop_width)
- nv = np.linspace(0, self.crop_height - 1, self.crop_height)
- u, v = np.meshgrid(nu, nv)
-
- self.u = u.flatten()
- self.v = v.flatten()
-
- def get_coords(self, gt):
- # get subpixel coordinates
- u = self.u + np.random.random_sample(self.u.size)
- v = self.v + np.random.random_sample(self.v.size)
-
- # use nearest neighbor to get gt for each samples
- d = gt[np.clip(np.rint(v).astype(np.uint16), 0, self.crop_height-1),
- np.clip(np.rint(u).astype(np.uint16), 0, self.crop_width-1)]
-
- # remove invalid depth values
- u = u[np.nonzero(d)]
- v = v[np.nonzero(d)]
- d = d[np.nonzero(d)]
-
- return np.stack((u, v, d), axis=-1)
-
-
- def get_boundaries(self, disp, th=1., dilation=10):
- edges_y = np.logical_or(np.pad(np.abs(disp[1:, :] - disp[:-1, :]) > th, ((1, 0), (0, 0))),
- np.pad(np.abs(disp[:-1, :] - disp[1:, :]) > th, ((0, 1), (0, 0))))
- edges_x = np.logical_or(np.pad(np.abs(disp[:, 1:] - disp[:, :-1]) > th, ((0, 0), (1, 0))),
- np.pad(np.abs(disp[:, :-1] - disp[:,1:]) > th, ((0, 0), (0, 1))))
- edges = np.logical_or(edges_y, edges_x).astype(np.float32)
-
- if dilation > 0:
- kernel = np.ones((dilation, dilation), np.uint8)
- edges = cv2.dilate(edges, kernel, iterations=1)
-
- return edges
-
-
- def __call__(self, results, disp_gt_copy):
- """Call function to load multiple types annotations.
-
- Args:
- results (dict): Result dict from :obj:`depth.CustomDataset`.
-
- Returns:
- dict: The dict contains loaded depth estimation annotations.
- """
-
- gt = results['depth']
- gt_squeeze = gt[:, :, 0]
- if self.sampling_strategy == "random":
- random_points = self.get_coords(gt_squeeze)
- idx = np.random.choice(random_points.shape[0], self.num_sample_inout)
- points = random_points[idx, :]
-
- elif self.sampling_strategy == "dda":
- disp_gt_squeeze = disp_gt_copy
- edges = self.get_boundaries(disp_gt_squeeze, dilation=self.dilation_factor)
- random_points = self.get_coords(gt_squeeze * (1. - edges))
- edge_points = self.get_coords(gt_squeeze * edges)
-
- # if edge points exist
- if edge_points.shape[0]>0 and random_points.shape[0]>0:
- # Check tot num of edge points
- cond = edges.sum()//2 - self.num_sample_inout//2 < 0
- tot= (self.num_sample_inout - int(edges.sum())//2, int(edges.sum())//2) if cond else \
- (self.num_sample_inout//2, self.num_sample_inout//2)
-
- idx = np.random.choice(random_points.shape[0], tot[0])
- idx_edges = np.random.choice(edge_points.shape[0], tot[1])
- points = np.concatenate([edge_points[idx_edges, :], random_points[idx, :]], 0)
- # use uniform sample otherwise
- else:
- random_points = self.get_coords(gt_squeeze)
- idx = np.random.choice(random_points.shape[0], self.num_sample_inout)
- points = random_points[idx, :]
-
-
- results['sample_points'] = points
-
- return results
-
-def _is_pil_image(img):
- return isinstance(img, Image.Image)
-
-def _is_numpy_image(img):
- return isinstance(img, np.ndarray) and (img.ndim in {2, 3})
-
-class ToTensor(object):
- def __init__(self, mode, do_normalize=False, size=None, sec_stage=False):
- self.mode = mode
- # don't do normalization as default
- self.normalize = transforms.Normalize(
- mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if do_normalize else nn.Identity()
- self.size = size
- if size is not None:
- # self.resize = transforms.Resize(size=size)
- net_h, net_w = size
- self.resize = Resize(net_w, net_h, keep_aspect_ratio=False, ensure_multiple_of=32, resize_method="minimal")
- else:
- self.resize = nn.Identity()
- self.sec_stage = sec_stage
-
- def __call__(self, sample):
- image, focal = sample['image'], sample['focal']
- crop_areas = sample.get('crop_area', None)
-
- if isinstance(image, list):
- # there must be crop_areas
- # only infer on eval sec_stage
- imgs_process = []
- crp_process = []
- for img, crp in zip(image, crop_areas):
- img = self.to_tensor(img)
- img = self.normalize(img)
- img = img.unsqueeze(dim=0)
- img = self.resize(img)
- img = img.squeeze(dim=0)
- imgs_process.append(img)
-
- crp = self.to_tensor(crp)
- crp = crp.unsqueeze(dim=0)
- crp = self.resize(crp)
- crp = crp.squeeze(dim=0)
- crp_process.append(crp)
-
- image = torch.cat(imgs_process, dim=0)
- crop_areas = torch.cat(crp_process, dim=0)
-
- img_temp = sample['img_temp']
- img_temp = self.to_tensor(img_temp)
- img_temp = self.normalize(img_temp)
- img_temp = img_temp.unsqueeze(dim=0)
- img_temp = self.resize(img_temp) #NOTE: hack
- img_temp = img_temp.squeeze(dim=0)
- image_raw = copy.deepcopy(img_temp)
-
-
- else:
- image = self.to_tensor(image)
- image = self.normalize(image)
-
- if crop_areas is not None:
- crop_areas = self.to_tensor(crop_areas)
- crop_areas = crop_areas.unsqueeze(dim=0)
- crop_areas = self.resize(crop_areas)
- crop_areas = crop_areas.squeeze(dim=0)
-
- if self.sec_stage:
- img_temp = sample['img_temp']
- img_temp = self.to_tensor(img_temp)
- img_temp = self.normalize(img_temp)
-
- img_temp = img_temp.unsqueeze(dim=0)
- img_temp = self.resize(img_temp)
- image_raw = img_temp.squeeze(dim=0)
-
- image = image.unsqueeze(dim=0)
- image = self.resize(image)
- image = image.squeeze(dim=0)
-
- else:
- # in the first stage, this hr info is reserved
- image_raw = copy.deepcopy(image)
- image = image.unsqueeze(dim=0)
- image = self.resize(image)
- image = image.squeeze(dim=0)
-
- if self.mode == 'test':
- return_dict = {'image': image, 'focal': focal}
- if crop_areas is not None:
- return_dict['crop_area'] = crop_areas
- return return_dict
-
-
- depth = sample['depth']
- depth = self.to_tensor(depth)
- depth_gt_temp = sample['depth_gt_temp']
- depth_gt_raw = self.to_tensor(depth_gt_temp)
-
- if self.mode == 'train':
- return_dict = {**sample, 'image': image, 'depth': depth, 'focal': focal, 'image_raw': image_raw, 'depth_raw': depth_gt_raw}
- if crop_areas is not None:
- return_dict['crop_area'] = crop_areas
- return return_dict
- else:
- has_valid_depth = sample['has_valid_depth']
- # image = self.resize(image)
- return_dict = {**sample, 'image': image, 'depth': depth, 'focal': focal, 'image_raw': image_raw,
- 'has_valid_depth': has_valid_depth, 'image_path': sample['image_path'], 'depth_path': sample['depth_path'],
- 'depth_raw': depth_gt_raw}
- if crop_areas is not None:
- return_dict['crop_area'] = crop_areas
- return return_dict
-
- def to_tensor(self, pic):
- if isinstance(pic, np.ndarray):
- img = torch.from_numpy(pic.transpose((2, 0, 1))) # img here
- return img
-
-
-def preprocessing_transforms(mode, sec_stage=False, **kwargs):
- return transforms.Compose([
- ToTensor(mode=mode, sec_stage=sec_stage, **kwargs)
- ])
-
-def remove_leading_slash(s):
- if s[0] == '/' or s[0] == '\\':
- return s[1:]
- return s
-
-class U4KDataset(Dataset):
- def __init__(self, config, mode, data_root, split):
- self.mode = mode
- self.config = config
- self.data_root = data_root
- self.split = split
-
- img_size = self.config.get("img_size", None)
- img_size = img_size if self.config.get(
- "do_input_resize", False) else None
- self.sec_stage = self.config.get("sec_stage", False)
- self.transform = preprocessing_transforms(mode, size=img_size, sec_stage=self.sec_stage)
-
- self.data_infos = self.load_data_list()
-
- self.sampled_training = self.config.get("sampled_training", False)
- if self.sampled_training:
- self.data_sampler = SampleDataPairs(
- num_sample_inout=config.num_sample_inout,
- sampling_strategy=config.sampling_strategy, # or 'dda'
- dilation_factor=config.dilation_factor,
- crop_size=config.transform_sample_gt_size)
-
- self.random_crop = self.config.get("random_crop", False)
- self.crop_size = [540, 960] # 1/4
- self.overlap = self.config.get("overlap", False)
-
- self.consistency_training = self.config.get("consistency_training", False)
- self.overlap_length_h = self.config.get("overlap_length_h", int(256))
- self.overlap_length_w = self.config.get("overlap_length_w", int(384))
- print("current overlap_length_h and overlap_length_w are {} and {}".format(self.overlap_length_h, self.overlap_length_w))
-
-
- def load_data_list(self):
- """Load annotation from directory.
- Args:
- data_root (str): Data root for img_dir/ann_dir.
- split (str|None): Split txt file. If split is specified, only file
- with suffix in the splits will be loaded. Otherwise, all images
- in img_dir/ann_dir will be loaded. Default: None
- Returns:
- list[dict]: All image info of dataset.
- """
- data_root = self.data_root
- split = self.split
-
- self.invalid_depth_num = 0
- img_infos = []
- if split is not None:
- with open(split) as f:
- for line in f:
- img_info_l = dict()
- # img_info_r = dict()
-
- img_l, img_r, depth_map_l, depth_map_r = line.strip().split(" ")
-
- # HACK: a hack to replace the png with raw to accelerate training
- img_l = img_l[:-3] + 'raw'
- # img_r = img_r[:-3] + 'raw'
-
- img_info_l['depth_map_path'] = osp.join(data_root, remove_leading_slash(depth_map_l))
- # img_info_r['depth_map_path'] = osp.join(data_root, remove_leading_slash(depth_map_r))
-
- img_info_l['img_path'] = osp.join(data_root, remove_leading_slash(img_l))
- # img_info_r['filename'] = osp.join(data_root, remove_leading_slash(img_r))
-
- img_info_l['depth_fields'] = []
-
- filename = img_info_l['depth_map_path']
- ext_name_l = filename.replace('Disp0', 'Extrinsics0')
- ext_name_l = ext_name_l.replace('npy', 'txt')
- ext_name_r = filename.replace('Disp0', 'Extrinsics1')
- ext_name_r = ext_name_r.replace('npy', 'txt')
- with open(ext_name_l, 'r') as f:
- ext_l = f.readlines()
- with open(ext_name_r, 'r') as f:
- ext_r = f.readlines()
- f = float(ext_l[0].split(' ')[0])
- img_info_l['focal'] = f
- base = abs(float(ext_l[1].split(' ')[3]) - float(ext_r[1].split(' ')[3]))
- img_info_l['depth_factor'] = base * f
-
- img_infos.append(img_info_l)
- # img_infos.append(img_info_r)
- else:
- raise NotImplementedError
-
- # github issue:: make sure the same order
- img_infos = sorted(img_infos, key=lambda x: x['img_path'])
- return img_infos
-
-
- def augment_image(self, image):
- # gamma augmentation
- gamma = random.uniform(0.9, 1.1)
- image_aug = image ** gamma
-
- # brightness augmentation
- if self.config.dataset == 'nyu':
- brightness = random.uniform(0.75, 1.25)
- else:
- brightness = random.uniform(0.9, 1.1)
- image_aug = image_aug * brightness
-
- # color augmentation
- colors = np.random.uniform(0.9, 1.1, size=3)
- white = np.ones((image.shape[0], image.shape[1]))
- color_image = np.stack([white * colors[i] for i in range(3)], axis=2)
- image_aug *= color_image
- image_aug = np.clip(image_aug, 0, 1)
-
- return image_aug
-
- def train_preprocess(self, image, depth_gt):
- if self.config.aug:
- # Random flipping
- do_flip = random.random()
- if do_flip > 0.5:
- image = (image[:, ::-1, :]).copy()
- depth_gt = (depth_gt[:, ::-1, :]).copy()
-
- # Random gamma, brightness, color augmentation
- do_augment = random.random()
- if do_augment > 0.5:
- image = self.augment_image(image)
-
- return image, depth_gt
-
- def get_crop_bbox(self, img):
- """Randomly get a crop bounding box."""
- margin_h = max(img.shape[0] - self.crop_size[0], 0)
- margin_w = max(img.shape[1] - self.crop_size[1], 0)
- offset_h = np.random.randint(0, margin_h + 1)
- offset_w = np.random.randint(0, margin_w + 1)
- crop_y1, crop_y2 = offset_h, offset_h + self.crop_size[0]
- crop_x1, crop_x2 = offset_w, offset_w + self.crop_size[1]
-
- return crop_y1, crop_y2, crop_x1, crop_x2
-
- def crop(self, img, crop_bbox, tmp=False):
- """Crop from ``img``"""
-
- crop_y1, crop_y2, crop_x1, crop_x2 = crop_bbox
- if tmp:
- templete = np.zeros((img.shape[0], img.shape[1], 1), dtype=np.float32)
- templete[crop_y1:crop_y2, crop_x1:crop_x2, :] = 1.0
- img = img[crop_y1:crop_y2, crop_x1:crop_x2, ...]
- return img, templete
-
- else:
- img = img[crop_y1:crop_y2, crop_x1:crop_x2, ...]
- return img
-
- def __getitem__(self, idx):
-
- img_file_path = self.data_infos[idx]['img_path']
- disp_path = self.data_infos[idx]['depth_map_path']
- depth_factor = self.data_infos[idx]['depth_factor']
-
- height=2160
- width=3840
- image = np.fromfile(open(img_file_path, 'rb'), dtype=np.uint8).reshape(height, width, 3) / 255.0
- if self.config.get("use_rgb", False):
- image = image.astype(np.float32)[:, :, ::-1].copy()
- elif self.config.get("use_brg", False):
- image = image.astype(np.float32)[:, :, [0, 2, 1]].copy()
- elif self.config.get("use_gbr", False):
- image = image.astype(np.float32)[:, :, [1, 0, 2]].copy()
- elif self.config.get("use_rbg", False):
- image = image.astype(np.float32)[:, :, [2, 0, 1]].copy()
- elif self.config.get("use_grb", False):
- image = image.astype(np.float32)[:, :, [1, 2, 0]].copy()
- else:
- image = image.astype(np.float32)
-
- disp_gt = np.expand_dims(np.load(disp_path, mmap_mode='c'), -1)
- disp_gt = disp_gt.astype(np.float32)
- disp_gt_copy = disp_gt[:, :, 0].copy()
-
- depth_gt = depth_factor / disp_gt
- depth_gt[depth_gt > self.config.max_depth] = self.config.max_depth # for vis
- focal = self.data_infos[idx]['focal']
-
-
- bbox = None
- bboxs_res = None
- crop_areas = None
- bboxs_roi = None # hack for infer
-
- if self.mode == 'train':
- image, depth_gt = self.train_preprocess(image, depth_gt)
- img_temp = copy.deepcopy(image)
- depth_gt_temp = copy.deepcopy(depth_gt)
-
- if self.random_crop: # use in sec_stage
- if self.consistency_training:
- crop_y1, crop_y2, crop_x1, crop_x2 = self.get_crop_bbox(image) # ensure the prob of crop is the same
- while True:
- # shift_x = random.randint(self.overlap_length//3, self.overlap_length)
- # shift_y = random.randint(self.overlap_length//3, self.overlap_length)
- shift_x = self.overlap_length_w
- shift_y = self.overlap_length_h
- if random.random() > 0.5:
- shift_x = shift_x * -1
- if random.random() > 0.5:
- shift_y = shift_y * -1
- crop_y1_shift, crop_y2_shift, crop_x1_shift, crop_x2_shift = crop_y1 + shift_y, crop_y2 + shift_y, crop_x1 + shift_x, crop_x2 + shift_x
- if crop_y1_shift > 0 and crop_x1_shift > 0 and crop_y2_shift < image.shape[0] and crop_x2_shift < image.shape[1]:
- break
- bbox_ori = (crop_y1, crop_y2, crop_x1, crop_x2)
- bbox_shift = (crop_y1_shift, crop_y2_shift, crop_x1_shift, crop_x2_shift)
- image_ori, crop_area_ori = self.crop(image, bbox_ori, tmp=True)
- image_shift, crop_area_shift = self.crop(image, bbox_shift, tmp=True)
- depth_gt_ori = self.crop(depth_gt, bbox_ori)
- depth_gt_shift = self.crop(depth_gt, bbox_shift)
- disp_gt_copy_ori = self.crop(disp_gt_copy, bbox_ori)
- disp_gt_copy_shift = self.crop(disp_gt_copy, bbox_shift)
-
- bboxs_ori = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384])
- bboxs_shift = torch.tensor([crop_x1_shift / width * 512, crop_y1_shift / height * 384, crop_x2_shift / width * 512, crop_y2_shift / height * 384])
- bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2])
- bboxs_raw_shift = torch.tensor([crop_x1_shift, crop_y1_shift, crop_x2_shift, crop_y2_shift])
-
-
- else:
- bbox = self.get_crop_bbox(image)
- image, crop_area = self.crop(image, bbox, tmp=True)
- depth_gt = self.crop(depth_gt, bbox)
- disp_gt_copy = self.crop(disp_gt_copy, bbox)
-
- crop_y1, crop_y2, crop_x1, crop_x2 = bbox
- bboxs_res = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384]) # coord in 384, 512
- bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2])
-
-
- mask = np.logical_and(depth_gt > self.config.min_depth,
- depth_gt < self.config.max_depth).squeeze()[None, ...]
- mask_raw = np.logical_and(depth_gt_temp > self.config.min_depth, depth_gt_temp < self.config.max_depth).squeeze()[None, ...]
- sample = {'image': image, 'depth': depth_gt, 'focal': focal, 'mask': mask, 'image_raw': image.copy(), 'mask_raw': mask_raw, 'image_path': img_file_path}
-
- if self.random_crop:
- if self.consistency_training:
- image = np.concatenate([image_ori, image_shift], axis=-1)
- depth_gt = np.concatenate([depth_gt_ori, depth_gt_shift], axis=-1)
- crop_area = np.concatenate([crop_area_ori, crop_area_shift], axis=-1)
- bboxs_res = torch.cat([bboxs_ori, bboxs_shift], dim=-1)
- bboxes_raw_res = torch.cat([bboxs_raw, bboxs_raw_shift], dim=-1)
- mask = np.logical_and(depth_gt > self.config.min_depth,
- depth_gt < self.config.max_depth)
-
- # hack the sample dict
- sample['image'] = image
- sample['depth'] = depth_gt
- sample['crop_area'] = crop_area
- sample['bbox'] = bboxs_res
- sample['bbox_raw'] = bboxes_raw_res
- sample['shift'] = torch.tensor([shift_y, shift_x]) # h direction, then w direction
- sample['mask'] = mask
-
- else:
- if bboxs_res is not None:
- sample['bbox'] = bboxs_res
- sample['bbox_raw'] = bboxs_raw
- sample['crop_area'] = crop_area
-
- if self.sampled_training:
- self.data_sampler(sample, disp_gt_copy)
-
- # update mask
- sample_points = sample['sample_points']
- sample_mask = np.logical_and(sample_points[:, -1] > self.config.min_depth,
- sample_points[:, -1] < self.config.max_depth).squeeze()[None, ...]
- sample['sample_mask'] = sample_mask
-
-
- else:
- # nothing needs to be changed for consistency training.
-
- img_temp = copy.deepcopy(image)
- depth_gt_temp = copy.deepcopy(depth_gt)
-
- if self.sec_stage:
- # x_start, y_start = [0, 540, 1080, 1620], [0, 960, 1920, 2880]
- x_start, y_start = [0 + 3 * self.overlap / 2, 540 + self.overlap / 2, 1080 - self.overlap / 2, 1620 - 3 * self.overlap / 2], \
- [0 + 3 * self.overlap / 2, 960 + self.overlap / 2, 1920 - self.overlap / 2, 2880 - 3 * self.overlap / 2]
- img_crops = []
- bboxs_roi = []
- crop_areas = []
- bboxs_raw_list = []
- for x in x_start:
- for y in y_start:
- bbox = (int(x), int(x+540), int(y), int(y+960))
- img_crop, crop_area = self.crop(image, bbox, tmp=True)
- img_crops.append(img_crop)
- crop_areas.append(crop_area)
- crop_y1, crop_y2, crop_x1, crop_x2 = bbox
- bbox_roi = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384])
- bboxs_roi.append(bbox_roi)
- bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2])
- bboxs_raw_list.append(bboxs_raw)
-
- image = img_crops
- bboxs_roi = torch.stack(bboxs_roi, dim=0)
- bboxs_raw = torch.stack(bboxs_raw_list, dim=0)
-
- # bbox = (820, 1360 ,1440, 2400) # a hack version for quick evaluation
- # image = self.crop(image, bbox)
- # depth_gt = self.crop(depth_gt, bbox)
- # disp_gt_copy = self.crop(disp_gt_copy, bbox)
-
- mask = np.logical_and(depth_gt > self.config.min_depth,
- depth_gt < self.config.max_depth).squeeze()[None, ...]
-
- disp_gt_edges = get_boundaries(disp_gt_copy, th=1, dilation=0)
-
- if self.mode == 'online_eval':
- sample = {'image': image, 'depth': depth_gt, 'focal': focal, 'has_valid_depth': True,
- 'image_path': img_file_path, 'depth_path': disp_path, 'depth_factor_path': depth_factor,
- 'mask': mask, 'image_raw': image.copy(), 'disp_gt_edges': disp_gt_edges, 'image_path': img_file_path}
- if bboxs_roi is not None:
- sample['bbox'] = bboxs_roi
- sample['bbox_raw'] = bboxs_raw
- if crop_areas is not None:
- sample['crop_area'] = crop_areas
-
-
- else:
- sample = {'image': image, 'focal': focal, 'image_raw': image.copy(), 'disp_gt_edges': disp_gt_edges, 'image_path': img_file_path}
- if bboxs_roi is not None:
- sample['bbox'] = bboxs_roi
- sample['bbox_raw'] = bboxs_raw
- if crop_areas is not None:
- sample['crop_area'] = crop_areas
-
- if self.transform:
- sample['img_temp'] = img_temp
- sample['depth_gt_temp'] = depth_gt_temp
- sample = self.transform(sample)
-
- sample['dataset'] = self.config.dataset
- return sample
-
- def __len__(self):
- return len(self.data_infos)
-
-
-def get_u4k_loader(config, mode, transform):
- if mode == 'train':
- dataset = U4KDataset(config, mode, config.data_path, config.filenames_train)
- dataset[0]
-
- if config.distributed:
- train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
- else:
- train_sampler = None
-
- dataloader = DataLoader(dataset,
- batch_size=config.batch_size,
- shuffle=(train_sampler is None),
- num_workers=config.workers,
- pin_memory=True,
- persistent_workers=True,
- sampler=train_sampler)
-
- elif mode == 'train_save':
- dataset = U4KDataset(config, 'online_eval', config.data_path, config.filenames_train)
-
- if config.distributed:
- train_sampler = None
- else:
- train_sampler = None
-
- dataloader = DataLoader(dataset, 1,
- shuffle=False,
- num_workers=1,
- pin_memory=False,
- sampler=train_sampler)
-
-
- elif mode == 'online_eval':
- dataset = U4KDataset(config, mode, config.data_path, config.filenames_val)
- # dataset = U4KDataset(config, mode, config.data_path, config.filenames_train)
-
-
- if config.distributed: # redundant. here only for readability and to be more explicit
- # Give whole test set to all processes (and report evaluation only on one) regardless
- eval_sampler = None
- else:
- eval_sampler = None
-
- dataloader = DataLoader(dataset, 1,
- shuffle=False,
- num_workers=1,
- pin_memory=False,
- sampler=eval_sampler)
-
- else:
- dataset = U4KDataset(config, mode, config.data_path, config.filenames_test)
- dataloader = DataLoader(dataset, 1, shuffle=False, num_workers=1)
-
- return dataloader
\ No newline at end of file
diff --git a/zoedepth/models/base_models/depth_anything.py b/zoedepth/models/base_models/depth_anything.py
new file mode 100644
index 0000000000000000000000000000000000000000..63f9a0b8f1c1b8507ca73d441e7f99019347cb7f
--- /dev/null
+++ b/zoedepth/models/base_models/depth_anything.py
@@ -0,0 +1,391 @@
+# MIT License
+
+# Copyright (c) 2022 Intelligent Systems Lab Org
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# File author: Shariq Farooq Bhat
+
+import torch
+import torch.nn as nn
+import numpy as np
+from torchvision.transforms import Normalize
+# from zoedepth.models.base_models.dpt_dinov2.dpt import DPT_DINOv2
+from depth_anything.dpt import DPT_DINOv2
+
+
+def denormalize(x):
+ """Reverses the imagenet normalization applied to the input.
+
+ Args:
+ x (torch.Tensor - shape(N,3,H,W)): input tensor
+
+ Returns:
+ torch.Tensor - shape(N,3,H,W): Denormalized input
+ """
+ mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(x.device)
+ std = torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(x.device)
+ return x * std + mean
+
+def get_activation(name, bank):
+ def hook(model, input, output):
+ bank[name] = output
+ return hook
+
+
+class Resize(object):
+ """Resize sample to given size (width, height).
+ """
+
+ def __init__(
+ self,
+ width,
+ height,
+ resize_target=True,
+ keep_aspect_ratio=False,
+ ensure_multiple_of=1,
+ resize_method="lower_bound",
+ ):
+ """Init.
+ Args:
+ width (int): desired output width
+ height (int): desired output height
+ resize_target (bool, optional):
+ True: Resize the full sample (image, mask, target).
+ False: Resize image only.
+ Defaults to True.
+ keep_aspect_ratio (bool, optional):
+ True: Keep the aspect ratio of the input sample.
+ Output sample might not have the given width and height, and
+ resize behaviour depends on the parameter 'resize_method'.
+ Defaults to False.
+ ensure_multiple_of (int, optional):
+ Output width and height is constrained to be multiple of this parameter.
+ Defaults to 1.
+ resize_method (str, optional):
+ "lower_bound": Output will be at least as large as the given size.
+ "upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
+ "minimal": Scale as least as possible. (Output size might be smaller than given size.)
+ Defaults to "lower_bound".
+ """
+ # print("Params passed to Resize transform:")
+ # print("\twidth: ", width)
+ # print("\theight: ", height)
+ # print("\tresize_target: ", resize_target)
+ # print("\tkeep_aspect_ratio: ", keep_aspect_ratio)
+ # print("\tensure_multiple_of: ", ensure_multiple_of)
+ # print("\tresize_method: ", resize_method)
+
+ self.__width = width
+ self.__height = height
+
+ self.__keep_aspect_ratio = keep_aspect_ratio
+ self.__multiple_of = ensure_multiple_of
+ self.__resize_method = resize_method
+
+ def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
+ y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+ if max_val is not None and y > max_val:
+ y = (np.floor(x / self.__multiple_of)
+ * self.__multiple_of).astype(int)
+
+ if y < min_val:
+ y = (np.ceil(x / self.__multiple_of)
+ * self.__multiple_of).astype(int)
+
+ return y
+
+ def get_size(self, width, height):
+ # determine new height and width
+ scale_height = self.__height / height
+ scale_width = self.__width / width
+
+ if self.__keep_aspect_ratio:
+ if self.__resize_method == "lower_bound":
+ # scale such that output size is lower bound
+ if scale_width > scale_height:
+ # fit width
+ scale_height = scale_width
+ else:
+ # fit height
+ scale_width = scale_height
+ elif self.__resize_method == "upper_bound":
+ # scale such that output size is upper bound
+ if scale_width < scale_height:
+ # fit width
+ scale_height = scale_width
+ else:
+ # fit height
+ scale_width = scale_height
+ elif self.__resize_method == "minimal":
+ # scale as least as possbile
+ if abs(1 - scale_width) < abs(1 - scale_height):
+ # fit width
+ scale_height = scale_width
+ else:
+ # fit height
+ scale_width = scale_height
+ else:
+ raise ValueError(
+ f"resize_method {self.__resize_method} not implemented"
+ )
+
+ if self.__resize_method == "lower_bound":
+ new_height = self.constrain_to_multiple_of(
+ scale_height * height, min_val=self.__height
+ )
+ new_width = self.constrain_to_multiple_of(
+ scale_width * width, min_val=self.__width
+ )
+ elif self.__resize_method == "upper_bound":
+ new_height = self.constrain_to_multiple_of(
+ scale_height * height, max_val=self.__height
+ )
+ new_width = self.constrain_to_multiple_of(
+ scale_width * width, max_val=self.__width
+ )
+ elif self.__resize_method == "minimal":
+ new_height = self.constrain_to_multiple_of(scale_height * height)
+ new_width = self.constrain_to_multiple_of(scale_width * width)
+ else:
+ raise ValueError(
+ f"resize_method {self.__resize_method} not implemented")
+
+ return (new_width, new_height)
+
+ def __call__(self, x):
+ width, height = self.get_size(*x.shape[-2:][::-1])
+ return nn.functional.interpolate(x, (int(height), int(width)), mode='bilinear', align_corners=True)
+
+class PrepForMidas(object):
+ def __init__(self, resize_mode="minimal", keep_aspect_ratio=True, img_size=384, do_resize=True):
+ if isinstance(img_size, int):
+ img_size = (img_size, img_size)
+ net_h, net_w = img_size
+ # self.normalization = Normalize(
+ # mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+ self.normalization = Normalize(
+ mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+ self.resizer = Resize(net_w, net_h, keep_aspect_ratio=keep_aspect_ratio, ensure_multiple_of=14, resize_method=resize_mode) \
+ if do_resize else nn.Identity()
+
+ def __call__(self, x):
+ return self.normalization(self.resizer(x))
+
+
+class DepthAnythingCore(nn.Module):
+ def __init__(self, midas, trainable=False, fetch_features=True, layer_names=('out_conv', 'l4_rn', 'r4', 'r3', 'r2', 'r1'), freeze_bn=False, keep_aspect_ratio=True, img_size=384, core_type='vits', **kwargs):
+ """Midas Base model used for multi-scale feature extraction.
+
+ Args:
+ midas (torch.nn.Module): Midas model.
+ trainable (bool, optional): Train midas model. Defaults to False.
+ fetch_features (bool, optional): Extract multi-scale features. Defaults to True.
+ layer_names (tuple, optional): Layers used for feature extraction. Order = (head output features, last layer features, ...decoder features). Defaults to ('out_conv', 'l4_rn', 'r4', 'r3', 'r2', 'r1').
+ freeze_bn (bool, optional): Freeze BatchNorm. Generally results in better finetuning performance. Defaults to False.
+ keep_aspect_ratio (bool, optional): Keep the aspect ratio of input images while resizing. Defaults to True.
+ img_size (int, tuple, optional): Input resolution. Defaults to 384.
+ """
+ super().__init__()
+ self.core_type = core_type
+ self.core = midas
+ self.output_channels = None
+ self.core_out = {}
+ self.trainable = trainable
+ self.fetch_features = fetch_features
+ # midas.scratch.output_conv = nn.Identity()
+ self.handles = []
+ # self.layer_names = ['out_conv','l4_rn', 'r4', 'r3', 'r2', 'r1']
+ self.layer_names = layer_names
+
+ self.set_trainable(trainable)
+ self.set_fetch_features(fetch_features)
+
+ self.prep = PrepForMidas(keep_aspect_ratio=keep_aspect_ratio,
+ img_size=img_size, do_resize=kwargs.get('do_resize', True))
+
+ if freeze_bn:
+ self.freeze_bn()
+
+ def set_trainable(self, trainable):
+ self.trainable = trainable
+ if trainable:
+ self.unfreeze()
+ else:
+ self.freeze()
+ return self
+
+ def set_fetch_features(self, fetch_features):
+ self.fetch_features = fetch_features
+ if fetch_features:
+ if len(self.handles) == 0:
+ self.attach_hooks(self.core)
+ else:
+ self.remove_hooks()
+ return self
+
+ def freeze(self):
+ for p in self.parameters():
+ p.requires_grad = False
+ self.trainable = False
+ return self
+
+ def unfreeze(self):
+ for p in self.parameters():
+ p.requires_grad = True
+ self.trainable = True
+ return self
+
+ def freeze_bn(self):
+ for m in self.modules():
+ if isinstance(m, nn.BatchNorm2d):
+ m.eval()
+ return self
+
+ def forward(self, x, denorm=False, return_rel_depth=False):
+ # print('input to midas:', x.shape)
+ with torch.no_grad():
+ if denorm:
+ x = denormalize(x)
+ x = self.prep(x)
+
+ with torch.set_grad_enabled(self.trainable):
+
+ rel_depth = self.core(x)
+ if not self.fetch_features:
+ return rel_depth
+ out = [self.core_out[k] for k in self.layer_names]
+
+ if return_rel_depth:
+ return rel_depth, out
+ return out
+
+ def get_rel_pos_params(self):
+ for name, p in self.core.pretrained.named_parameters():
+ if "pos_embed" in name:
+ yield p
+
+ def get_enc_params_except_rel_pos(self):
+ for name, p in self.core.pretrained.named_parameters():
+ if "pos_embed" not in name:
+ yield p
+
+ def freeze_encoder(self, freeze_rel_pos=False):
+ if freeze_rel_pos:
+ for p in self.core.pretrained.parameters():
+ p.requires_grad = False
+ else:
+ for p in self.get_enc_params_except_rel_pos():
+ p.requires_grad = False
+ return self
+
+ def attach_hooks(self, midas):
+ if len(self.handles) > 0:
+ self.remove_hooks()
+ if "out_conv" in self.layer_names:
+ self.handles.append(list(midas.depth_head.scratch.output_conv2.children())[
+ 1].register_forward_hook(get_activation("out_conv", self.core_out)))
+ if "r4" in self.layer_names:
+ self.handles.append(midas.depth_head.scratch.refinenet4.register_forward_hook(
+ get_activation("r4", self.core_out)))
+ if "r3" in self.layer_names:
+ self.handles.append(midas.depth_head.scratch.refinenet3.register_forward_hook(
+ get_activation("r3", self.core_out)))
+ if "r2" in self.layer_names:
+ self.handles.append(midas.depth_head.scratch.refinenet2.register_forward_hook(
+ get_activation("r2", self.core_out)))
+ if "r1" in self.layer_names:
+ self.handles.append(midas.depth_head.scratch.refinenet1.register_forward_hook(
+ get_activation("r1", self.core_out)))
+ if "l4_rn" in self.layer_names:
+ self.handles.append(midas.depth_head.scratch.layer4_rn.register_forward_hook(
+ get_activation("l4_rn", self.core_out)))
+
+ return self
+
+ def remove_hooks(self):
+ for h in self.handles:
+ h.remove()
+ return self
+
+ def __del__(self):
+ self.remove_hooks()
+
+ def set_output_channels(self):
+ if self.core_type == 'vits':
+ self.output_channels = [64, 64, 64, 64, 64]
+ elif self.core_type == 'vitb':
+ self.output_channels = [128, 128, 128, 128, 128]
+ elif self.core_type == 'vitl':
+ self.output_channels = [256, 256, 256, 256, 256]
+
+ @staticmethod
+ def build(midas_model_type="dinov2_large", train_midas=False, use_pretrained_midas=True, fetch_features=False, freeze_bn=True, force_keep_ar=False, force_reload=False, **kwargs):
+ if "img_size" in kwargs:
+ kwargs = DepthAnythingCore.parse_img_size(kwargs)
+ img_size = kwargs.pop("img_size", [384, 384])
+
+ if midas_model_type == 'vits':
+ depth_anything = DPT_DINOv2(encoder=midas_model_type, features=64, out_channels=[48, 96, 192, 384], use_clstoken=False)
+ state_dict = torch.load('/ibex/ai/home/liz0l/codes/ZoeDepth/depth_anything_vits14.pth', map_location='cpu')
+ elif midas_model_type == 'vitb':
+ depth_anything = DPT_DINOv2(encoder=midas_model_type, features=128, out_channels=[96, 192, 384, 768], use_clstoken=False)
+ state_dict = torch.load('/ibex/ai/home/liz0l/codes/ZoeDepth/depth_anything_vitb14.pth', map_location='cpu')
+ elif midas_model_type == 'vitl':
+ depth_anything = DPT_DINOv2(encoder=midas_model_type, features=256, out_channels=[256, 512, 1024, 1024], use_clstoken=False)
+ state_dict = torch.load('/ibex/ai/home/liz0l/codes/ZoeDepth/depth_anything_vitl14.pth', map_location='cpu')
+ else:
+ raise NotImplementedError
+
+ depth_anything.load_state_dict(state_dict)
+
+ kwargs.update({'keep_aspect_ratio': force_keep_ar})
+
+ depth_anything_core = DepthAnythingCore(depth_anything, trainable=train_midas, fetch_features=fetch_features,
+ freeze_bn=freeze_bn, img_size=img_size, core_type=midas_model_type, **kwargs)
+
+ depth_anything_core.set_output_channels()
+ return depth_anything_core
+
+ @staticmethod
+ def parse_img_size(config):
+ assert 'img_size' in config
+ if isinstance(config['img_size'], str):
+ assert "," in config['img_size'], "img_size should be a string with comma separated img_size=H,W"
+ config['img_size'] = list(map(int, config['img_size'].split(",")))
+ assert len(
+ config['img_size']) == 2, "img_size should be a string with comma separated img_size=H,W"
+ elif isinstance(config['img_size'], int):
+ config['img_size'] = [config['img_size'], config['img_size']]
+ else:
+ assert isinstance(config['img_size'], list) and len(
+ config['img_size']) == 2, "img_size should be a list of H,W"
+ return config
+
+
+nchannels2models = {
+ tuple([256]*5): ["DPT_BEiT_L_384", "DPT_BEiT_L_512", "DPT_BEiT_B_384", "DPT_SwinV2_L_384", "DPT_SwinV2_B_384", "DPT_SwinV2_T_256", "DPT_Large", "DPT_Hybrid"],
+ (512, 256, 128, 64, 64): ["MiDaS_small"]
+}
+
+# Model name to number of output channels
+MIDAS_SETTINGS = {m: k for k, v in nchannels2models.items()
+ for m in v
+ }
\ No newline at end of file
diff --git a/zoedepth/models/base_models/midas.py b/zoedepth/models/base_models/midas.py
index e26f8589502f8298bde8820262083f54b254f70e..7ad698d2906b405520a551ddccaa92eb81edc897 100644
--- a/zoedepth/models/base_models/midas.py
+++ b/zoedepth/models/base_models/midas.py
@@ -82,13 +82,13 @@ class Resize(object):
"minimal": Scale as least as possible. (Output size might be smaller than given size.)
Defaults to "lower_bound".
"""
- print("Params passed to Resize transform:")
- print("\twidth: ", width)
- print("\theight: ", height)
- print("\tresize_target: ", resize_target)
- print("\tkeep_aspect_ratio: ", keep_aspect_ratio)
- print("\tensure_multiple_of: ", ensure_multiple_of)
- print("\tresize_method: ", resize_method)
+ # print("Params passed to Resize transform:")
+ # print("\twidth: ", width)
+ # print("\theight: ", height)
+ # print("\tresize_target: ", resize_target)
+ # print("\tkeep_aspect_ratio: ", keep_aspect_ratio)
+ # print("\tensure_multiple_of: ", ensure_multiple_of)
+ # print("\tresize_method: ", resize_method)
self.__width = width
self.__height = height
@@ -170,7 +170,7 @@ class Resize(object):
def __call__(self, x):
width, height = self.get_size(*x.shape[-2:][::-1])
- return nn.functional.interpolate(x, (height, width), mode='bilinear', align_corners=True)
+ return nn.functional.interpolate(x, (int(height), int(width)), mode='bilinear', align_corners=True)
class PrepForMidas(object):
def __init__(self, resize_mode="minimal", keep_aspect_ratio=True, img_size=384, do_resize=True):
@@ -261,10 +261,7 @@ class MidasCore(nn.Module):
x = denormalize(x)
x = self.prep(x)
# print("Shape after prep: ", x.shape)
-
with torch.set_grad_enabled(self.trainable):
-
- # print("Input size to Midascore", x.shape)
rel_depth = self.core(x)
# print("Output from midas shape", rel_depth.shape)
if not self.fetch_features:
@@ -337,9 +334,11 @@ class MidasCore(nn.Module):
if "img_size" in kwargs:
kwargs = MidasCore.parse_img_size(kwargs)
img_size = kwargs.pop("img_size", [384, 384])
- print("img_size", img_size)
- midas = torch.hub.load("intel-isl/MiDaS", midas_model_type,
- pretrained=use_pretrained_midas, force_reload=force_reload)
+ # print("img_size", img_size)
+ # midas = torch.hub.load("intel-isl/MiDaS", midas_model_type,
+ # pretrained=use_pretrained_midas, force_reload=force_reload)
+ midas = torch.hub.load("AyaanShah2204/MiDaS", midas_model_type,
+ pretrained=use_pretrained_midas, force_reload=force_reload) # switcher to a better version?
kwargs.update({'keep_aspect_ratio': force_keep_ar})
midas_core = MidasCore(midas, trainable=train_midas, fetch_features=fetch_features,
freeze_bn=freeze_bn, img_size=img_size, **kwargs)
diff --git a/zoedepth/models/builder.py b/zoedepth/models/builder.py
index d3c18c96179da6056029e5407e1446a17abba42d..4363d59689158912a412feb5c296b4a72bc2c608 100644
--- a/zoedepth/models/builder.py
+++ b/zoedepth/models/builder.py
@@ -22,8 +22,6 @@
# File author: Shariq Farooq Bhat
-# This file may include modifications from author Zhenyu Li
-
from importlib import import_module
from zoedepth.models.depth_model import DepthModel
diff --git a/zoedepth/models/layers/dist_layers.py b/zoedepth/models/layers/dist_layers.py
index 5165fb3eb09c34f9fd281f1040afc9602055a047..572fc0e389cc970c96bb1d4a0c18d9f6e1b42d56 100644
--- a/zoedepth/models/layers/dist_layers.py
+++ b/zoedepth/models/layers/dist_layers.py
@@ -121,6 +121,7 @@ class ConditionalLogBinomial(nn.Module):
return self.log_binomial_transform(p, t)
+
class ConditionalLogBinomialV2(nn.Module):
def __init__(self, in_features, condition_dim, n_classes=256, bottleneck_factor=2, p_eps=1e-4, max_temp=50, min_temp=1e-7, act=torch.softmax):
"""Conditional Log Binomial distribution
@@ -163,4 +164,4 @@ class ConditionalLogBinomialV2(nn.Module):
"""
pt = self.mlp(torch.concat((x, cond), dim=1))
prob, shift = pt[:, :self.n_classes, ...], pt[:, self.n_classes:, ...]
- return prob, shift
+ return prob, shift
\ No newline at end of file
diff --git a/zoedepth/models/layers/fusion_network.py b/zoedepth/models/layers/fusion_network.py
deleted file mode 100644
index 3ecef6d4427201e99be3f163984909646e3efec6..0000000000000000000000000000000000000000
--- a/zoedepth/models/layers/fusion_network.py
+++ /dev/null
@@ -1,203 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Zhenyu Li
-
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from zoedepth.models.layers.swin_layers import G2LFusion
-from zoedepth.models.layers.transformer import TransformerEncoder, TransformerEncoderLayer
-from torchvision.ops import roi_align as torch_roi_align
-
-class DoubleConvWOBN(nn.Module):
- """(convolution => [BN] => ReLU) * 2"""
-
- def __init__(self, in_channels, out_channels, mid_channels=None):
- super().__init__()
- if not mid_channels:
- mid_channels = out_channels
- self.double_conv = nn.Sequential(
- nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=True),
- # nn.BatchNorm2d(mid_channels),
- nn.ReLU(inplace=True),
- nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=True),
- # nn.BatchNorm2d(mid_channels),
- nn.ReLU(inplace=True))
-
- def forward(self, x):
- return self.double_conv(x)
-
-class DoubleConv(nn.Module):
- """(convolution => [BN] => ReLU) * 2"""
-
- def __init__(self, in_channels, out_channels, mid_channels=None):
- super().__init__()
- if not mid_channels:
- mid_channels = out_channels
- self.double_conv = nn.Sequential(
- nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
- nn.BatchNorm2d(mid_channels),
- nn.ReLU(inplace=True),
- nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
- nn.BatchNorm2d(out_channels),
- nn.ReLU(inplace=True)
- )
-
- def forward(self, x):
- return self.double_conv(x)
-
-
-class Down(nn.Module):
- """Downscaling with maxpool then double conv"""
-
- def __init__(self, in_channels, out_channels):
- super().__init__()
- self.maxpool_conv = nn.Sequential(
- nn.MaxPool2d(2),
- DoubleConv(in_channels, out_channels)
- )
-
- def forward(self, x):
- return self.maxpool_conv(x)
-
-class Upv1(nn.Module):
- """Upscaling then double conv"""
-
- def __init__(self, in_channels, out_channels, mid_channels=None):
- super().__init__()
- self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
-
- if mid_channels is not None:
- self.conv = DoubleConvWOBN(in_channels, out_channels, mid_channels)
- else:
- self.conv = DoubleConvWOBN(in_channels, out_channels, in_channels)
-
- def forward(self, x1, x2):
- x1 = self.up(x1)
- x = torch.cat([x2, x1], dim=1)
- return self.conv(x)
-
-class UNetv1(nn.Module):
- def __init__(self, n_channels, g2l, pos_embed=False, use_area_prior=True):
- super(UNetv1, self).__init__()
- self.n_channels = n_channels
-
- self.inc = DoubleConv(n_channels, 32)
- self.down1 = Down(32, 256)
- self.down2 = Down(256, 256)
- self.down3 = Down(256, 256)
- self.down4 = Down(256, 256)
- self.down5 = Down(256, 256)
-
- self.up1 = Upv1(256+256+256, 256, 384)
- self.up2 = Upv1(256+256+256, 256, 384)
- self.up3 = Upv1(256+256+256, 256, 384)
- self.up4 = Upv1(256+256+256, 256, 384)
- self.up5 = Upv1(256+32+256, 32, 272)
-
- self.g2l = g2l
-
- if self.g2l:
- self.g2l_att = nn.ModuleList()
- win = 12
- in_channels = [32, 256, 256, 256, 256, 256]
- crf_dims = [32, 256, 256, 256, 256, 256]
-
- self.g2l5 = G2LFusion(input_dim=in_channels[5], embed_dim=crf_dims[5], window_size=win, num_heads=32, depth=4, num_patches=12*16)
- self.g2l4 = G2LFusion(input_dim=in_channels[4], embed_dim=crf_dims[4], window_size=win, num_heads=32, depth=4, num_patches=24*32)
- self.g2l3 = G2LFusion(input_dim=in_channels[3], embed_dim=crf_dims[3], window_size=win, num_heads=16, depth=3, num_patches=48*64)
- self.g2l2 = G2LFusion(input_dim=in_channels[2], embed_dim=crf_dims[2], window_size=win, num_heads=16, depth=3, num_patches=96*128)
- self.g2l1 = G2LFusion(input_dim=in_channels[1], embed_dim=crf_dims[1], window_size=win, num_heads=8, depth=2, num_patches=192*256)
- self.g2l0 = G2LFusion(input_dim=in_channels[0], embed_dim=crf_dims[0], window_size=win, num_heads=8, depth=2, num_patches=384*512)
-
- self.conv5 = DoubleConvWOBN(in_channels[4] * 2, in_channels[4], in_channels[4])
- self.conv4 = DoubleConvWOBN(in_channels[4] * 2, in_channels[4], in_channels[4])
- self.conv3 = DoubleConvWOBN(in_channels[3] * 2, in_channels[3], in_channels[3])
- self.conv2 = DoubleConvWOBN(in_channels[2] * 2, in_channels[2], in_channels[2])
- self.conv1 = DoubleConvWOBN(in_channels[1] * 2, in_channels[1], in_channels[1])
- self.conv0 = DoubleConvWOBN(in_channels[0] * 2, in_channels[0], in_channels[0])
-
- def forward(self,
- input_tensor,
- guide_plus,
- guide_cat,
- crop_area_resize=None,
- bbox=None,
- fine_feat_crop=None,
- coarse_feat_whole=None,
- coarse_feat_whole_hack=None,
- coarse_feat_crop=None):
-
- # apply unscaled feat to swin
- if coarse_feat_whole_hack is not None:
- coarse_feat_whole = coarse_feat_whole_hack
-
- if crop_area_resize is None:
- not_use_prior = True
- else:
- not_use_prior = False
-
- x1 = self.inc(input_tensor)
- x2 = self.down1(x1)
- x3 = self.down2(x2)
- x4 = self.down3(x3)
- x5 = self.down4(x4)
- x6 = self.down5(x5)
- if self.g2l:
- g2l_feat5 = self.g2l5(coarse_feat_whole[0], crop_area_resize[0])
- g2l_feat5 = torch_roi_align(g2l_feat5, bbox, (12, 16), 12/384, aligned=True)
- x6 = self.conv5(torch.cat([x6, g2l_feat5], dim=1))
-
- x5 = self.up1(torch.cat([x6, guide_cat[0]], dim=1), x5)
- if self.g2l:
- g2l_feat4 = self.g2l4(coarse_feat_whole[1], crop_area_resize[1])
- g2l_feat4 = torch_roi_align(g2l_feat4, bbox, (24, 32), 24/384, aligned=True)
- x5 = self.conv4(torch.cat([x5, g2l_feat4], dim=1))
-
- x4 = self.up2(torch.cat([x5, guide_cat[1]], dim=1), x4)
- if self.g2l:
- g2l_feat3 = self.g2l3(coarse_feat_whole[2], crop_area_resize[2])
- g2l_feat3 = torch_roi_align(g2l_feat3, bbox, (48, 64), 48/384, aligned=True)
- x4 = self.conv3(torch.cat([x4, g2l_feat3], dim=1))
-
- x3 = self.up3(torch.cat([x4, guide_cat[2]], dim=1), x3)
- if self.g2l:
- g2l_feat2 = self.g2l2(coarse_feat_whole[3], crop_area_resize[3])
- g2l_feat2 = torch_roi_align(g2l_feat2, bbox, (96, 128), 96/384, aligned=True)
- x3 = self.conv2(torch.cat([x3, g2l_feat2], dim=1))
-
- x2 = self.up4(torch.cat([x3, guide_cat[3]], dim=1), x2)
- if self.g2l:
- g2l_feat1 = self.g2l1(coarse_feat_whole[4], crop_area_resize[4])
- g2l_feat1 = torch_roi_align(g2l_feat1, bbox, (192, 256), 192/384, aligned=True)
- x2 = self.conv1(torch.cat([x2, g2l_feat1], dim=1))
-
- x1 = self.up5(torch.cat([x2, guide_cat[4]], dim=1), x1)
- if self.g2l:
- g2l_feat0 = self.g2l0(coarse_feat_whole[5], crop_area_resize[5])
- g2l_feat0 = torch_roi_align(g2l_feat0, bbox, (384, 512), 384/384, aligned=True)
- x1 = self.conv0(torch.cat([x1, g2l_feat0], dim=1))
-
- output = [x1, x2, x3, x4, x5, x6]
- return output
\ No newline at end of file
diff --git a/zoedepth/models/layers/patch_transformer.py b/zoedepth/models/layers/patch_transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..99d9e51a06b981bae45ce7dd64eaef19a4121991
--- /dev/null
+++ b/zoedepth/models/layers/patch_transformer.py
@@ -0,0 +1,91 @@
+# MIT License
+
+# Copyright (c) 2022 Intelligent Systems Lab Org
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# File author: Shariq Farooq Bhat
+
+import torch
+import torch.nn as nn
+
+
+class PatchTransformerEncoder(nn.Module):
+ def __init__(self, in_channels, patch_size=10, embedding_dim=128, num_heads=4, use_class_token=False):
+ """ViT-like transformer block
+
+ Args:
+ in_channels (int): Input channels
+ patch_size (int, optional): patch size. Defaults to 10.
+ embedding_dim (int, optional): Embedding dimension in transformer model. Defaults to 128.
+ num_heads (int, optional): number of attention heads. Defaults to 4.
+ use_class_token (bool, optional): Whether to use extra token at the start for global accumulation (called as "class token"). Defaults to False.
+ """
+ super(PatchTransformerEncoder, self).__init__()
+ self.use_class_token = use_class_token
+ encoder_layers = nn.TransformerEncoderLayer(
+ embedding_dim, num_heads, dim_feedforward=1024)
+ self.transformer_encoder = nn.TransformerEncoder(
+ encoder_layers, num_layers=4) # takes shape S,N,E
+
+ self.embedding_convPxP = nn.Conv2d(in_channels, embedding_dim,
+ kernel_size=patch_size, stride=patch_size, padding=0)
+
+ def positional_encoding_1d(self, sequence_length, batch_size, embedding_dim, device='cpu'):
+ """Generate positional encodings
+
+ Args:
+ sequence_length (int): Sequence length
+ embedding_dim (int): Embedding dimension
+
+ Returns:
+ torch.Tensor SBE: Positional encodings
+ """
+ position = torch.arange(
+ 0, sequence_length, dtype=torch.float32, device=device).unsqueeze(1)
+ index = torch.arange(
+ 0, embedding_dim, 2, dtype=torch.float32, device=device).unsqueeze(0)
+ div_term = torch.exp(index * (-torch.log(torch.tensor(10000.0, device=device)) / embedding_dim))
+ pos_encoding = position * div_term
+ pos_encoding = torch.cat([torch.sin(pos_encoding), torch.cos(pos_encoding)], dim=1)
+ pos_encoding = pos_encoding.unsqueeze(1).repeat(1, batch_size, 1)
+ return pos_encoding
+
+
+ def forward(self, x):
+ """Forward pass
+
+ Args:
+ x (torch.Tensor - NCHW): Input feature tensor
+
+ Returns:
+ torch.Tensor - SNE: Transformer output embeddings. S - sequence length (=HW/patch_size^2), N - batch size, E - embedding dim
+ """
+ embeddings = self.embedding_convPxP(x).flatten(
+ 2) # .shape = n,c,s = n, embedding_dim, s
+ if self.use_class_token:
+ # extra special token at start ?
+ embeddings = nn.functional.pad(embeddings, (1, 0))
+
+ # change to S,N,E format required by transformer
+ embeddings = embeddings.permute(2, 0, 1)
+ S, N, E = embeddings.shape
+ embeddings = embeddings + self.positional_encoding_1d(S, N, E, device=embeddings.device)
+ x = self.transformer_encoder(embeddings) # .shape = S, N, E
+ return x
diff --git a/zoedepth/models/model_io.py b/zoedepth/models/model_io.py
index ff0e10ed7597dffa3123219d8be69795f2b0438b..78b6579631dd847ac76651238cb5a948b5a66286 100644
--- a/zoedepth/models/model_io.py
+++ b/zoedepth/models/model_io.py
@@ -46,7 +46,7 @@ def load_state_dict(model, state_dict):
state[k] = v
- model.load_state_dict(state, strict=False)
+ model.load_state_dict(state)
print("Loaded successfully")
return model
diff --git a/zoedepth/models/zoedepth/__init__.py b/zoedepth/models/zoedepth/__init__.py
index b972c255db238aabe44d955c55f5c8221dcffede..cc33f737d238766559f0e3a8def3c0b568f23b7f 100644
--- a/zoedepth/models/zoedepth/__init__.py
+++ b/zoedepth/models/zoedepth/__init__.py
@@ -22,7 +22,7 @@
# File author: Shariq Farooq Bhat
-from .zoedepth_v1 import ZoeDepth
+from .zoedepth_v1 import ZoeDepth
all_versions = {
"v1": ZoeDepth,
diff --git a/zoedepth/models/zoedepth/color_channel/config_zoedepth_rgb.json b/zoedepth/models/zoedepth/color_channel/config_zoedepth_rgb.json
deleted file mode 100644
index 61f6520afc4c034304fbe852b06c2b08001a22c4..0000000000000000000000000000000000000000
--- a/zoedepth/models/zoedepth/color_channel/config_zoedepth_rgb.json
+++ /dev/null
@@ -1,66 +0,0 @@
-{
- "model": {
- "name": "ZoeDepth",
- "version_name": "v1",
- "n_bins": 64,
- "bin_embedding_dim": 128,
- "bin_centers_type": "softplus",
- "n_attractors":[16, 8, 4, 1],
- "attractor_alpha": 1000,
- "attractor_gamma": 2,
- "attractor_kind" : "mean",
- "attractor_type" : "inv",
- "midas_model_type" : "DPT_BEiT_L_384",
- "min_temp": 0.0212,
- "max_temp": 50.0,
- "output_distribution": "logbinomial",
- "memory_efficient": true,
- "inverse_midas": false,
- "img_size": [384, 512],
- "sampled_training": false,
- "do_resize": false // do resize in dataloader to speed up
- },
-
- "train": {
- "use_rbg": true,
- "train_midas": true,
- "use_pretrained_midas": true,
- "trainer": "zoedepth",
- // "epochs": 5,
- // "epochs": 8,
- // "epochs": 12,
- "epochs": 16,
- "bs": 16,
- // "optim_kwargs": {"lr": 0.0001, "wd": 0.01},
- "optim_kwargs": {"lr": 0.000161, "wd": 0.01},
- "sched_kwargs": {"div_factor": 1, "final_div_factor": 10000, "pct_start": 0.5, "three_phase":false, "cycle_momentum": true},
- // "sched_kwargs": {"div_factor": 1, "final_div_factor": 10000, "pct_start": 0.7, "three_phase":false, "cycle_momentum": true},
- "same_lr": false,
- "w_si": 1,
- "w_domain": 0.2,
- "w_reg": 0,
- "w_grad": 0,
- "avoid_boundary": false,
- "random_crop": false,
- "input_width": 640,
- "input_height": 480,
- "midas_lr_factor": 1,
- "encoder_lr_factor":10,
- "pos_enc_lr_factor":10,
- "freeze_midas_bn": true
-
- },
-
- "infer":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt",
- "force_keep_ar": true
- },
-
- "eval":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt"
- }
-}
\ No newline at end of file
diff --git a/zoedepth/models/zoedepth/config_zoedepth.json b/zoedepth/models/zoedepth/config_zoedepth.json
index 70b77dd0dd419eff1a54209bc7786bedc945fa1c..99beb2dcd886006ba87805bddbe408b6d5fdff78 100644
--- a/zoedepth/models/zoedepth/config_zoedepth.json
+++ b/zoedepth/models/zoedepth/config_zoedepth.json
@@ -16,15 +16,11 @@
"output_distribution": "logbinomial",
"memory_efficient": true,
"inverse_midas": false,
- "img_size": [384, 512],
- "sampled_training": false,
- // "do_resize": false // do resize in dataloader to speed up
- "do_resize": true
+ "img_size": [384, 512]
},
"train": {
"train_midas": true,
- "use_rgb": true,
"use_pretrained_midas": true,
"trainer": "zoedepth",
"epochs": 5,
diff --git a/zoedepth/models/zoedepth/config_zoedepth_fine.json b/zoedepth/models/zoedepth/config_zoedepth_fine.json
deleted file mode 100644
index ced6717389d0394b6ae869dcb84d00f8353566f4..0000000000000000000000000000000000000000
--- a/zoedepth/models/zoedepth/config_zoedepth_fine.json
+++ /dev/null
@@ -1,66 +0,0 @@
-{
- "model": {
- "name": "ZoeDepth",
- "version_name": "v1",
- "n_bins": 64,
- "bin_embedding_dim": 128,
- "bin_centers_type": "softplus",
- "n_attractors":[16, 8, 4, 1],
- "attractor_alpha": 1000,
- "attractor_gamma": 2,
- "attractor_kind" : "mean",
- "attractor_type" : "inv",
- "midas_model_type" : "DPT_BEiT_L_384",
- "min_temp": 0.0212,
- "max_temp": 50.0,
- "output_distribution": "logbinomial",
- "memory_efficient": true,
- "inverse_midas": false,
- "img_size": [384, 512],
- "sampled_training": false,
- // "do_resize": false // do resize in dataloader to speed up
- "do_resize": false,
-
- "do_normalize": true,
- "do_input_resize": true
- },
-
- "train": {
- "train_midas": true,
- "use_rgb": true,
- "use_pretrained_midas": true,
- "trainer": "zoedepth_custom",
- "epochs": 5,
- "bs": 16,
- "optim_kwargs": {"lr": 0.000161, "wd": 0.01},
- "sched_kwargs": {"div_factor": 1, "final_div_factor": 10000, "pct_start": 0.7, "three_phase":false, "cycle_momentum": true},
- "same_lr": false,
- "w_si": 1,
- "w_domain": 0.2,
- "w_reg": 0,
- "w_grad": 0,
- "avoid_boundary": false,
- "random_crop": true,
- "input_width": 160,
- "input_height": 120,
- "midas_lr_factor": 1,
- "encoder_lr_factor":10,
- "pos_enc_lr_factor":10,
- "freeze_midas_bn": true,
- "sec_stage": true
-
- },
-
- "infer":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt",
- "force_keep_ar": true
- },
-
- "eval":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt"
- }
-}
\ No newline at end of file
diff --git a/zoedepth/models/zoedepth/zoedepth_v1.py b/zoedepth/models/zoedepth/zoedepth_v1.py
index f57d00c90e96a189b3a9da5d63613ba1ce315cd0..091e2a2f0e1c11288c523d30e20b3923b47ff85c 100644
--- a/zoedepth/models/zoedepth/zoedepth_v1.py
+++ b/zoedepth/models/zoedepth/zoedepth_v1.py
@@ -28,16 +28,17 @@ import torch
import torch.nn as nn
from zoedepth.models.depth_model import DepthModel
from zoedepth.models.base_models.midas import MidasCore
+from zoedepth.models.base_models.depth_anything import DepthAnythingCore
from zoedepth.models.layers.attractor import AttractorLayer, AttractorLayerUnnormed
from zoedepth.models.layers.dist_layers import ConditionalLogBinomial
from zoedepth.models.layers.localbins_layers import (Projector, SeedBinRegressor,
SeedBinRegressorUnnormed)
from zoedepth.models.model_io import load_state_from_resource
-
+from mmengine import print_log
class ZoeDepth(DepthModel):
def __init__(self, core, n_bins=64, bin_centers_type="softplus", bin_embedding_dim=128, min_depth=1e-3, max_depth=10,
- n_attractors=[16, 8, 4, 1], attractor_alpha=300, attractor_gamma=2, attractor_kind='sum', attractor_type='exp', min_temp=5, max_temp=50, train_midas=True,
+ n_attractors=[16, 8, 4, 1], attractor_alpha=300, attractor_gamma=2, attractor_kind='sum', attractor_type='exp', min_temp=5, max_temp=50, train_midas=True, depth_anything=False,
midas_lr_factor=10, encoder_lr_factor=10, pos_enc_lr_factor=10, inverse_midas=False, **kwargs):
"""ZoeDepth model. This is the version of ZoeDepth that has a single metric head
@@ -63,6 +64,7 @@ class ZoeDepth(DepthModel):
"""
super().__init__()
+ self.depth_anything = depth_anything
self.core = core
self.max_depth = max_depth
self.min_depth = min_depth
@@ -82,7 +84,6 @@ class ZoeDepth(DepthModel):
N_MIDAS_OUT = 32
btlnck_features = self.core.output_channels[0]
num_out_features = self.core.output_channels[1:]
-
self.conv2 = nn.Conv2d(btlnck_features, btlnck_features,
kernel_size=1, stride=1, padding=0) # btlnck conv
@@ -120,10 +121,8 @@ class ZoeDepth(DepthModel):
# use log binomial instead of softmax
self.conditional_log_binomial = ConditionalLogBinomial(
last_in, bin_embedding_dim, n_classes=n_bins, min_temp=min_temp, max_temp=max_temp)
-
- self.collect_feat = {}
-
- def forward(self, x, return_final_centers=False, denorm=False, return_probs=False, **kwargs):
+
+ def forward(self, x, return_final_centers=False, denorm=False, return_probs=False, hack_feature=None, only_encoder=False, **kwargs):
"""
Args:
x (torch.Tensor): Input image tensor of shape (B, C, H, W)
@@ -139,18 +138,40 @@ class ZoeDepth(DepthModel):
- probs (torch.Tensor): Output probability distribution of shape (B, n_bins, H, W). Present only if return_probs is True
"""
- b, c, h, w = x.shape
- # print("input shape ", x.shape)
- self.orig_input_width = w
- self.orig_input_height = h
- rel_depth, out = self.core(x, denorm=denorm, return_rel_depth=True)
- # print("output shapes", rel_depth.shape, out.shape)
-
- outconv_activation = out[0]
- btlnck = out[1]
- x_blocks = out[2:]
-
+ temp_features = {}
+
+ if hack_feature is None:
+ b, c, h, w = x.shape
+ # print("input shape ", x.shape)
+ self.orig_input_width = w
+ self.orig_input_height = h
+ rel_depth, out = self.core(x, denorm=denorm, return_rel_depth=True)
+
+ outconv_activation = out[0]
+ btlnck = out[1]
+ x_blocks = out[2:]
+
+ # prepare encoder features
+ encoder_feat_list = [btlnck]
+ for i in x_blocks:
+ encoder_feat_list.append(i)
+ encoder_feat_list.append(outconv_activation)
+
+ if only_encoder:
+ # only use zoe encoder in this case
+ return encoder_feat_list
+
+ else:
+ rel_depth, out = hack_feature[0], hack_feature[1] # forward encoder outputs as inputs in zoedepth
+
+ outconv_activation = out[-1]
+ btlnck = out[0]
+ x_blocks = out[1:-1]
+
+
+
x_d0 = self.conv2(btlnck)
+ temp_features['x_d0'] = x_d0
x = x_d0
_, seed_b_centers = self.seed_bin_regressor(x)
@@ -163,14 +184,17 @@ class ZoeDepth(DepthModel):
prev_b_embedding = self.seed_projector(x)
# unroll this loop for better performance
- for projector, attractor, x in zip(self.projectors, self.attractors, x_blocks):
+ # for projector, attractor, x in zip(self.projectors, self.attractors, x_blocks):
+ for idx, (projector, attractor, x) in enumerate(zip(self.projectors, self.attractors, x_blocks)):
b_embedding = projector(x)
+ temp_features['x_blocks_feat_{}'.format(idx)] = x
b, b_centers = attractor(
b_embedding, b_prev, prev_b_embedding, interpolate=True)
b_prev = b.clone()
prev_b_embedding = b_embedding.clone()
last = outconv_activation
+ temp_features['midas_final_feat'] = last
if self.inverse_midas:
# invert depth followed by normalization
@@ -181,6 +205,7 @@ class ZoeDepth(DepthModel):
rel_cond = rel_depth.unsqueeze(1)
rel_cond = nn.functional.interpolate(
rel_cond, size=last.shape[2:], mode='bilinear', align_corners=True)
+ temp_features['last'] = last
last = torch.cat([last, rel_cond], dim=1)
b_embedding = nn.functional.interpolate(
@@ -197,11 +222,14 @@ class ZoeDepth(DepthModel):
output = dict(metric_depth=out)
if return_final_centers or return_probs:
output['bin_centers'] = b_centers
+ output['b_embedding'] = b_embedding
if return_probs:
output['probs'] = x
- output['coarse_depth_pred'] = out
+ output['temp_features'] = temp_features
+ # output['encoder_feats'] = encoder_feat_list
+
return output
def get_lr_params(self, lr):
@@ -212,39 +240,73 @@ class ZoeDepth(DepthModel):
Returns:
list : list of parameters to optimize and their learning rates, in the format required by torch optimizers.
"""
- param_conf = []
- if self.train_midas:
- if self.encoder_lr_factor > 0:
- param_conf.append({'params': self.core.get_enc_params_except_rel_pos(
- ), 'lr': lr / self.encoder_lr_factor})
-
- if self.pos_enc_lr_factor > 0:
+ if self.depth_anything:
+ param_conf = []
+ if self.train_midas:
+ if self.encoder_lr_factor > 0:
+ param_conf.append({'params': self.core.get_enc_params_except_rel_pos(
+ ), 'lr': lr / self.encoder_lr_factor})
+
+ if self.pos_enc_lr_factor > 0:
+ param_conf.append(
+ {'params': self.core.get_rel_pos_params(), 'lr': lr / self.pos_enc_lr_factor})
+
+ # midas_params = self.core.core.scratch.parameters()
+ midas_params = self.core.core.depth_head.parameters()
+ midas_lr_factor = self.midas_lr_factor
param_conf.append(
- {'params': self.core.get_rel_pos_params(), 'lr': lr / self.pos_enc_lr_factor})
+ {'params': midas_params, 'lr': lr / midas_lr_factor})
- midas_params = self.core.core.scratch.parameters()
- midas_lr_factor = self.midas_lr_factor
- param_conf.append(
- {'params': midas_params, 'lr': lr / midas_lr_factor})
+ remaining_modules = []
+ for name, child in self.named_children():
+ if name != 'core':
+ remaining_modules.append(child)
+ remaining_params = itertools.chain(
+ *[child.parameters() for child in remaining_modules])
- remaining_modules = []
- for name, child in self.named_children():
- if name != 'core':
- remaining_modules.append(child)
- remaining_params = itertools.chain(
- *[child.parameters() for child in remaining_modules])
+ param_conf.append({'params': remaining_params, 'lr': lr})
+
+ else:
+ param_conf = []
+ if self.train_midas:
+ if self.encoder_lr_factor > 0:
+ param_conf.append({'params': self.core.get_enc_params_except_rel_pos(
+ ), 'lr': lr / self.encoder_lr_factor})
+
+ if self.pos_enc_lr_factor > 0:
+ param_conf.append(
+ {'params': self.core.get_rel_pos_params(), 'lr': lr / self.pos_enc_lr_factor})
+
+ midas_params = self.core.core.scratch.parameters()
+ midas_lr_factor = self.midas_lr_factor
+ param_conf.append(
+ {'params': midas_params, 'lr': lr / midas_lr_factor})
- param_conf.append({'params': remaining_params, 'lr': lr})
+ remaining_modules = []
+ for name, child in self.named_children():
+ if name != 'core':
+ remaining_modules.append(child)
+ remaining_params = itertools.chain(
+ *[child.parameters() for child in remaining_modules])
+
+ param_conf.append({'params': remaining_params, 'lr': lr})
return param_conf
@staticmethod
def build(midas_model_type="DPT_BEiT_L_384", pretrained_resource=None, use_pretrained_midas=False, train_midas=False, freeze_midas_bn=True, **kwargs):
- core = MidasCore.build(midas_model_type=midas_model_type, use_pretrained_midas=use_pretrained_midas,
+ # core = MidasCore.build(midas_model_type=midas_model_type, use_pretrained_midas=use_pretrained_midas,
+ # train_midas=train_midas, fetch_features=True, freeze_bn=freeze_midas_bn, **kwargs)
+ if midas_model_type == 'vits' or midas_model_type == 'vitb' or midas_model_type == 'vitl':
+ core = DepthAnythingCore.build(midas_model_type=midas_model_type, use_pretrained_midas=use_pretrained_midas,
train_midas=train_midas, fetch_features=True, freeze_bn=freeze_midas_bn, **kwargs)
+ else:
+ core = MidasCore.build(midas_model_type=midas_model_type, use_pretrained_midas=use_pretrained_midas,
+ train_midas=train_midas, fetch_features=True, freeze_bn=freeze_midas_bn, **kwargs)
model = ZoeDepth(core, **kwargs)
if pretrained_resource:
assert isinstance(pretrained_resource, str), "pretrained_resource must be a string"
+ print_log("Loading deepnet from {}".format(pretrained_resource), logger='current')
model = load_state_from_resource(model, pretrained_resource)
return model
diff --git a/zoedepth/models/zoedepth_custom/configs/config_zoedepth_custom_coarse_stage1.json b/zoedepth/models/zoedepth_custom/configs/config_zoedepth_custom_coarse_stage1.json
deleted file mode 100644
index e93da621ac468a6240bb056ef4635c419c46206a..0000000000000000000000000000000000000000
--- a/zoedepth/models/zoedepth_custom/configs/config_zoedepth_custom_coarse_stage1.json
+++ /dev/null
@@ -1,72 +0,0 @@
-{
- "model": {
- "name": "ZoeDepthCustom",
- "version_name": "custom",
- "n_bins": 64,
- "bin_embedding_dim": 128,
- "bin_centers_type": "softplus",
- "n_attractors":[16, 8, 4, 1],
- "attractor_alpha": 1000,
- "attractor_gamma": 2,
- "attractor_kind" : "mean",
- "attractor_type" : "inv",
- "midas_model_type" : "DPT_BEiT_L_384",
- "min_temp": 0.0212,
- "max_temp": 50.0,
- "output_distribution": "logbinomial",
- "memory_efficient": true,
- "inverse_midas": false,
- "img_size": [384, 512],
- "do_resize": false, // do resize in dataloader to speed up
- "raw_depth_shape": [2160, 3840],
- "sr_ratio": 1,
- "sampled_training": false,
- "transform_sample_gt_size": [2160, 3840],
- "sample_feat_level": 4,
- "use_hr": false,
- "baseline": true
- },
-
- "train": {
- "use_rgb": true,
- "train_midas": true,
- "use_pretrained_midas": true,
- "trainer": "zoedepth_custom",
- // "epochs": 5,
- // "epochs": 12,
- "epochs": 16,
- // "epochs": 20,
- "bs": 16,
- // "optim_kwargs": {"lr": 0.0001, "wd": 0.01},
- "optim_kwargs": {"lr": 0.000161, "wd": 0.01},
- "sched_kwargs": {"div_factor": 1, "final_div_factor": 10000, "pct_start": 0.5, "three_phase":false, "cycle_momentum": true},
- // "sched_kwargs": {"div_factor": 1, "final_div_factor": 10000, "pct_start": 0.7, "three_phase":false, "cycle_momentum": true},
- "same_lr": false,
- "w_si": 1,
- "w_domain": 0.2,
- "w_reg": 0,
- "w_grad": 0,
- "avoid_boundary": false,
- "random_crop": false,
- "input_width": 640,
- "input_height": 480,
- "midas_lr_factor": 1,
- "encoder_lr_factor":10,
- "pos_enc_lr_factor":10,
- "freeze_midas_bn": true,
- },
-
- "infer":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt",
- "force_keep_ar": true
- },
-
- "eval":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt"
- }
-}
-
diff --git a/zoedepth/models/zoedepth_custom/configs/config_zoedepth_custom_fine_stage2.json b/zoedepth/models/zoedepth_custom/configs/config_zoedepth_custom_fine_stage2.json
deleted file mode 100644
index caa79553f8f35d60b42d72253a97f8079c90a7e0..0000000000000000000000000000000000000000
--- a/zoedepth/models/zoedepth_custom/configs/config_zoedepth_custom_fine_stage2.json
+++ /dev/null
@@ -1,74 +0,0 @@
-{
- "model": {
- "name": "ZoeDepthCustom",
- "version_name": "custom",
- "n_bins": 64,
- "bin_embedding_dim": 128,
- "bin_centers_type": "softplus",
- "n_attractors":[16, 8, 4, 1],
- "attractor_alpha": 1000,
- "attractor_gamma": 2,
- "attractor_kind" : "mean",
- "attractor_type" : "inv",
- "midas_model_type" : "DPT_BEiT_L_384",
- "min_temp": 0.0212,
- "max_temp": 50.0,
- "output_distribution": "logbinomial",
- "memory_efficient": true,
- "inverse_midas": false,
- "img_size": [384, 512],
- "do_resize": false, // do resize in dataloader to speed up
- "raw_depth_shape": [2160, 3840],
- "sr_ratio": 1,
- "sampled_training": false,
- "transform_sample_gt_size": [2160, 3840],
- "sample_feat_level": 4,
- "use_hr": false,
- "baseline": true
- },
-
- "train": {
- "use_rgb": true,
- "use_blur": false,
- "train_midas": true,
- "use_pretrained_midas": true,
- "trainer": "zoedepth_custom",
- // "epochs": 5,
- // "epochs": 12,
- "epochs": 16,
- // "epochs": 20,
- "bs": 16,
- // "optim_kwargs": {"lr": 0.0001, "wd": 0.01},
- "optim_kwargs": {"lr": 0.000161, "wd": 0.01},
- "sched_kwargs": {"div_factor": 1, "final_div_factor": 10000, "pct_start": 0.5, "three_phase":false, "cycle_momentum": true},
- // "sched_kwargs": {"div_factor": 1, "final_div_factor": 10000, "pct_start": 0.7, "three_phase":false, "cycle_momentum": true},
- "same_lr": false,
- "w_si": 1,
- "w_domain": 0.2,
- "w_reg": 0,
- "w_grad": 0,
- "avoid_boundary": false,
- "random_crop": true,
- "input_width": 640,
- "input_height": 480,
- "midas_lr_factor": 1,
- "encoder_lr_factor":10,
- "pos_enc_lr_factor":10,
- "freeze_midas_bn": true,
- "sec_stage": true
- },
-
- "infer":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt",
- "force_keep_ar": true
- },
-
- "eval":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt"
- }
-}
-
diff --git a/zoedepth/models/zoedepth_custom/configs/config_zoedepth_patchfusion.json b/zoedepth/models/zoedepth_custom/configs/config_zoedepth_patchfusion.json
deleted file mode 100644
index 33d5bd96012540496a080825aed76675fbb6d3b1..0000000000000000000000000000000000000000
--- a/zoedepth/models/zoedepth_custom/configs/config_zoedepth_patchfusion.json
+++ /dev/null
@@ -1,75 +0,0 @@
-{
- "model": {
- "name": "PatchFusion",
- "version_name": "patchfusion",
- "n_bins": 64,
- "bin_embedding_dim": 128,
- "bin_centers_type": "softplus",
- "n_attractors":[16, 8, 4, 1],
- "attractor_alpha": 1000,
- "attractor_gamma": 2,
- "attractor_kind" : "mean",
- "attractor_type" : "inv",
- "midas_model_type" : "DPT_BEiT_L_384",
- "min_temp": 0.0212,
- "max_temp": 50.0,
- "output_distribution": "logbinomial",
- "memory_efficient": true,
- "inverse_midas": false,
- "img_size": [384, 512],
- "do_resize": false, // do resize in dataloader to speed up
- "raw_depth_shape": [2160, 3840],
- "sr_ratio": 1,
- "sampled_training": false,
- "transform_sample_gt_size": [2160, 3840],
- "sample_feat_level": 4,
- "use_hr": false,
- "baseline": true,
- "condition": true,
- "freeze": true,
- "g2l": true,
- "use_fusion_network": true,
- "use_area_prior": true
- },
-
- "train": {
- "use_rgb": true,
- "use_blur": false,
- "train_midas": true,
- "use_pretrained_midas": true,
- "trainer": "zoedepth_custom",
- "epochs": 12,
- "bs": 16,
- "optim_kwargs": {"lr": 0.0001, "wd": 0.01},
- "sched_kwargs": {"div_factor": 10, "final_div_factor": 10000, "pct_start": 0.5, "three_phase":false, "cycle_momentum": true},
- "same_lr": false,
- "w_si": 1,
- "w_domain": 0.2,
- "w_reg": 0,
- "w_grad": 0,
- "avoid_boundary": false,
- "random_crop": true,
- "input_width": 640,
- "input_height": 480,
- "midas_lr_factor": 1,
- "encoder_lr_factor":10,
- "pos_enc_lr_factor":10,
- "freeze_midas_bn": true,
- "sec_stage": true,
- "multi_consistency": true
- },
-
- "infer":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt",
- "force_keep_ar": true
- },
-
- "eval":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt"
- }
-}
-
diff --git a/zoedepth/models/zoedepth_custom/configs/config_zoedepth_patchfusion_finetune.json b/zoedepth/models/zoedepth_custom/configs/config_zoedepth_patchfusion_finetune.json
deleted file mode 100644
index d1e574488daf4a6b9ac004c77766e2c1de99afbc..0000000000000000000000000000000000000000
--- a/zoedepth/models/zoedepth_custom/configs/config_zoedepth_patchfusion_finetune.json
+++ /dev/null
@@ -1,82 +0,0 @@
-{
- "model": {
- "name": "PatchFusion",
- "version_name": "patchfusion",
- "n_bins": 64,
- "bin_embedding_dim": 128,
- "bin_centers_type": "softplus",
- "n_attractors":[16, 8, 4, 1],
- "attractor_alpha": 1000,
- "attractor_gamma": 2,
- "attractor_kind" : "mean",
- "attractor_type" : "inv",
- "midas_model_type" : "DPT_BEiT_L_384",
- "min_temp": 0.0212,
- "max_temp": 50.0,
- "output_distribution": "logbinomial",
- "memory_efficient": true,
- "inverse_midas": false,
- "img_size": [384, 512],
- "do_resize": false, // do resize in dataloader to speed up
- "raw_depth_shape": [2160, 3840], // 540, 960
- "sr_ratio": 1,
- "sampled_training": false,
- "transform_sample_gt_size": [2160, 3840],
- "sample_feat_level": 4,
- "use_hr": false,
- "baseline": true,
- "condition": true,
- "freeze": true,
- "g2l": true,
- "use_fusion_network": true,
- "use_area_prior": true,
- "consistency_training": true,
- // "consistency_target": "final_feat",
- "consistency_target": "mix",
- },
-
- "train": {
- "use_rgb": true,
- "use_blur": false,
- "train_midas": true,
- "use_pretrained_midas": true,
- "trainer": "zoedepth_custom",
- "epochs": 12,
- "bs": 16,
- "optim_kwargs": {"lr": 0.0001, "wd": 0.01},
- "sched_kwargs": {"div_factor": 10, "final_div_factor": 10000, "pct_start": 0.5, "three_phase":false, "cycle_momentum": true},
- "same_lr": false,
- "w_si": 1,
- "w_domain": 0.2,
- "w_reg": 0,
- "w_grad": 0,
- "avoid_boundary": false,
- "random_crop": true,
- "input_width": 640,
- "input_height": 480,
- "midas_lr_factor": 1,
- "encoder_lr_factor":10,
- "pos_enc_lr_factor":10,
- "freeze_midas_bn": true,
- "sec_stage": true,
- "multi_consistency": true,
- "w_consistency": 0.1, // consistency weight here
- "overlap_length_h": 135,
- "overlap_length_w": 240,
- "w_p": 0.1 // weight of pred depth in consistency loss
- },
-
- "infer":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt",
- "force_keep_ar": true
- },
-
- "eval":{
- "train_midas": false,
- "use_pretrained_midas": false,
- "pretrained_resource" : "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt"
- }
-}
-
diff --git a/zoedepth/models/zoedepth_custom/patchfusion.py b/zoedepth/models/zoedepth_custom/patchfusion.py
deleted file mode 100644
index 75b25b8b50f5fef9f968fb697f914e46b0a9c502..0000000000000000000000000000000000000000
--- a/zoedepth/models/zoedepth_custom/patchfusion.py
+++ /dev/null
@@ -1,610 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Zhenyu Li
-
-import itertools
-
-import math
-import copy
-import torch
-import torch.nn as nn
-import numpy as np
-
-from zoedepth.models.depth_model import DepthModel
-from zoedepth.models.base_models.midas import MidasCore
-from zoedepth.models.layers.attractor import AttractorLayer, AttractorLayerUnnormed
-from zoedepth.models.layers.dist_layers import ConditionalLogBinomial, ConditionalLogBinomialV2
-from zoedepth.models.layers.localbins_layers import (Projector, SeedBinRegressor, SeedBinRegressorUnnormed)
-from zoedepth.models.model_io import load_state_from_resource
-from torchvision.transforms import Normalize
-from torchvision.ops import roi_align as torch_roi_align
-from zoedepth.utils.misc import generatemask
-
-from zoedepth.models.layers.transformer import TransformerDecoderLayer, TransformerEncoderLayer, TransformerEncoder
-
-from zoedepth.utils.misc import colorize, colors
-import matplotlib.pyplot as plt
-
-from zoedepth.models.layers.fusion_network import UNetv1
-import matplotlib.pyplot as plt
-
-import os
-import torch.distributed as dist
-import torch.nn.functional as F
-
-def check_keywords_in_name(name, keywords=()):
- isin = False
- for keyword in keywords:
- if keyword in name:
- isin = True
- return isin
-
-def get_activation(name, bank):
- # input of forward_hook will be a function of model/inp/oup
- def hook(module, input, output):
- bank[name] = output
- return hook
-
-def get_input(name, bank):
- # input of forward_hook will be a function of model/inp/oup
- def hook(module, input, output):
- bank[name] = input
- return hook
-
-class AttributeDict(dict):
- def __getattr__(self, key):
- try:
- return self[key]
- except KeyError:
- raise AttributeError(key)
-
- def __setattr__(self, key, value):
- self[key] = value
-
- def __delattr__(self, key):
- try:
- del self[key]
- except KeyError:
- raise AttributeError(key)
-
-class PatchFusion(DepthModel):
- def __init__(self, coarse_model, fine_model, n_bins=64, bin_centers_type="softplus", bin_embedding_dim=128, min_depth=1e-3, max_depth=10,
- n_attractors=[16, 8, 4, 1], attractor_alpha=300, attractor_gamma=2, attractor_kind='sum', attractor_type='exp', min_temp=5, max_temp=50, train_midas=True,
- midas_lr_factor=10, encoder_lr_factor=10, pos_enc_lr_factor=10, inverse_midas=False, sr_ratio=1,
- raw_depth_shape=(2160, 3840), transform_sample_gt_size=(2160, 3840), representation='',
- fetch_features=True, sample_feat_level=3, use_hr=False, deform=False, wo_bins=False, baseline=False,
- condition=True, freeze=False, g2l=False, use_fusion_network=False, use_area_prior=False,
- unet_version='v1', consistency_training=False, consistency_target='unet_feat', pos_embed=False, **kwargs):
- """ZoeDepth model. This is the version of ZoeDepth that has a single metric head
-
- Args:
- core (models.base_models.midas.MidasCore): The base midas model that is used for extraction of "relative" features
- n_bins (int, optional): Number of bin centers. Defaults to 64.
- bin_centers_type (str, optional): "normed" or "softplus". Activation type used for bin centers. For "normed" bin centers, linear normalization trick is applied. This results in bounded bin centers.
- For "softplus", softplus activation is used and thus are unbounded. Defaults to "softplus".
- bin_embedding_dim (int, optional): bin embedding dimension. Defaults to 128.
- min_depth (float, optional): Lower bound for normed bin centers. Defaults to 1e-3.
- max_depth (float, optional): Upper bound for normed bin centers. Defaults to 10.
- n_attractors (List[int], optional): Number of bin attractors at decoder layers. Defaults to [16, 8, 4, 1].
- attractor_alpha (int, optional): Proportional attractor strength. Refer to models.layers.attractor for more details. Defaults to 300.
- attractor_gamma (int, optional): Exponential attractor strength. Refer to models.layers.attractor for more details. Defaults to 2.
- attractor_kind (str, optional): Attraction aggregation "sum" or "mean". Defaults to 'sum'.
- attractor_type (str, optional): Type of attractor to use; "inv" (Inverse attractor) or "exp" (Exponential attractor). Defaults to 'exp'.
- min_temp (int, optional): Lower bound for temperature of output probability distribution. Defaults to 5.
- max_temp (int, optional): Upper bound for temperature of output probability distribution. Defaults to 50.
- train_midas (bool, optional): Whether to train "core", the base midas model. Defaults to True.
- midas_lr_factor (int, optional): Learning rate reduction factor for base midas model except its encoder and positional encodings. Defaults to 10.
- encoder_lr_factor (int, optional): Learning rate reduction factor for the encoder in midas model. Defaults to 10.
- pos_enc_lr_factor (int, optional): Learning rate reduction factor for positional encodings in the base midas model. Defaults to 10.
-
- sr_ratio: sr ratio during infer
- raw_depth_shape: raw depth shape during infer. times sr_ratio will be the target resolution. Used to sample points during training
- transform_sample_gt_size: training depth shape # influenced by crop shape which is not included in this pipeline right now
- representation: I use it to test the "bilap head" and a discarded idea
- fetch_features: if fetch feats. Default=True
- """
- super().__init__()
-
- self.coarse_model = coarse_model
- self.fine_model = fine_model
- self.max_depth = max_depth
- self.min_depth = min_depth
- self.min_temp = min_temp
- self.bin_centers_type = bin_centers_type
-
- self.midas_lr_factor = midas_lr_factor
- self.encoder_lr_factor = encoder_lr_factor
- self.pos_enc_lr_factor = pos_enc_lr_factor
- self.train_midas = train_midas
- self.inverse_midas = inverse_midas
-
- if bin_centers_type == "normed":
- SeedBinRegressorLayer = SeedBinRegressor
- Attractor = AttractorLayer
- elif bin_centers_type == "softplus": # default
- SeedBinRegressorLayer = SeedBinRegressorUnnormed
- Attractor = AttractorLayerUnnormed
- elif bin_centers_type == "hybrid1":
- SeedBinRegressorLayer = SeedBinRegressor
- Attractor = AttractorLayerUnnormed
- elif bin_centers_type == "hybrid2":
- SeedBinRegressorLayer = SeedBinRegressorUnnormed
- Attractor = AttractorLayer
- else:
- raise ValueError(
- "bin_centers_type should be one of 'normed', 'softplus', 'hybrid1', 'hybrid2'")
-
- N_MIDAS_OUT = 32
- btlnck_features = self.fine_model.core.output_channels[0]
- num_out_features = self.fine_model.core.output_channels[1:] # all of them are the same
-
- self.seed_bin_regressor = SeedBinRegressorLayer(
- btlnck_features, n_bins=n_bins, min_depth=min_depth, max_depth=max_depth)
- self.seed_projector = Projector(btlnck_features, bin_embedding_dim)
- self.projectors = nn.ModuleList([
- Projector(num_out, bin_embedding_dim)
- for num_out in num_out_features
- ])
- # 1000, 2, inv, mean
- self.attractors = nn.ModuleList([
- Attractor(bin_embedding_dim, n_bins, n_attractors=n_attractors[i], min_depth=min_depth, max_depth=max_depth,
- alpha=attractor_alpha, gamma=attractor_gamma, kind=attractor_kind, attractor_type=attractor_type)
- for i in range(len(num_out_features))
- ])
-
- last_in = N_MIDAS_OUT + 1 # +1 for relative depth
-
- # use log binomial instead of softmax
- self.conditional_log_binomial = ConditionalLogBinomial(
- last_in, bin_embedding_dim, n_classes=n_bins, min_temp=min_temp, max_temp=max_temp)
-
- self.handles = []
- self.hook_feats = {}
- self.set_fetch_features(fetch_features)
-
- # settings for patchfusion
- self.use_area_prior = use_area_prior
- self.g2l = g2l
-
- self.fusion_conv_list = nn.ModuleList()
- for i in range(6):
- if i == 5:
- layer = nn.Conv2d(N_MIDAS_OUT * 2, N_MIDAS_OUT, 3, 1, 1)
- else:
- layer = nn.Conv2d(btlnck_features * 2, btlnck_features, 3, 1, 1)
- self.fusion_conv_list.append(layer)
-
- self.coarse_input_proj = nn.ModuleList()
- for i in range(6):
- if i == 4:
- layer = nn.Conv2d(N_MIDAS_OUT, N_MIDAS_OUT, 3, 1, 1)
- else:
- layer = nn.Conv2d(btlnck_features, btlnck_features, 3, 1, 1)
- self.coarse_input_proj.append(layer)
-
- self.fine_input_proj = nn.ModuleList()
- for i in range(6):
- if i == 4:
- layer = nn.Conv2d(N_MIDAS_OUT, N_MIDAS_OUT, 3, 1, 1)
- else:
- layer = nn.Conv2d(btlnck_features, btlnck_features, 3, 1, 1)
- self.fine_input_proj.append(layer)
-
- self.coarse_depth_proj = nn.Conv2d(1, last_in, 3, 1, 1)
- self.fine_depth_proj = nn.Conv2d(1, last_in, 3, 1, 1)
- # self.coarse_depth_proj = nn.Conv2d(1, 1, 3, 1, 1)
- # self.fine_depth_proj = nn.Conv2d(1, 1, 3, 1, 1)
- # self.init_weight()
-
- self.freeze = freeze
- if self.freeze:
- # Freeze the parameters of sub_model
- for param in self.fine_model.parameters():
- param.requires_grad = False
-
- for param in self.coarse_model.parameters():
- param.requires_grad = False
-
- # Set the sub_model to evaluation mode
- self.fine_model.eval()
- self.coarse_model.eval()
-
- self.use_fusion_network = use_fusion_network
- if self.use_fusion_network:
- self.fusion_extractor = UNetv1(5, self.g2l, pos_embed, use_area_prior)
-
- self.consistency_training = consistency_training
- if self.consistency_training:
- if consistency_target == 'mix':
- consistency_target = 'unet_feat'
- self.consistency_target = consistency_target
- print("current consistency target is {}".format(consistency_target))
-
- self.consistency_projs = nn.ModuleList()
- if self.consistency_target == 'unet_feat':
- for i in range(6):
- if i == 5:
- layer = nn.Conv2d(N_MIDAS_OUT, N_MIDAS_OUT, 1, 1, 0)
- layer = nn.Identity()
- else:
- layer = nn.Conv2d(btlnck_features, btlnck_features, 1, 1, 0)
- layer = nn.Identity()
- self.consistency_projs.append(layer)
-
- if self.consistency_target == 'final_feat':
- layer = nn.Conv2d(64, 64, 1, 1, 0) # 192
- layer = nn.Identity()
- self.consistency_projs.append(layer)
- layer = nn.Conv2d(32, 32, 1, 1, 0) # 384
- layer = nn.Identity()
- self.consistency_projs.append(layer)
- layer = nn.Conv2d(128, 128, 1, 1, 0) # 192
- layer = nn.Identity()
- self.consistency_projs.append(layer)
-
- def init_weight(self):
- for m in self.coarse_input_proj:
- torch.nn.init.constant_(m.weight, 0)
- torch.nn.init.constant_(m.bias, 0)
-
- for m in self.fine_input_proj:
- torch.nn.init.constant_(m.weight, 0)
- torch.nn.init.constant_(m.bias, 0)
-
- torch.nn.init.constant_(self.coarse_depth_proj.weight, 0)
- torch.nn.init.constant_(self.fine_depth_proj.weight, 0)
-
- def get_lr_params(self, lr):
- """
- Learning rate configuration for different layers of the model
- Args:
- lr (float) : Base learning rate
- Returns:
- list : list of parameters to optimize and their learning rates, in the format required by torch optimizers.
- """
-
- # return self.fusion_extractor.parameters()
-
- param_conf = []
- param_conf_coarse_model = self.coarse_model.get_lr_params(lr)
- param_conf_fine_model = self.fine_model.get_lr_params(lr)
- param_conf.extend(param_conf_coarse_model)
- param_conf.extend(param_conf_fine_model)
-
-
- skip_list = {'absolute_pos_embed'}
- skip_keywords = {'relative_position_bias_table'}
- skip_hack = {'g2l0', 'g2l1', 'g2l2', 'g2l3', 'g2l4', 'g2l5'}
-
- no_decay = []
- has_decay = []
- fusion_enc = []
-
- for name, param in self.named_parameters():
- if 'coarse_model' not in name and 'fine_model' not in name:
- if len(param.shape) == 1 or (name.endswith(".bias") and check_keywords_in_name(name, skip_hack)) or check_keywords_in_name(name, skip_list) or \
- check_keywords_in_name(name, skip_keywords):
- # print("no decay: {}".format(name))
- no_decay.append(param)
- else:
- # print("has decay: {}".format(name))
- has_decay.append(param)
-
- param_conf.append({'params': has_decay, 'lr': lr})
- param_conf.append({'params': no_decay, 'weight_decay': 0., 'lr': lr})
- # param_conf.append({'params': fusion_enc, 'lr': lr / self.encoder_lr_factor})
- param_conf.append({'params': fusion_enc, 'lr': lr})
-
- return param_conf
-
- def forward(
- self,
- x,
- sampled_depth=None,
- mode='train',
- return_final_centers=False,
- denorm=False,
- return_probs=False,
- image_raw=None,
- bbox=None,
- crop_area=None,
- shift=None,
- bbox_raw=None,
- iter_prior=None,
- previous_info=None,
- **kwargs):
- """
- Args:
- x (torch.Tensor): Input image tensor of shape (B, C, H, W)
- return_final_centers (bool, optional): Whether to return the final bin centers. Defaults to False.
- denorm (bool, optional): Whether to denormalize the input image. This reverses ImageNet normalization as midas normalization is different. Defaults to False.
- return_probs (bool, optional): Whether to return the output probability distribution. Defaults to False.
-
- Returns:
- dict: Dictionary containing the following keys:
- - rel_depth (torch.Tensor): Relative depth map of shape (B, H, W)
- - metric_depth (torch.Tensor): Metric depth map of shape (B, 1, H, W)
- - bin_centers (torch.Tensor): Bin centers of shape (B, n_bins). Present only if return_final_centers is True
- - probs (torch.Tensor): Output probability distribution of shape (B, n_bins, H, W). Present only if return_probs is True
-
- """
-
- if self.consistency_training and mode == 'train':
- split_x = torch.split(x, 3, dim=1)
- x = torch.cat(split_x, dim=0)
- image_raw = torch.cat([image_raw, image_raw], dim=0)
- split_bbox = torch.split(bbox, 4, dim=-1)
- bbox = torch.cat(split_bbox, dim=0)
- split_bbox = torch.split(bbox_raw, 4, dim=-1)
- bbox_raw = torch.cat(split_bbox, dim=0)
- crop_area = torch.split(crop_area, 1, dim=1)
- crop_area = torch.cat(crop_area, dim=0)
-
- crop_input = x
-
- # coarse forward
- if self.freeze:
- with torch.no_grad():
- if self.fine_model.training:
- self.fine_model.eval()
- self.coarse_model.eval()
-
- if previous_info is None:
- previous_info = dict()
- whole_depth_pred = self.coarse_model(image_raw)['metric_depth']
- previous_info['whole_depth_pred'] = whole_depth_pred
- previous_info['coarse_model.hook_feats'] = self.coarse_model.hook_feats
-
- else:
- whole_depth_pred = previous_info['whole_depth_pred']
- self.coarse_model.hook_feats = dict()
- self.coarse_model.hook_feats = previous_info['coarse_model.hook_feats']
-
- fine_depth_pred = self.fine_model(x)['metric_depth']
-
- whole_depth_pred = nn.functional.interpolate(
- whole_depth_pred, (2160, 3840), mode='bilinear', align_corners=True)
-
- else:
- whole_depth_pred = self.coarse_model(image_raw)['metric_depth']
- fine_depth_pred = self.fine_model(x)['metric_depth']
-
- coarse_model_midas_enc_feats = [
- self.coarse_input_proj[5](self.coarse_model.hook_feats['x_d0']),
- self.coarse_input_proj[0](self.coarse_model.hook_feats['x_blocks_feat_0']),
- self.coarse_input_proj[1](self.coarse_model.hook_feats['x_blocks_feat_1']),
- self.coarse_input_proj[2](self.coarse_model.hook_feats['x_blocks_feat_2']),
- self.coarse_input_proj[3](self.coarse_model.hook_feats['x_blocks_feat_3']),
- self.coarse_input_proj[4](self.coarse_model.hook_feats['midas_final_feat'])] # 384
-
- if self.g2l:
- coarse_model_midas_enc_feats_g2l = coarse_model_midas_enc_feats
-
- if self.use_area_prior:
- crop_area_resize = [
- nn.functional.interpolate(crop_area, (12, 16), mode='bilinear', align_corners=True),
- nn.functional.interpolate(crop_area, (24, 32), mode='bilinear', align_corners=True),
- nn.functional.interpolate(crop_area, (48, 64), mode='bilinear', align_corners=True),
- nn.functional.interpolate(crop_area, (96, 128), mode='bilinear', align_corners=True),
- nn.functional.interpolate(crop_area, (192, 256), mode='bilinear', align_corners=True),
- nn.functional.interpolate(crop_area, (384, 512), mode='bilinear', align_corners=True)]
- else:
- crop_area_resize = None
-
- inds = torch.arange(bbox.shape[0]).to(bbox.device).unsqueeze(dim=-1)
- bbox = torch.cat((inds, bbox), dim=-1)
- coarse_model_midas_enc_roi_feats = [
- torch_roi_align(coarse_model_midas_enc_feats[0], bbox, (12, 16), 12/384, aligned=True),
- torch_roi_align(coarse_model_midas_enc_feats[1], bbox, (24, 32), 24/384, aligned=True),
- torch_roi_align(coarse_model_midas_enc_feats[2], bbox, (48, 64), 48/384, aligned=True),
- torch_roi_align(coarse_model_midas_enc_feats[3], bbox, (96, 128), 96/384, aligned=True),
- torch_roi_align(coarse_model_midas_enc_feats[4], bbox, (192, 256), 192/384, aligned=True),
- torch_roi_align(coarse_model_midas_enc_feats[5], bbox, (384, 512), 384/384, aligned=True)
- ]
-
- # whole_depth_roi_pred = torch_roi_align(whole_depth_pred, bbox, (384, 512), 384/384)
- # back to full resolution to avoid potential misalignment
- bbox_hack = copy.deepcopy(bbox)
- bbox_hack[:, 1] = bbox[:, 1] / 512 * 3840 # scale back to full resolution coord
- bbox_hack[:, 2] = bbox[:, 2] / 384 * 2160
- bbox_hack[:, 3] = bbox[:, 3] / 512 * 3840
- bbox_hack[:, 4] = bbox[:, 4] / 384 * 2160
- whole_depth_roi_pred = torch_roi_align(whole_depth_pred, bbox_hack, (384, 512), 1, aligned=True)
-
- fine_model_midas_enc_feats = [
- self.fine_input_proj[5](self.fine_model.hook_feats['x_d0']),
- self.fine_input_proj[0](self.fine_model.hook_feats['x_blocks_feat_0']),
- self.fine_input_proj[1](self.fine_model.hook_feats['x_blocks_feat_1']),
- self.fine_input_proj[2](self.fine_model.hook_feats['x_blocks_feat_2']),
- self.fine_input_proj[3](self.fine_model.hook_feats['x_blocks_feat_3']),
- self.fine_input_proj[4](self.fine_model.hook_feats['midas_final_feat'])] # 384
-
- x_plane = []
- x_blocks = []
- feat_plus_list = []
- feat_cat_list = []
- res_pool = [(24, 32), (48, 64), (96, 128), (192, 256), (384, 512)]
- for l_i, (f_ca, f_c_roi, f_f) in enumerate(zip(coarse_model_midas_enc_feats, coarse_model_midas_enc_roi_feats, fine_model_midas_enc_feats)):
- feat_cat = self.fusion_conv_list[l_i](torch.cat([f_c_roi, f_f], dim=1))
- feat_plus = f_c_roi + f_f
- feat_cat_list.append(feat_cat)
- feat_plus_list.append(feat_plus)
-
- if iter_prior is not None:
- input_tensor = torch.cat([whole_depth_roi_pred, iter_prior, crop_input], dim=1)
- else:
- input_tensor = torch.cat([whole_depth_roi_pred, fine_depth_pred, crop_input], dim=1)
- output = self.fusion_extractor(
- input_tensor = input_tensor,
- guide_plus = feat_plus_list,
- guide_cat = feat_cat_list,
- bbox = bbox,
- crop_area_resize = crop_area_resize,
- fine_feat_crop = fine_model_midas_enc_feats,
- coarse_feat_whole = coarse_model_midas_enc_feats,
- coarse_feat_crop = coarse_model_midas_enc_roi_feats,
- coarse_feat_whole_hack=None)[::-1] # low -> high
-
- x_blocks = output
- x = x_blocks[0]
- x_blocks = x_blocks[1:]
-
- if self.consistency_training:
- if self.consistency_target == 'unet_feat':
- proj_feat_list = []
- for idx, feat in enumerate(output):
- proj_feat = self.consistency_projs[idx](feat)
- proj_feat_list.append(proj_feat)
-
- # NOTE: below is ZoeDepth implementation
- # # new last
- # last = coarse_model_midas_enc_roi_feats[-1] + fine_model_midas_enc_feats[-1]
- last = x_blocks[-1] # have already been fused in x_blocks
- self.hook_feats['midas_final_feat'] = last
-
- bs, c, h, w = last.shape
- rel_cond = torch.zeros((bs, 1, h, w), device=last.device)
-
- self.hook_feats['rel_depth'] = rel_cond # skip this
-
- _, seed_b_centers = self.seed_bin_regressor(x)
-
- if self.bin_centers_type == 'normed' or self.bin_centers_type == 'hybrid2':
- b_prev = (seed_b_centers - self.min_depth) / \
- (self.max_depth - self.min_depth)
- else:
- b_prev = seed_b_centers
-
- prev_b_embedding = self.seed_projector(x)
-
- # unroll this loop for better performance
- for idx, (projector, attractor, x) in enumerate(zip(self.projectors, self.attractors, x_blocks)):
- b_embedding = projector(x)
- self.hook_feats['x_blocks_feat_{}'.format(idx)] = x
- b, b_centers = attractor(
- b_embedding, b_prev, prev_b_embedding, interpolate=True)
- b_prev = b.clone()
- prev_b_embedding = b_embedding.clone()
-
- self.hook_feats['b_centers'] = b_centers
-
-
- if self.consistency_training:
- if self.consistency_target == 'final_feat':
- proj_feat_1 = self.consistency_projs[0](b_centers)
- proj_feat_2 = self.consistency_projs[1](last)
- proj_feat_3 = self.consistency_projs[2](b_embedding)
- proj_feat_list = [proj_feat_1, proj_feat_2, proj_feat_3]
-
- rel_cond = nn.functional.interpolate(
- rel_cond, size=last.shape[2:], mode='bilinear', align_corners=True)
- last = torch.cat([last, rel_cond], dim=1) # + self.coarse_depth_proj(whole_depth_roi_pred) + self.fine_depth_proj(fine_depth_pred)
- b_embedding = nn.functional.interpolate(
- b_embedding, last.shape[-2:], mode='bilinear', align_corners=True)
- # till here, we have features (attached with a relative depth prediction) and embeddings
- # post process
- # final_pred = out * self.blur_mask + whole_depth_roi_pred * (1-self.blur_mask)
- # out = F.interpolate(out, (540, 960), mode='bilinear', align_corners=True)
- x = self.conditional_log_binomial(last, b_embedding)
- b_centers = nn.functional.interpolate(
- b_centers, x.shape[-2:], mode='bilinear', align_corners=True)
-
- out = torch.sum(x * b_centers, dim=1, keepdim=True)
-
-
- final_pred = out
- output = dict(metric_depth=final_pred)
-
- output['coarse_depth_pred'] = whole_depth_pred
- output['fine_depth_pred'] = fine_depth_pred
- output['coarse_depth_pred_roi'] = whole_depth_roi_pred
- if self.consistency_training:
- if self.consistency_target == 'final_feat' or self.consistency_target == 'unet_feat':
- output['temp_features'] = proj_feat_list
-
- output['previous_info'] = previous_info
-
- return output
-
- @staticmethod
- def build(midas_model_type="DPT_BEiT_L_384", pretrained_resource=None, use_pretrained_midas=False, train_midas=False, freeze_midas_bn=True, coarse_model_path=None, fine_model_path=None, **kwargs):
- from zoedepth.models.zoedepth_custom.zoedepth_custom import ZoeDepthCustom
-
- print("build pretrained condition model from {}".format(coarse_model_path))
- coarse_model = ZoeDepthCustom.build(
- midas_model_type=midas_model_type,
- pretrained_resource=coarse_model_path,
- # pretrained_resource="",
- use_pretrained_midas=use_pretrained_midas,
- # use_pretrained_midas=False,
- train_midas=train_midas,
- freeze_midas_bn=freeze_midas_bn,
- **kwargs)
-
- print("build pretrained condition model from {}".format(fine_model_path))
- fine_model = ZoeDepthCustom.build(
- midas_model_type=midas_model_type,
- pretrained_resource=fine_model_path,
- # pretrained_resource="",
- use_pretrained_midas=use_pretrained_midas,
- # use_pretrained_midas=False,
- train_midas=train_midas,
- freeze_midas_bn=freeze_midas_bn,
- **kwargs)
-
- model = PatchFusion(coarse_model, fine_model, **kwargs)
- if pretrained_resource:
- assert isinstance(pretrained_resource, str), "pretrained_resource must be a string"
- model = load_state_from_resource(model, pretrained_resource)
-
- return model
-
- @staticmethod
- def build_from_config(config):
- return PatchFusion.build(**config)
-
- def remove_hooks(self):
- for h in self.handles:
- h.remove()
- return self
-
- def set_fetch_features(self, fetch_features):
- self.fetch_features = fetch_features
- if fetch_features:
- if len(self.handles) == 0:
- self.attach_hooks()
- else:
- self.remove_hooks()
- return self
-
- def attach_hooks(self):
-
- self.handles.append(self.seed_projector.register_forward_hook(get_activation("seed_projector", self.hook_feats)))
-
- for idx, proj in enumerate(self.projectors):
- self.handles.append(proj.register_forward_hook(get_activation("projector_{}".format(idx), self.hook_feats)))
-
- for idx, proj in enumerate(self.attractors):
- self.handles.append(proj.register_forward_hook(get_activation("attractor_{}".format(idx), self.hook_feats)))
-
- return self
diff --git a/zoedepth/models/zoedepth_custom/zoedepth_custom.py b/zoedepth/models/zoedepth_custom/zoedepth_custom.py
deleted file mode 100644
index e810cdd09842133cd6a311768df328215e7309d1..0000000000000000000000000000000000000000
--- a/zoedepth/models/zoedepth_custom/zoedepth_custom.py
+++ /dev/null
@@ -1,335 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Shariq Farooq Bhat
-
-import itertools
-
-import math
-import torch
-import torch.nn as nn
-import numpy as np
-
-from zoedepth.models.depth_model import DepthModel
-from zoedepth.models.base_models.midas import MidasCore
-from zoedepth.models.layers.attractor import AttractorLayer, AttractorLayerUnnormed
-from zoedepth.models.layers.dist_layers import ConditionalLogBinomial, ConditionalLogBinomialV2
-from zoedepth.models.layers.localbins_layers import (Projector, SeedBinRegressor,
- SeedBinRegressorUnnormed)
-from zoedepth.models.model_io import load_state_from_resource
-from torchvision.transforms import Normalize
-
-def get_activation(name, bank):
- # input of forward_hook will be a function of model/inp/oup
- def hook(module, input, output):
- bank[name] = output
- return hook
-
-class ZoeDepthCustom(DepthModel):
- def __init__(self, core, n_bins=64, bin_centers_type="softplus", bin_embedding_dim=128, min_depth=1e-3, max_depth=10,
- n_attractors=[16, 8, 4, 1], attractor_alpha=300, attractor_gamma=2, attractor_kind='sum', attractor_type='exp', min_temp=5, max_temp=50, train_midas=True,
- midas_lr_factor=10, encoder_lr_factor=10, pos_enc_lr_factor=10, inverse_midas=False, sr_ratio=1,
- raw_depth_shape=(2160, 3840), transform_sample_gt_size=(2160, 3840), representation='',
- fetch_features=True, sample_feat_level=3, use_hr=False, deform=False, wo_bins=False, baseline=False, **kwargs):
- """ZoeDepth model. This is the version of ZoeDepth that has a single metric head
-
- Args:
- core (models.base_models.midas.MidasCore): The base midas model that is used for extraction of "relative" features
- n_bins (int, optional): Number of bin centers. Defaults to 64.
- bin_centers_type (str, optional): "normed" or "softplus". Activation type used for bin centers. For "normed" bin centers, linear normalization trick is applied. This results in bounded bin centers.
- For "softplus", softplus activation is used and thus are unbounded. Defaults to "softplus".
- bin_embedding_dim (int, optional): bin embedding dimension. Defaults to 128.
- min_depth (float, optional): Lower bound for normed bin centers. Defaults to 1e-3.
- max_depth (float, optional): Upper bound for normed bin centers. Defaults to 10.
- n_attractors (List[int], optional): Number of bin attractors at decoder layers. Defaults to [16, 8, 4, 1].
- attractor_alpha (int, optional): Proportional attractor strength. Refer to models.layers.attractor for more details. Defaults to 300.
- attractor_gamma (int, optional): Exponential attractor strength. Refer to models.layers.attractor for more details. Defaults to 2.
- attractor_kind (str, optional): Attraction aggregation "sum" or "mean". Defaults to 'sum'.
- attractor_type (str, optional): Type of attractor to use; "inv" (Inverse attractor) or "exp" (Exponential attractor). Defaults to 'exp'.
- min_temp (int, optional): Lower bound for temperature of output probability distribution. Defaults to 5.
- max_temp (int, optional): Upper bound for temperature of output probability distribution. Defaults to 50.
- train_midas (bool, optional): Whether to train "core", the base midas model. Defaults to True.
- midas_lr_factor (int, optional): Learning rate reduction factor for base midas model except its encoder and positional encodings. Defaults to 10.
- encoder_lr_factor (int, optional): Learning rate reduction factor for the encoder in midas model. Defaults to 10.
- pos_enc_lr_factor (int, optional): Learning rate reduction factor for positional encodings in the base midas model. Defaults to 10.
-
- sr_ratio: sr ratio during infer
- raw_depth_shape: raw depth shape during infer. times sr_ratio will be the target resolution. Used to sample points during training
- transform_sample_gt_size: training depth shape # influenced by crop shape which is not included in this pipeline right now
- representation: I use it to test the "bilap head" and a discarded idea
- fetch_features: if fetch feats. Default=True
- """
- super().__init__()
-
- self.core = core
- self.max_depth = max_depth
- self.min_depth = min_depth
- self.min_temp = min_temp
- self.bin_centers_type = bin_centers_type
-
- self.midas_lr_factor = midas_lr_factor
- self.encoder_lr_factor = encoder_lr_factor
- self.pos_enc_lr_factor = pos_enc_lr_factor
- self.train_midas = train_midas
- self.inverse_midas = inverse_midas
-
- if self.encoder_lr_factor <= 0:
- self.core.freeze_encoder(
- freeze_rel_pos=self.pos_enc_lr_factor <= 0)
-
- N_MIDAS_OUT = 32
- btlnck_features = self.core.output_channels[0]
- num_out_features = self.core.output_channels[1:]
-
- self.conv2 = nn.Conv2d(btlnck_features, btlnck_features,
- kernel_size=1, stride=1, padding=0) # btlnck conv
-
- if bin_centers_type == "normed":
- SeedBinRegressorLayer = SeedBinRegressor
- Attractor = AttractorLayer
- elif bin_centers_type == "softplus": # default
- SeedBinRegressorLayer = SeedBinRegressorUnnormed
- Attractor = AttractorLayerUnnormed
- elif bin_centers_type == "hybrid1":
- SeedBinRegressorLayer = SeedBinRegressor
- Attractor = AttractorLayerUnnormed
- elif bin_centers_type == "hybrid2":
- SeedBinRegressorLayer = SeedBinRegressorUnnormed
- Attractor = AttractorLayer
- else:
- raise ValueError(
- "bin_centers_type should be one of 'normed', 'softplus', 'hybrid1', 'hybrid2'")
-
- self.seed_bin_regressor = SeedBinRegressorLayer(
- btlnck_features, n_bins=n_bins, min_depth=min_depth, max_depth=max_depth)
- self.seed_projector = Projector(btlnck_features, bin_embedding_dim)
- self.projectors = nn.ModuleList([
- Projector(num_out, bin_embedding_dim)
- for num_out in num_out_features
- ])
- # 1000, 2, inv, mean
- self.attractors = nn.ModuleList([
- Attractor(bin_embedding_dim, n_bins, n_attractors=n_attractors[i], min_depth=min_depth, max_depth=max_depth,
- alpha=attractor_alpha, gamma=attractor_gamma, kind=attractor_kind, attractor_type=attractor_type)
- for i in range(len(num_out_features))
- ])
-
- last_in = N_MIDAS_OUT + 1 # +1 for relative depth
-
- # use log binomial instead of softmax
- self.conditional_log_binomial = ConditionalLogBinomial(
- last_in, bin_embedding_dim, n_classes=n_bins, min_temp=min_temp, max_temp=max_temp)
-
- self.handles = []
- self.hook_feats = {}
- self.set_fetch_features(fetch_features)
-
- self.baseline = baseline
-
- def init_weight(self):
- if self.deform:
- for m in self.mlp_feat_offset.modules():
- if isinstance(m, nn.Conv1d):
- torch.nn.init.constant_(m.weight, 0)
-
- if self.representation == 'biLaplacian':
- for m in self.proj_x_block:
- torch.nn.init.constant_(m.weight, 0)
-
- def forward(self, x, sampled_depth=None, mode='train', return_final_centers=False, denorm=False, return_probs=False, image_raw=None, **kwargs):
- """
- Args:
- x (torch.Tensor): Input image tensor of shape (B, C, H, W)
- return_final_centers (bool, optional): Whether to return the final bin centers. Defaults to False.
- denorm (bool, optional): Whether to denormalize the input image. This reverses ImageNet normalization as midas normalization is different. Defaults to False.
- return_probs (bool, optional): Whether to return the output probability distribution. Defaults to False.
-
- Returns:
- dict: Dictionary containing the following keys:
- - rel_depth (torch.Tensor): Relative depth map of shape (B, H, W)
- - metric_depth (torch.Tensor): Metric depth map of shape (B, 1, H, W)
- - bin_centers (torch.Tensor): Bin centers of shape (B, n_bins). Present only if return_final_centers is True
- - probs (torch.Tensor): Output probability distribution of shape (B, n_bins, H, W). Present only if return_probs is True
-
- """
-
- b, c, h, w = x.shape
- # print("input shape ", x.shape)
- self.orig_input_width = w
- self.orig_input_height = h
- rel_depth, out = self.core(x, denorm=denorm, return_rel_depth=True)
- # print("output shapes", rel_depth.shape, out.shape)
-
- outconv_activation = out[0]
- btlnck = out[1]
- x_blocks = out[2:]
-
- x_d0 = self.conv2(btlnck)
- self.hook_feats['x_d0'] = x_d0
- x = x_d0
-
- last = outconv_activation
- self.hook_feats['midas_final_feat'] = last
-
- if self.inverse_midas:
- # invert depth followed by normalization
- rel_depth = 1.0 / (rel_depth + 1e-6)
- rel_depth = (rel_depth - rel_depth.min()) / \
- (rel_depth.max() - rel_depth.min())
- # concat rel depth with last. First interpolate rel depth to last size
-
- rel_cond = rel_depth.unsqueeze(1)
- self.hook_feats['rel_depth'] = rel_cond
-
- _, seed_b_centers = self.seed_bin_regressor(x)
-
- if self.bin_centers_type == 'normed' or self.bin_centers_type == 'hybrid2':
- b_prev = (seed_b_centers - self.min_depth) / \
- (self.max_depth - self.min_depth)
- else:
- b_prev = seed_b_centers
-
- prev_b_embedding = self.seed_projector(x)
-
- # unroll this loop for better performance
-
- for idx, (projector, attractor, x) in enumerate(zip(self.projectors, self.attractors, x_blocks)):
- b_embedding = projector(x)
- self.hook_feats['x_blocks_feat_{}'.format(idx)] = x
- # self.hook_feats['attractor_{}_0'.format(idx)] = b_embedding
- # self.hook_feats['attractor_{}_1'.format(idx)] = b_prev
- # self.hook_feats['attractor_{}_2'.format(idx)] = prev_b_embedding
- b, b_centers = attractor(
- b_embedding, b_prev, prev_b_embedding, interpolate=True)
- b_prev = b.clone()
- prev_b_embedding = b_embedding.clone()
-
- self.hook_feats['b_centers'] = b_centers
-
- if self.baseline:
- rel_cond = nn.functional.interpolate(
- rel_cond, size=last.shape[2:], mode='bilinear', align_corners=True)
- last = torch.cat([last, rel_cond], dim=1)
- b_embedding = nn.functional.interpolate(
- b_embedding, last.shape[-2:], mode='bilinear', align_corners=True)
- # till here, we have features (attached with a relative depth prediction) and embeddings
- x = self.conditional_log_binomial(last, b_embedding)
- b_centers = nn.functional.interpolate(
- b_centers, x.shape[-2:], mode='bilinear', align_corners=True)
- self.hook_feats['b_centers'] = b_centers
- out = torch.sum(x * b_centers, dim=1, keepdim=True)
- output = dict(metric_depth=out)
- return output
-
- # if mode == 'train':
- # if self.wo_bins and self.representation == 'biLaplacian':
- # pred_depth_dict = self.implicit_function_train_process(sampled_depth, mode=mode)
- # output = pred_depth_dict
- # else:
- # pred_depth = self.implicit_function_train_process(sampled_depth, mode=mode)
- # output = dict(metric_depth=pred_depth)
-
- # elif mode == 'eval':
- # pred_depth = self.implicit_function_infer_process(mode=mode)
- # output = dict(metric_depth=pred_depth)
- # else:
- # raise NotImplementedError
-
- return output
-
- # change param lr later
- def get_lr_params(self, lr):
- """
- Learning rate configuration for different layers of the model
- Args:
- lr (float) : Base learning rate
- Returns:
- list : list of parameters to optimize and their learning rates, in the format required by torch optimizers.
- """
- param_conf = []
- if self.train_midas:
- if self.encoder_lr_factor > 0:
- param_conf.append({'params': self.core.get_enc_params_except_rel_pos(
- ), 'lr': lr / self.encoder_lr_factor})
-
- if self.pos_enc_lr_factor > 0:
- param_conf.append(
- {'params': self.core.get_rel_pos_params(), 'lr': lr / self.pos_enc_lr_factor})
-
- midas_params = self.core.core.scratch.parameters()
- midas_lr_factor = self.midas_lr_factor
- param_conf.append(
- {'params': midas_params, 'lr': lr / midas_lr_factor})
-
- remaining_modules = []
- for name, child in self.named_children():
- if name != 'core':
- remaining_modules.append(child)
- remaining_params = itertools.chain(
- *[child.parameters() for child in remaining_modules])
-
- param_conf.append({'params': remaining_params, 'lr': lr})
-
- return param_conf
-
- @staticmethod
- def build(midas_model_type="DPT_BEiT_L_384", pretrained_resource=None, use_pretrained_midas=False, train_midas=False, freeze_midas_bn=True, **kwargs):
- core = MidasCore.build(midas_model_type=midas_model_type, use_pretrained_midas=use_pretrained_midas,
- train_midas=train_midas, fetch_features=True, freeze_bn=freeze_midas_bn, **kwargs)
- model = ZoeDepthCustom(core, **kwargs)
- if pretrained_resource:
- assert isinstance(pretrained_resource, str), "pretrained_resource must be a string"
- model = load_state_from_resource(model, pretrained_resource)
- # a = torch.ones((2, 3, 100, 200))
- # b = torch.ones((2, 500, 3))
- # model(a, b)
- return model
-
- @staticmethod
- def build_from_config(config):
- return ZoeDepthCustom.build(**config)
-
- def remove_hooks(self):
- for h in self.handles:
- h.remove()
- return self
-
- def set_fetch_features(self, fetch_features):
- self.fetch_features = fetch_features
- if fetch_features:
- if len(self.handles) == 0:
- self.attach_hooks()
- else:
- self.remove_hooks()
- return self
-
- def attach_hooks(self):
-
- self.handles.append(self.seed_projector.register_forward_hook(get_activation("seed_projector", self.hook_feats)))
-
- for idx, proj in enumerate(self.projectors):
- self.handles.append(proj.register_forward_hook(get_activation("projector_{}".format(idx), self.hook_feats)))
-
- for idx, proj in enumerate(self.attractors):
- self.handles.append(proj.register_forward_hook(get_activation("attractor_{}".format(idx), self.hook_feats)))
-
- return self
diff --git a/zoedepth/models/zoedepth_custom/__init__.py b/zoedepth/models/zoedepth_nk/__init__.py
similarity index 87%
rename from zoedepth/models/zoedepth_custom/__init__.py
rename to zoedepth/models/zoedepth_nk/__init__.py
index 2de3a6af3d3834d7895f2eb1d4c6f17ba77ef9f5..513a278b939c10c010e3c0250ec73544d5663886 100644
--- a/zoedepth/models/zoedepth_custom/__init__.py
+++ b/zoedepth/models/zoedepth_nk/__init__.py
@@ -20,14 +20,12 @@
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
-# File author: Zhenyu Li
+# File author: Shariq Farooq Bhat
-from .zoedepth_custom import ZoeDepthCustom
-from .patchfusion import PatchFusion
+from .zoedepth_nk_v1 import ZoeDepthNK
all_versions = {
- "custom": ZoeDepthCustom,
- "patchfusion": PatchFusion
+ "v1": ZoeDepthNK,
}
get_version = lambda v : all_versions[v]
\ No newline at end of file
diff --git a/zoedepth/models/zoedepth_nk/config_zoedepth_nk.json b/zoedepth/models/zoedepth_nk/config_zoedepth_nk.json
new file mode 100644
index 0000000000000000000000000000000000000000..42bab2a3ad159a09599a5aba270c491021a3cf1a
--- /dev/null
+++ b/zoedepth/models/zoedepth_nk/config_zoedepth_nk.json
@@ -0,0 +1,67 @@
+{
+ "model": {
+ "name": "ZoeDepthNK",
+ "version_name": "v1",
+ "bin_conf" : [
+ {
+ "name": "nyu",
+ "n_bins": 64,
+ "min_depth": 1e-3,
+ "max_depth": 10.0
+ },
+ {
+ "name": "kitti",
+ "n_bins": 64,
+ "min_depth": 1e-3,
+ "max_depth": 80.0
+ }
+ ],
+ "bin_embedding_dim": 128,
+ "bin_centers_type": "softplus",
+ "n_attractors":[16, 8, 4, 1],
+ "attractor_alpha": 1000,
+ "attractor_gamma": 2,
+ "attractor_kind" : "mean",
+ "attractor_type" : "inv",
+ "min_temp": 0.0212,
+ "max_temp": 50.0,
+ "memory_efficient": true,
+ "midas_model_type" : "DPT_BEiT_L_384",
+ "img_size": [384, 512]
+ },
+
+ "train": {
+ "train_midas": true,
+ "use_pretrained_midas": true,
+ "trainer": "zoedepth_nk",
+ "epochs": 5,
+ "bs": 16,
+ "optim_kwargs": {"lr": 0.0002512, "wd": 0.01},
+ "sched_kwargs": {"div_factor": 1, "final_div_factor": 10000, "pct_start": 0.7, "three_phase":false, "cycle_momentum": true},
+ "same_lr": false,
+ "w_si": 1,
+ "w_domain": 100,
+ "avoid_boundary": false,
+ "random_crop": false,
+ "input_width": 640,
+ "input_height": 480,
+ "w_grad": 0,
+ "w_reg": 0,
+ "midas_lr_factor": 10,
+ "encoder_lr_factor":10,
+ "pos_enc_lr_factor":10
+ },
+
+ "infer": {
+ "train_midas": false,
+ "pretrained_resource": "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_NK.pt",
+ "use_pretrained_midas": false,
+ "force_keep_ar": true
+ },
+
+ "eval": {
+ "train_midas": false,
+ "pretrained_resource": "url::https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_NK.pt",
+ "use_pretrained_midas": false
+ }
+}
\ No newline at end of file
diff --git a/zoedepth/models/zoedepth_nk/zoedepth_nk_v1.py b/zoedepth/models/zoedepth_nk/zoedepth_nk_v1.py
new file mode 100644
index 0000000000000000000000000000000000000000..7368ae8031188a9f946d9d3f29633c96e791e68e
--- /dev/null
+++ b/zoedepth/models/zoedepth_nk/zoedepth_nk_v1.py
@@ -0,0 +1,333 @@
+# MIT License
+
+# Copyright (c) 2022 Intelligent Systems Lab Org
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# File author: Shariq Farooq Bhat
+
+import itertools
+
+import torch
+import torch.nn as nn
+
+from zoedepth.models.depth_model import DepthModel
+from zoedepth.models.base_models.midas import MidasCore
+from zoedepth.models.layers.attractor import AttractorLayer, AttractorLayerUnnormed
+from zoedepth.models.layers.dist_layers import ConditionalLogBinomial
+from zoedepth.models.layers.localbins_layers import (Projector, SeedBinRegressor,
+ SeedBinRegressorUnnormed)
+from zoedepth.models.layers.patch_transformer import PatchTransformerEncoder
+from zoedepth.models.model_io import load_state_from_resource
+
+
+class ZoeDepthNK(DepthModel):
+ def __init__(self, core, bin_conf, bin_centers_type="softplus", bin_embedding_dim=128,
+ n_attractors=[16, 8, 4, 1], attractor_alpha=300, attractor_gamma=2, attractor_kind='sum', attractor_type='exp',
+ min_temp=5, max_temp=50,
+ memory_efficient=False, train_midas=True,
+ is_midas_pretrained=True, midas_lr_factor=1, encoder_lr_factor=10, pos_enc_lr_factor=10, inverse_midas=False, **kwargs):
+ """ZoeDepthNK model. This is the version of ZoeDepth that has two metric heads and uses a learned router to route to experts.
+
+ Args:
+ core (models.base_models.midas.MidasCore): The base midas model that is used for extraction of "relative" features
+
+ bin_conf (List[dict]): A list of dictionaries that contain the bin configuration for each metric head. Each dictionary should contain the following keys:
+ "name" (str, typically same as the dataset name), "n_bins" (int), "min_depth" (float), "max_depth" (float)
+
+ The length of this list determines the number of metric heads.
+ bin_centers_type (str, optional): "normed" or "softplus". Activation type used for bin centers. For "normed" bin centers, linear normalization trick is applied. This results in bounded bin centers.
+ For "softplus", softplus activation is used and thus are unbounded. Defaults to "normed".
+ bin_embedding_dim (int, optional): bin embedding dimension. Defaults to 128.
+
+ n_attractors (List[int], optional): Number of bin attractors at decoder layers. Defaults to [16, 8, 4, 1].
+ attractor_alpha (int, optional): Proportional attractor strength. Refer to models.layers.attractor for more details. Defaults to 300.
+ attractor_gamma (int, optional): Exponential attractor strength. Refer to models.layers.attractor for more details. Defaults to 2.
+ attractor_kind (str, optional): Attraction aggregation "sum" or "mean". Defaults to 'sum'.
+ attractor_type (str, optional): Type of attractor to use; "inv" (Inverse attractor) or "exp" (Exponential attractor). Defaults to 'exp'.
+
+ min_temp (int, optional): Lower bound for temperature of output probability distribution. Defaults to 5.
+ max_temp (int, optional): Upper bound for temperature of output probability distribution. Defaults to 50.
+
+ memory_efficient (bool, optional): Whether to use memory efficient version of attractor layers. Memory efficient version is slower but is recommended incase of multiple metric heads in order save GPU memory. Defaults to False.
+
+ train_midas (bool, optional): Whether to train "core", the base midas model. Defaults to True.
+ is_midas_pretrained (bool, optional): Is "core" pretrained? Defaults to True.
+ midas_lr_factor (int, optional): Learning rate reduction factor for base midas model except its encoder and positional encodings. Defaults to 10.
+ encoder_lr_factor (int, optional): Learning rate reduction factor for the encoder in midas model. Defaults to 10.
+ pos_enc_lr_factor (int, optional): Learning rate reduction factor for positional encodings in the base midas model. Defaults to 10.
+
+ """
+
+ super().__init__()
+
+ self.core = core
+ self.bin_conf = bin_conf
+ self.min_temp = min_temp
+ self.max_temp = max_temp
+ self.memory_efficient = memory_efficient
+ self.train_midas = train_midas
+ self.is_midas_pretrained = is_midas_pretrained
+ self.midas_lr_factor = midas_lr_factor
+ self.encoder_lr_factor = encoder_lr_factor
+ self.pos_enc_lr_factor = pos_enc_lr_factor
+ self.inverse_midas = inverse_midas
+
+ N_MIDAS_OUT = 32
+ btlnck_features = self.core.output_channels[0]
+ num_out_features = self.core.output_channels[1:]
+ # self.scales = [16, 8, 4, 2] # spatial scale factors
+
+ self.conv2 = nn.Conv2d(
+ btlnck_features, btlnck_features, kernel_size=1, stride=1, padding=0)
+
+ # Transformer classifier on the bottleneck
+ self.patch_transformer = PatchTransformerEncoder(
+ btlnck_features, 1, 128, use_class_token=True)
+ self.mlp_classifier = nn.Sequential(
+ nn.Linear(128, 128),
+ nn.ReLU(),
+ nn.Linear(128, 2)
+ )
+
+ if bin_centers_type == "normed":
+ SeedBinRegressorLayer = SeedBinRegressor
+ Attractor = AttractorLayer
+ elif bin_centers_type == "softplus":
+ SeedBinRegressorLayer = SeedBinRegressorUnnormed
+ Attractor = AttractorLayerUnnormed
+ elif bin_centers_type == "hybrid1":
+ SeedBinRegressorLayer = SeedBinRegressor
+ Attractor = AttractorLayerUnnormed
+ elif bin_centers_type == "hybrid2":
+ SeedBinRegressorLayer = SeedBinRegressorUnnormed
+ Attractor = AttractorLayer
+ else:
+ raise ValueError(
+ "bin_centers_type should be one of 'normed', 'softplus', 'hybrid1', 'hybrid2'")
+ self.bin_centers_type = bin_centers_type
+ # We have bins for each bin conf.
+ # Create a map (ModuleDict) of 'name' -> seed_bin_regressor
+ self.seed_bin_regressors = nn.ModuleDict(
+ {conf['name']: SeedBinRegressorLayer(btlnck_features, conf["n_bins"], mlp_dim=bin_embedding_dim//2, min_depth=conf["min_depth"], max_depth=conf["max_depth"])
+ for conf in bin_conf}
+ )
+
+ self.seed_projector = Projector(
+ btlnck_features, bin_embedding_dim, mlp_dim=bin_embedding_dim//2)
+ self.projectors = nn.ModuleList([
+ Projector(num_out, bin_embedding_dim, mlp_dim=bin_embedding_dim//2)
+ for num_out in num_out_features
+ ])
+
+ # Create a map (ModuleDict) of 'name' -> attractors (ModuleList)
+ self.attractors = nn.ModuleDict(
+ {conf['name']: nn.ModuleList([
+ Attractor(bin_embedding_dim, n_attractors[i],
+ mlp_dim=bin_embedding_dim, alpha=attractor_alpha,
+ gamma=attractor_gamma, kind=attractor_kind,
+ attractor_type=attractor_type, memory_efficient=memory_efficient,
+ min_depth=conf["min_depth"], max_depth=conf["max_depth"])
+ for i in range(len(n_attractors))
+ ])
+ for conf in bin_conf}
+ )
+
+ last_in = N_MIDAS_OUT
+ # conditional log binomial for each bin conf
+ self.conditional_log_binomial = nn.ModuleDict(
+ {conf['name']: ConditionalLogBinomial(last_in, bin_embedding_dim, conf['n_bins'], bottleneck_factor=4, min_temp=self.min_temp, max_temp=self.max_temp)
+ for conf in bin_conf}
+ )
+
+ def forward(self, x, return_final_centers=False, denorm=False, return_probs=False, **kwargs):
+ """
+ Args:
+ x (torch.Tensor): Input image tensor of shape (B, C, H, W). Assumes all images are from the same domain.
+ return_final_centers (bool, optional): Whether to return the final centers of the attractors. Defaults to False.
+ denorm (bool, optional): Whether to denormalize the input image. Defaults to False.
+ return_probs (bool, optional): Whether to return the probabilities of the bins. Defaults to False.
+
+ Returns:
+ dict: Dictionary of outputs with keys:
+ - "rel_depth": Relative depth map of shape (B, 1, H, W)
+ - "metric_depth": Metric depth map of shape (B, 1, H, W)
+ - "domain_logits": Domain logits of shape (B, 2)
+ - "bin_centers": Bin centers of shape (B, N, H, W). Present only if return_final_centers is True
+ - "probs": Bin probabilities of shape (B, N, H, W). Present only if return_probs is True
+ """
+ b, c, h, w = x.shape
+ self.orig_input_width = w
+ self.orig_input_height = h
+ rel_depth, out = self.core(x, denorm=denorm, return_rel_depth=True)
+
+ outconv_activation = out[0]
+ btlnck = out[1]
+ x_blocks = out[2:]
+
+ x_d0 = self.conv2(btlnck)
+ x = x_d0
+
+ # Predict which path to take
+ embedding = self.patch_transformer(x)[0] # N, E
+ domain_logits = self.mlp_classifier(embedding) # N, 2
+ domain_vote = torch.softmax(domain_logits.sum(
+ dim=0, keepdim=True), dim=-1) # 1, 2
+
+ # Get the path
+ bin_conf_name = ["nyu", "kitti"][torch.argmax(
+ domain_vote, dim=-1).squeeze().item()]
+
+ try:
+ conf = [c for c in self.bin_conf if c.name == bin_conf_name][0]
+ except IndexError:
+ raise ValueError(
+ f"bin_conf_name {bin_conf_name} not found in bin_confs")
+
+ min_depth = conf['min_depth']
+ max_depth = conf['max_depth']
+
+ seed_bin_regressor = self.seed_bin_regressors[bin_conf_name]
+ _, seed_b_centers = seed_bin_regressor(x)
+ if self.bin_centers_type == 'normed' or self.bin_centers_type == 'hybrid2':
+ b_prev = (seed_b_centers - min_depth)/(max_depth - min_depth)
+ else:
+ b_prev = seed_b_centers
+ prev_b_embedding = self.seed_projector(x)
+
+ attractors = self.attractors[bin_conf_name]
+ for projector, attractor, x in zip(self.projectors, attractors, x_blocks):
+ b_embedding = projector(x)
+ b, b_centers = attractor(
+ b_embedding, b_prev, prev_b_embedding, interpolate=True)
+ b_prev = b
+ prev_b_embedding = b_embedding
+
+ last = outconv_activation
+
+ b_centers = nn.functional.interpolate(
+ b_centers, last.shape[-2:], mode='bilinear', align_corners=True)
+ b_embedding = nn.functional.interpolate(
+ b_embedding, last.shape[-2:], mode='bilinear', align_corners=True)
+
+ clb = self.conditional_log_binomial[bin_conf_name]
+ x = clb(last, b_embedding)
+
+ # Now depth value is Sum px * cx , where cx are bin_centers from the last bin tensor
+ # print(x.shape, b_centers.shape)
+ # b_centers = nn.functional.interpolate(b_centers, x.shape[-2:], mode='bilinear', align_corners=True)
+ out = torch.sum(x * b_centers, dim=1, keepdim=True)
+
+ output = dict(domain_logits=domain_logits, metric_depth=out)
+ if return_final_centers or return_probs:
+ output['bin_centers'] = b_centers
+
+ if return_probs:
+ output['probs'] = x
+ return output
+
+ def get_lr_params(self, lr):
+ """
+ Learning rate configuration for different layers of the model
+
+ Args:
+ lr (float) : Base learning rate
+ Returns:
+ list : list of parameters to optimize and their learning rates, in the format required by torch optimizers.
+ """
+ param_conf = []
+ if self.train_midas:
+ def get_rel_pos_params():
+ for name, p in self.core.core.pretrained.named_parameters():
+ if "relative_position" in name:
+ yield p
+
+ def get_enc_params_except_rel_pos():
+ for name, p in self.core.core.pretrained.named_parameters():
+ if "relative_position" not in name:
+ yield p
+
+ encoder_params = get_enc_params_except_rel_pos()
+ rel_pos_params = get_rel_pos_params()
+ midas_params = self.core.core.scratch.parameters()
+ midas_lr_factor = self.midas_lr_factor if self.is_midas_pretrained else 1.0
+ param_conf.extend([
+ {'params': encoder_params, 'lr': lr / self.encoder_lr_factor},
+ {'params': rel_pos_params, 'lr': lr / self.pos_enc_lr_factor},
+ {'params': midas_params, 'lr': lr / midas_lr_factor}
+ ])
+
+ remaining_modules = []
+ for name, child in self.named_children():
+ if name != 'core':
+ remaining_modules.append(child)
+ remaining_params = itertools.chain(
+ *[child.parameters() for child in remaining_modules])
+ param_conf.append({'params': remaining_params, 'lr': lr})
+ return param_conf
+
+ def get_conf_parameters(self, conf_name):
+ """
+ Returns parameters of all the ModuleDicts children that are exclusively used for the given bin configuration
+ """
+ params = []
+ for name, child in self.named_children():
+ if isinstance(child, nn.ModuleDict):
+ for bin_conf_name, module in child.items():
+ if bin_conf_name == conf_name:
+ params += list(module.parameters())
+ return params
+
+ def freeze_conf(self, conf_name):
+ """
+ Freezes all the parameters of all the ModuleDicts children that are exclusively used for the given bin configuration
+ """
+ for p in self.get_conf_parameters(conf_name):
+ p.requires_grad = False
+
+ def unfreeze_conf(self, conf_name):
+ """
+ Unfreezes all the parameters of all the ModuleDicts children that are exclusively used for the given bin configuration
+ """
+ for p in self.get_conf_parameters(conf_name):
+ p.requires_grad = True
+
+ def freeze_all_confs(self):
+ """
+ Freezes all the parameters of all the ModuleDicts children
+ """
+ for name, child in self.named_children():
+ if isinstance(child, nn.ModuleDict):
+ for bin_conf_name, module in child.items():
+ for p in module.parameters():
+ p.requires_grad = False
+
+ @staticmethod
+ def build(midas_model_type="DPT_BEiT_L_384", pretrained_resource=None, use_pretrained_midas=False, train_midas=False, freeze_midas_bn=True, **kwargs):
+ core = MidasCore.build(midas_model_type=midas_model_type, use_pretrained_midas=use_pretrained_midas,
+ train_midas=train_midas, fetch_features=True, freeze_bn=freeze_midas_bn, **kwargs)
+ model = ZoeDepthNK(core, **kwargs)
+ if pretrained_resource:
+ assert isinstance(pretrained_resource, str), "pretrained_resource must be a string"
+ model = load_state_from_resource(model, pretrained_resource)
+ return model
+
+ @staticmethod
+ def build_from_config(config):
+ return ZoeDepthNK.build(**config)
diff --git a/zoedepth/trainers/base_trainer.py b/zoedepth/trainers/base_trainer.py
index 6009b9e5f64b32b019fcdae3798c59c6741c84ca..33fbbea3a7d49efe11b005adb5127f441eabfaf6 100644
--- a/zoedepth/trainers/base_trainer.py
+++ b/zoedepth/trainers/base_trainer.py
@@ -22,8 +22,6 @@
# File author: Shariq Farooq Bhat
-# This file may include modifications from author Zhenyu Li
-
import os
import uuid
import warnings
@@ -63,25 +61,6 @@ class BaseTrainer:
self.optimizer = self.init_optimizer()
self.scheduler = self.init_scheduler()
- # import matplotlib.pyplot as plt
- # lrs = []
- # momentums = []
- # for e in range(self.config.epochs):
- # for s in range(len(self.train_loader)):
- # self.scheduler.step()
- # self.optimizer.step()
- # lr = self.scheduler.get_last_lr()[2]
- # lrs.append(lr)
- # print(self.optimizer.param_groups[0]['betas'])
- # momentum = self.optimizer.param_groups[0]['betas'][0]
- # momentums.append(momentum)
-
- # step = [_ for _ in range(len(lrs))]
- # plt.scatter(step, momentums)
- # plt.savefig("debug.png")
- # exit(100)
-
-
def resize_to_target(self, prediction, target):
if prediction.shape[2:] != target.shape[-2:]:
prediction = nn.functional.interpolate(
@@ -135,8 +114,7 @@ class BaseTrainer:
lrs = [l['lr'] for l in self.optimizer.param_groups]
return optim.lr_scheduler.OneCycleLR(self.optimizer, lrs, epochs=self.config.epochs, steps_per_epoch=len(self.train_loader),
cycle_momentum=self.config.cycle_momentum,
- base_momentum=self.config.get('base_momentum', 0.85), max_momentum=self.config.get('max_momentum', 0.95), div_factor=self.config.div_factor,
- final_div_factor=self.config.final_div_factor, pct_start=self.config.pct_start, three_phase=self.config.three_phase)
+ base_momentum=0.85, max_momentum=0.95, div_factor=self.config.div_factor, final_div_factor=self.config.final_div_factor, pct_start=self.config.pct_start, three_phase=self.config.three_phase)
def train_on_batch(self, batch, train_step):
raise NotImplementedError
@@ -184,6 +162,7 @@ class BaseTrainer:
if self.config.prefetch:
+
for i, batch in tqdm(enumerate(self.train_loader), desc=f"Prefetching...",
total=self.iters_per_epoch) if is_rank_zero(self.config) else enumerate(self.train_loader):
pass
@@ -196,10 +175,6 @@ class BaseTrainer:
break
self.epoch = epoch
- # self.model.eval()
- # metrics, test_losses = self.validate()
- # print(metrics)
- # exit(100)
################################# Train loop ##########################################################
if self.should_log:
wandb.log({"Epoch": epoch}, step=self.step)
@@ -221,12 +196,7 @@ class BaseTrainer:
if self.should_log and self.step % 50 == 0:
wandb.log({f"Train/{name}": loss.item()
- for name, loss in losses.items()}, step=self.step)
- # current_lr = self.optimizer.param_groups[0]['lr']
- current_lr = self.scheduler.get_last_lr()[0]
- wandb.log({f"Train/LR": current_lr}, step=self.step)
- momentum = self.optimizer.param_groups[0]['betas'][0]
- wandb.log({f"Train/momentum": momentum}, step=self.step)
+ for name, loss in losses.items()}, step=self.step)
self.step += 1
@@ -298,10 +268,7 @@ class BaseTrainer:
if metrics:
metrics_avg.update(metrics)
- r1, r2 = metrics_avg.get_value(), losses_avg.get_value()
- if self.should_log and self.config.get("debug", False):
- print(r1)
- return r1, r2
+ return metrics_avg.get_value(), losses_avg.get_value()
def save_checkpoint(self, filename):
if not self.should_write:
@@ -319,31 +286,23 @@ class BaseTrainer:
"epoch": self.epoch
}, fpath)
- def log_images(self, rgb: Dict[str, list] = {}, depth: Dict[str, list] = {}, scalar_field: Dict[str, list] = {}, prefix="", scalar_cmap="turbo_r", min_depth=None, max_depth=None):
+ def log_images(self, rgb: Dict[str, list] = {}, depth: Dict[str, list] = {}, scalar_field: Dict[str, list] = {}, prefix="", scalar_cmap="jet", min_depth=None, max_depth=None):
if not self.should_log:
return
- # if min_depth is None:
- # try:
- # min_depth = self.config.min_depth
- # max_depth = self.config.max_depth
- # except AttributeError:
- # min_depth = None
- # max_depth = None
-
- depths_gt = depth['GT']
- invalid_mask = torch.logical_or(depths_gt<=min_depth, depths_gt>=max_depth).detach().cpu().squeeze().numpy()
- min_depth = None
- max_depth = None
-
- depth = {k: colorize(v, vmin=min_depth, vmax=max_depth, invalid_mask=invalid_mask)
+ if min_depth is None:
+ try:
+ min_depth = self.config.min_depth
+ max_depth = self.config.max_depth
+ except AttributeError:
+ min_depth = None
+ max_depth = None
+
+ depth = {k: colorize(v, vmin=min_depth, vmax=max_depth)
for k, v in depth.items()}
- # depth = {k: colorize(v, vmin=0, vmax=80, invalid_mask=invalid_mask)
- # for k, v in depth.items()}
scalar_field = {k: colorize(
v, vmin=None, vmax=None, cmap=scalar_cmap) for k, v in scalar_field.items()}
images = {**rgb, **depth, **scalar_field}
-
wimages = {
prefix+"Predictions": [wandb.Image(v, caption=k) for k, v in images.items()]}
wandb.log(wimages, step=self.step)
diff --git a/zoedepth/trainers/loss.py b/zoedepth/trainers/loss.py
index 4381f5f956c546211f48c7c762f9e9a314ccab04..0c5a1c15cdf5628c1474c566fdc6e58159d7f5ab 100644
--- a/zoedepth/trainers/loss.py
+++ b/zoedepth/trainers/loss.py
@@ -22,21 +22,15 @@
# File author: Shariq Farooq Bhat
-# This file may include modifications from author Zhenyu Li
-
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.cuda.amp as amp
import numpy as np
-from torch.autograd import Variable
-from math import exp
-import matplotlib.pyplot as plt
KEY_OUTPUT = 'metric_depth'
-# import kornia
-import copy
+
def extract_key(prediction, key):
if isinstance(prediction, dict):
@@ -53,8 +47,6 @@ class SILogLoss(nn.Module):
self.beta = beta
def forward(self, input, target, mask=None, interpolate=True, return_interpolated=False):
- hack_input = input
-
input = extract_key(input, KEY_OUTPUT)
if input.shape[-1] != target.shape[-1] and interpolate:
input = nn.functional.interpolate(
@@ -86,12 +78,6 @@ class SILogLoss(nn.Module):
loss = 10 * torch.sqrt(Dg)
if torch.isnan(loss):
- if input.numel() == 0:
- loss = torch.mean(hack_input) * 0
- if not return_interpolated:
- return loss
- return loss, intr_input
-
print("Nan SILog loss")
print("input:", input.shape)
print("target:", target.shape)
@@ -121,33 +107,6 @@ def grad_mask(mask):
return mask[..., 1:, 1:] & mask[..., 1:, :-1] & mask[..., :-1, 1:]
-# class GradL1Loss(nn.Module):
-# """Gradient loss"""
-# def __init__(self):
-# super(GradL1Loss, self).__init__()
-# self.name = 'GradL1'
-
-# def forward(self, input, target, mask=None, interpolate=True, return_interpolated=False):
-# input = extract_key(input, KEY_OUTPUT)
-# if input.shape[-1] != target.shape[-1] and interpolate:
-# input = nn.functional.interpolate(
-# input, target.shape[-2:], mode='bilinear', align_corners=True)
-# intr_input = input
-# else:
-# intr_input = input
-
-# grad_gt = grad(target)
-# grad_pred = grad(input)
-# mask_g = grad_mask(mask)
-
-# loss = nn.functional.l1_loss(grad_pred[0][mask_g], grad_gt[0][mask_g])
-# loss = loss + \
-# nn.functional.l1_loss(grad_pred[1][mask_g], grad_gt[1][mask_g])
-# if not return_interpolated:
-# return loss
-# return loss, intr_input
-
-
class GradL1Loss(nn.Module):
"""Gradient loss"""
def __init__(self):
@@ -319,7 +278,6 @@ def compute_scale_and_shift(prediction, target, mask):
x_1[valid] = (-a_01[valid] * b_0[valid] + a_00[valid] * b_1[valid]) / det[valid]
return x_0, x_1
-
class ScaleAndShiftInvariantLoss(nn.Module):
def __init__(self):
super().__init__()
@@ -347,23 +305,6 @@ class ScaleAndShiftInvariantLoss(nn.Module):
return loss, intr_input
-class BudgetConstraint(nn.Module):
- """
- Given budget constraint to reduce expected inference FLOPs in the Dynamic Network.
- """
- def __init__(self, loss_mu, flops_all, warm_up=True):
- super().__init__()
- self.loss_mu = loss_mu
- self.flops_all = flops_all
- self.warm_up = warm_up
-
- def forward(self, flops_expt, warm_up_rate=1.0):
- if self.warm_up:
- warm_up_rate = min(1.0, warm_up_rate)
- else:
- warm_up_rate = 1.0
- losses = warm_up_rate * ((flops_expt / self.flops_all - self.loss_mu)**2)
- return losses
if __name__ == '__main__':
@@ -373,483 +314,3 @@ if __name__ == '__main__':
d = torch.Tensor([6.59, 3.8, 10.0])
print(celoss.dequantize_depth(celoss.quantize_depth(d)))
-
-
-
-class HistogramMatchingLoss(nn.Module):
- def __init__(self, min_depth, max_depth, bins=512):
- super(HistogramMatchingLoss, self).__init__()
- self.name = 'HistogramMatchingLoss'
- self.min_depth = min_depth
- self.max_depth = max_depth
- self.bins = bins
-
- def forward(self, input, target, mask, interpolate=True):
- if input.shape[-1] != mask.shape[-1] and interpolate:
- input = nn.functional.interpolate(
- input, mask.shape[-2:], mode='bilinear', align_corners=True)
-
- if target.shape[-1] != mask.shape[-1] and interpolate:
- target = nn.functional.interpolate(
- target, mask.shape[-2:], mode='bilinear', align_corners=True)
-
- input[~mask] = 0
- target[~mask] = 0
-
-
- pred_hist = torch.histc(input, bins=self.bins, min=self.min_depth, max=self.max_depth)
- gt_hist = torch.histc(target, bins=self.bins, min=self.min_depth, max=self.max_depth)
-
- pred_hist /= pred_hist.sum(dim=0, keepdim=True)
- gt_hist /= gt_hist.sum(dim=0, keepdim=True)
-
- # print(pred_hist.shape)
- # print(pred_hist)
- # _pred_hist = pred_hist.detach().cpu().numpy()
- # _gt_hist = gt_hist.detach().cpu().numpy()
- # plt.subplot(2, 1, 1)
- # plt.bar(range(len(_pred_hist)), _pred_hist)
- # plt.subplot(2, 1, 2)
- # plt.bar(range(len(_gt_hist)), _gt_hist)
- # plt.savefig('./debug_scale.png')
-
- # Compute cumulative histograms (CDF)
- cdf_pred = torch.cumsum(pred_hist, dim=0)
- cdf_gt = torch.cumsum(gt_hist, dim=0)
-
- # Compute Earth Mover's Distance (EMD) between the CDFs
- loss = torch.mean(torch.abs(cdf_pred - cdf_gt))
- # loss = torch.mean(torch.sqrt((pred_hist - gt_hist)**2))
- # loss = F.kl_div(torch.log(pred_hist + 1e-10), gt_hist, reduction='mean')
-
- return loss
-
-
-
-
-def gaussian(window_size, sigma):
- gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
- return gauss/gauss.sum()
-
-def create_window(window_size, channel):
- _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
- _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
- window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
- return window
-
-def _ssim(img1, img2, window, window_size, channel, size_average = True):
- mu1 = F.conv2d(img1, window, padding = window_size//2, groups = channel)
- mu2 = F.conv2d(img2, window, padding = window_size//2, groups = channel)
-
- mu1_sq = mu1.pow(2)
- mu2_sq = mu2.pow(2)
- mu1_mu2 = mu1*mu2
-
- sigma1_sq = F.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq
- sigma2_sq = F.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq
- sigma12 = F.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2
-
- C1 = 0.01**2
- C2 = 0.03**2
-
- ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
-
- if size_average:
- return ssim_map.mean()
- else:
- return ssim_map.mean(1).mean(1).mean(1)
-
-class SSIM(torch.nn.Module):
- def __init__(self, window_size = 11, size_average = True):
- super(SSIM, self).__init__()
- self.window_size = window_size
- self.size_average = size_average
- self.channel = 1
- self.window = create_window(window_size, self.channel)
-
- def forward(self, img1, img2, mask, interpolate=True):
- if img1.shape[-1] != mask.shape[-1] and interpolate:
- img1 = nn.functional.interpolate(
- img1, mask.shape[-2:], mode='bilinear', align_corners=True)
-
- if img2.shape[-1] != mask.shape[-1] and interpolate:
- img2 = nn.functional.interpolate(
- img2, mask.shape[-2:], mode='bilinear', align_corners=True)
-
- img1[~mask] = 0
- img2[~mask] = 0
-
- (_, channel, _, _) = img1.size()
-
- if channel == self.channel and self.window.data.type() == img1.data.type():
- window = self.window
- else:
- window = create_window(self.window_size, channel)
-
- if img1.is_cuda:
- window = window.cuda(img1.get_device())
- window = window.type_as(img1)
-
- self.window = window
- self.channel = channel
-
-
- loss = _ssim(img1, img2, window, self.window_size, channel, self.size_average)
- return loss
-
-def ssim(img1, img2, window_size = 11, size_average = True):
- (_, channel, _, _) = img1.size()
- window = create_window(window_size, channel)
-
- if img1.is_cuda:
- window = window.cuda(img1.get_device())
- window = window.type_as(img1)
-
- return _ssim(img1, img2, window, window_size, channel, size_average)
-
-class ConsistencyLoss(nn.Module):
- def __init__(self, target, focus_flatten=False, wp=1) -> None:
- super().__init__()
- self.name = 'Consistency'
- self.target = target
- self.mode = 'no-resize'
- # self.mode = 'resize'
- self.focus_flatten = focus_flatten
- self.wp = wp
-
- def gradient_y(self, img):
- # gy = torch.cat([F.conv2d(img[:, i, :, :].unsqueeze(0), torch.Tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]]).view((1, 1, 3, 3)).to(img.device), padding=1) for i in range(img.shape[1])], 1)
- gy = F.conv2d(img, torch.Tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]]).view((1, 1, 3, 3)).to(img.device), padding=1)
- return gy
-
- def gradient_x(self, img):
- # gx = torch.cat([F.conv2d(img[:, i, :, :].unsqueeze(0), torch.Tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]]).view((1, 1, 3, 3)).to(img.device), padding=1) for i in range(img.shape[1])], 1)
- gx = F.conv2d(img, torch.Tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]]).view((1, 1, 3, 3)).to(img.device), padding=1)
- return gx
-
- def forward(self, depth_preds, shifts, mask, temp_features, pred_f=None):
-
- common_area_1_list = []
- common_area_2_list = []
-
- if self.focus_flatten:
- # only consider flatten place
- grad = kornia.filters.spatial_gradient(pred_f.detach())
- grad_x, grad_y = grad[:, :, 0, :, :], grad[:, :, 1, :, :]
- grad = torch.sqrt(grad_x ** 2 + grad_y ** 2)
- grad_ext = grad > 0.05 # over 5cm
- grad_ext = grad_ext.float()
- grad_blur = kornia.filters.gaussian_blur2d(grad_ext, (11, 11), (3, 3))
- grad_ext = grad_blur > 0 # over 5cm
- grad_ext = grad_blur == 0
- mask = torch.logical_and(mask, grad_ext)
-
-
- if self.target == "mix":
- ## for feature
- bs, c, h, w = depth_preds.shape
- split_depth = torch.split(depth_preds, bs//2, dim=0)
- split_mask = torch.split(F.interpolate(mask.float(), (384, 512)).bool(), bs//2, dim=0)
-
- feat_ori_list = []
- feat_shift_list = []
- multi_level_mask = []
-
- for idx, feature in enumerate(temp_features): # multi-level
- split_feat = torch.split(feature, bs//2, dim=0)
-
- _, _, h, w = split_feat[0].shape
- feat_ori_list.append(split_feat[0])
- feat_shift_list.append(split_feat[1])
-
- mask_ori_cur_scale = F.interpolate(split_mask[0].float(), (h, w)).bool()
- multi_level_mask.append(mask_ori_cur_scale)
-
- for idx_out, (feat_ori_cur_level, feat_shift_cur_level, mask_ori_cur_level) in enumerate(zip(feat_ori_list, feat_shift_list, multi_level_mask)): # iter multi-scale
- scale_factor = 2 ** (5 - idx_out)
- _, _, cur_scale_h, cur_scale_w = feat_ori_cur_level.shape
- scale_factor = int(384 / cur_scale_h)
-
- for idx_in, (feat_ori, feat_shift, mask_ori, shift_bs) in enumerate(zip(feat_ori_cur_level, feat_shift_cur_level, mask_ori_cur_level, shifts)): # iter bs (paired feat)
- c, _, _ = feat_ori.shape
- mask_ori = mask_ori.repeat(c, 1, 1)
- shift_h, shift_w = int(shift_bs[0] * (384/540) / scale_factor), int(shift_bs[1]* (512/960) / scale_factor)
-
- if shift_h >= 0 and shift_w >= 0:
- common_area_1 = feat_ori[:, shift_h:, shift_w:]
- common_area_2 = feat_shift[:, :-shift_h, :-shift_w]
- mask_common = mask_ori[:, shift_h:, shift_w:]
- elif shift_h >= 0 and shift_w <= 0:
- common_area_1 = feat_ori[:, shift_h:, :-abs(shift_w)]
- common_area_2 = feat_shift[:, :-shift_h, abs(shift_w):]
- mask_common = mask_ori[:, shift_h:, :-abs(shift_w)]
- elif shift_h <= 0 and shift_w <= 0:
- common_area_1 = feat_ori[:, :-abs(shift_h), :-abs(shift_w)]
- common_area_2 = feat_shift[:, abs(shift_h):, abs(shift_w):]
- mask_common = mask_ori[:, :-abs(shift_h), :-abs(shift_w)]
- elif shift_h <= 0 and shift_w >= 0:
- common_area_1 = feat_ori[:, :-abs(shift_h):, shift_w:]
- common_area_2 = feat_shift[:, abs(shift_h):, :-shift_w]
- mask_common = mask_ori[:, :-abs(shift_h):, shift_w:]
- else:
- print("can you really reach here?")
-
- common_area_masked_1 = common_area_1[mask_common].flatten()
- common_area_masked_2 = common_area_2[mask_common].flatten()
- common_area_1_list.append(common_area_masked_1)
- common_area_2_list.append(common_area_masked_2)
-
- common_area_1 = torch.cat(common_area_1_list)
- common_area_2 = torch.cat(common_area_2_list)
- if common_area_1.numel() == 0 or common_area_2.numel() == 0:
- consistency_loss = torch.Tensor([0]).squeeze()
- else:
- consistency_loss = F.mse_loss(common_area_1, common_area_2)
- consistency_loss_feat = consistency_loss
-
-
- common_area_1_list = []
- common_area_2_list = []
-
- ## for pred
- bs, c, h, w = depth_preds.shape
- split_depth = torch.split(depth_preds, bs//2, dim=0)
- split_mask = torch.split(mask, bs//2, dim=0)
-
- for shift, depth_ori, depth_shift, mask_ori, mask_shift in zip(shifts, split_depth[0], split_depth[1], split_mask[0], split_mask[1]):
- shift_h, shift_w = shift[0], shift[1]
- if shift_h >= 0 and shift_w >= 0:
- common_area_1 = depth_ori[:, shift_h:, shift_w:]
- common_area_2 = depth_shift[:, :-shift_h, :-shift_w]
- mask_common = mask_ori[:, shift_h:, shift_w:]
- # mask_debug = mask_shift[:, :-shift_h, :-shift_w]
- elif shift_h >= 0 and shift_w <= 0:
- common_area_1 = depth_ori[:, shift_h:, :-abs(shift_w)]
- common_area_2 = depth_shift[:, :-shift_h, abs(shift_w):]
- mask_common = mask_ori[:, shift_h:, :-abs(shift_w)]
- # mask_debug = mask_shift[:, :-shift_h, abs(shift_w):]
- elif shift_h <= 0 and shift_w <= 0:
- common_area_1 = depth_ori[:, :-abs(shift_h), :-abs(shift_w)]
- common_area_2 = depth_shift[:, abs(shift_h):, abs(shift_w):]
- mask_common = mask_ori[:, :-abs(shift_h), :-abs(shift_w)]
- # mask_debug = mask_shift[:, abs(shift_h):, abs(shift_w):]
- elif shift_h <= 0 and shift_w >= 0:
- common_area_1 = depth_ori[:, :-abs(shift_h):, shift_w:]
- common_area_2 = depth_shift[:, abs(shift_h):, :-shift_w]
- mask_common = mask_ori[:, :-abs(shift_h):, shift_w:]
- # mask_debug = mask_shift[:, abs(shift_h):, :-shift_w]
- else:
- print("can you really reach here?")
-
- common_area_1 = common_area_1[mask_common].flatten()
- common_area_2 = common_area_2[mask_common].flatten()
- common_area_1_list.append(common_area_1)
- common_area_2_list.append(common_area_2)
-
- common_area_1 = torch.cat(common_area_1_list)
- common_area_2 = torch.cat(common_area_2_list)
- if common_area_1.numel() == 0 or common_area_2.numel() == 0:
- consistency_loss = torch.Tensor([0]).squeeze()
- else:
- # pred_hist = torch.histc(common_area_1, bins=512, min=0, max=80)
- # gt_hist = torch.histc(common_area_2, bins=512, min=0, max=80)
-
- # pred_hist /= pred_hist.sum(dim=0, keepdim=True)
- # gt_hist /= gt_hist.sum(dim=0, keepdim=True)
-
- # # Compute cumulative histograms (CDF)
- # cdf_pred = torch.cumsum(pred_hist, dim=0)
- # cdf_gt = torch.cumsum(gt_hist, dim=0)
-
- # # Compute Earth Mover's Distance (EMD) between the CDFs
- # consistency_loss = torch.mean(torch.abs(cdf_pred - cdf_gt))
- consistency_loss = F.mse_loss(common_area_1, common_area_2)
- consistency_loss_pred = consistency_loss
-
- consistency_loss = consistency_loss_pred * self.wp + consistency_loss_feat
- return consistency_loss
-
- elif 'feat' in self.target:
- if self.mode == 'resize':
- bs, c, h, w = depth_preds.shape
- split_depth = torch.split(depth_preds, bs//2, dim=0)
- split_mask = torch.split(mask, bs//2, dim=0)
-
- feat_ori_list = []
- feat_shift_list = []
-
- for idx, feature in enumerate(temp_features): # multi-level
- if idx < 4:
- continue
-
- split_feat = torch.split(feature, bs//2, dim=0)
- f = F.interpolate(split_feat[0], (h, w), mode='bilinear', align_corners=True)
- feat_ori_list.append(f)
- f = F.interpolate(split_feat[1], (h, w), mode='bilinear', align_corners=True)
- feat_shift_list.append(f)
-
-
- for idx_out, (feat_ori_cur_level, feat_shift_cur_level) in enumerate(zip(feat_ori_list, feat_shift_list)): # iter multi-scale
- scale_factor = 2 ** (5 - idx_out)
-
- for idx_in, (feat_ori, feat_shift, mask_ori, shift_bs) in enumerate(zip(feat_ori_cur_level, feat_shift_cur_level, split_mask[0], shifts)): # iter bs (paired feat)
- c, h, w = feat_ori.shape
- mask_ori = mask_ori.repeat(c, 1, 1)
- shift_h, shift_w = shift_bs[0], shift_bs[1]
-
- if shift_h >= 0 and shift_w >= 0:
- common_area_1 = feat_ori[:, shift_h:, shift_w:]
- common_area_2 = feat_shift[:, :-shift_h, :-shift_w]
- mask_common = mask_ori[:, shift_h:, shift_w:]
- elif shift_h >= 0 and shift_w <= 0:
- common_area_1 = feat_ori[:, shift_h:, :-abs(shift_w)]
- common_area_2 = feat_shift[:, :-shift_h, abs(shift_w):]
- mask_common = mask_ori[:, shift_h:, :-abs(shift_w)]
- elif shift_h <= 0 and shift_w <= 0:
- common_area_1 = feat_ori[:, :-abs(shift_h), :-abs(shift_w)]
- common_area_2 = feat_shift[:, abs(shift_h):, abs(shift_w):]
- mask_common = mask_ori[:, :-abs(shift_h), :-abs(shift_w)]
- elif shift_h <= 0 and shift_w >= 0:
- common_area_1 = feat_ori[:, :-abs(shift_h):, shift_w:]
- common_area_2 = feat_shift[:, abs(shift_h):, :-shift_w]
- mask_common = mask_ori[:, :-abs(shift_h):, shift_w:]
- else:
- print("can you really reach here?")
-
- common_area_masked_1 = common_area_1[mask_common].flatten()
- common_area_masked_2 = common_area_2[mask_common].flatten()
- # common_area_masked_1 = common_area_1.flatten()
- # common_area_masked_2 = common_area_2.flatten()
- common_area_1_list.append(common_area_masked_1)
- common_area_2_list.append(common_area_masked_2)
-
- common_area_1 = torch.cat(common_area_1_list)
- common_area_2 = torch.cat(common_area_2_list)
- if common_area_1.numel() == 0 or common_area_2.numel() == 0:
- consistency_loss = torch.Tensor([0]).squeeze()
- else:
- consistency_loss = F.mse_loss(common_area_1, common_area_2)
-
- return consistency_loss
-
-
- else:
- bs, c, h, w = depth_preds.shape
- split_depth = torch.split(depth_preds, bs//2, dim=0)
- mask = F.interpolate(mask.float(), (384, 512)).bool() # back to 384, 512
- split_mask = torch.split(mask, bs//2, dim=0)
-
- feat_ori_list = []
- feat_shift_list = []
- multi_level_mask = []
-
- for idx, feature in enumerate(temp_features): # multi-level
- split_feat = torch.split(feature, bs//2, dim=0)
-
- _, _, h, w = split_feat[0].shape
- feat_ori_list.append(split_feat[0])
- feat_shift_list.append(split_feat[1])
-
- mask_ori_cur_scale = F.interpolate(split_mask[0].float(), (h, w)).bool()
- multi_level_mask.append(mask_ori_cur_scale)
-
- for idx_out, (feat_ori_cur_level, feat_shift_cur_level, mask_ori_cur_level) in enumerate(zip(feat_ori_list, feat_shift_list, multi_level_mask)): # iter multi-scale
- scale_factor = 2 ** (5 - idx_out)
- _, _, cur_scale_h, cur_scale_w = feat_ori_cur_level.shape
- scale_factor = int(384 / cur_scale_h)
-
- for idx_in, (feat_ori, feat_shift, mask_ori, shift_bs) in enumerate(zip(feat_ori_cur_level, feat_shift_cur_level, mask_ori_cur_level, shifts)): # iter bs (paired feat)
- c, _, _ = feat_ori.shape
- mask_ori = mask_ori.repeat(c, 1, 1)
- shift_h, shift_w = int(shift_bs[0] * (384/540) / scale_factor), int(shift_bs[1]* (512/960) / scale_factor)
-
- if shift_h >= 0 and shift_w >= 0:
- common_area_1 = feat_ori[:, shift_h:, shift_w:]
- common_area_2 = feat_shift[:, :-shift_h, :-shift_w]
- mask_common = mask_ori[:, shift_h:, shift_w:]
- elif shift_h >= 0 and shift_w <= 0:
- common_area_1 = feat_ori[:, shift_h:, :-abs(shift_w)]
- common_area_2 = feat_shift[:, :-shift_h, abs(shift_w):]
- mask_common = mask_ori[:, shift_h:, :-abs(shift_w)]
- elif shift_h <= 0 and shift_w <= 0:
- common_area_1 = feat_ori[:, :-abs(shift_h), :-abs(shift_w)]
- common_area_2 = feat_shift[:, abs(shift_h):, abs(shift_w):]
- mask_common = mask_ori[:, :-abs(shift_h), :-abs(shift_w)]
- elif shift_h <= 0 and shift_w >= 0:
- common_area_1 = feat_ori[:, :-abs(shift_h):, shift_w:]
- common_area_2 = feat_shift[:, abs(shift_h):, :-shift_w]
- mask_common = mask_ori[:, :-abs(shift_h):, shift_w:]
- else:
- print("can you really reach here?")
-
- common_area_masked_1 = common_area_1[mask_common].flatten()
- common_area_masked_2 = common_area_2[mask_common].flatten()
- common_area_1_list.append(common_area_masked_1)
- common_area_2_list.append(common_area_masked_2)
-
- common_area_1 = torch.cat(common_area_1_list)
- common_area_2 = torch.cat(common_area_2_list)
- if common_area_1.numel() == 0 or common_area_2.numel() == 0:
- consistency_loss = torch.Tensor([0]).squeeze()
- else:
- consistency_loss = F.mse_loss(common_area_1, common_area_2)
- return consistency_loss
-
- elif self.target == 'pred':
- bs, c, h, w = depth_preds.shape
- split_depth = torch.split(depth_preds, bs//2, dim=0)
- split_mask = torch.split(mask, bs//2, dim=0)
-
- for shift, depth_ori, depth_shift, mask_ori, mask_shift in zip(shifts, split_depth[0], split_depth[1], split_mask[0], split_mask[1]):
- shift_h, shift_w = shift[0], shift[1]
- if shift_h >= 0 and shift_w >= 0:
- common_area_1 = depth_ori[:, shift_h:, shift_w:]
- common_area_2 = depth_shift[:, :-shift_h, :-shift_w]
- mask_common = mask_ori[:, shift_h:, shift_w:]
- # mask_debug = mask_shift[:, :-shift_h, :-shift_w]
- elif shift_h >= 0 and shift_w <= 0:
- common_area_1 = depth_ori[:, shift_h:, :-abs(shift_w)]
- common_area_2 = depth_shift[:, :-shift_h, abs(shift_w):]
- mask_common = mask_ori[:, shift_h:, :-abs(shift_w)]
- # mask_debug = mask_shift[:, :-shift_h, abs(shift_w):]
- elif shift_h <= 0 and shift_w <= 0:
- common_area_1 = depth_ori[:, :-abs(shift_h), :-abs(shift_w)]
- common_area_2 = depth_shift[:, abs(shift_h):, abs(shift_w):]
- mask_common = mask_ori[:, :-abs(shift_h), :-abs(shift_w)]
- # mask_debug = mask_shift[:, abs(shift_h):, abs(shift_w):]
- elif shift_h <= 0 and shift_w >= 0:
- common_area_1 = depth_ori[:, :-abs(shift_h):, shift_w:]
- common_area_2 = depth_shift[:, abs(shift_h):, :-shift_w]
- mask_common = mask_ori[:, :-abs(shift_h):, shift_w:]
- # mask_debug = mask_shift[:, abs(shift_h):, :-shift_w]
- else:
- print("can you really reach here?")
-
- common_area_1 = common_area_1[mask_common].flatten()
- common_area_2 = common_area_2[mask_common].flatten()
- common_area_1_list.append(common_area_1)
- common_area_2_list.append(common_area_2)
-
- common_area_1 = torch.cat(common_area_1_list)
- common_area_2 = torch.cat(common_area_2_list)
- if common_area_1.numel() == 0 or common_area_2.numel() == 0:
- consistency_loss = torch.Tensor([0]).squeeze()
- else:
- # pred_hist = torch.histc(common_area_1, bins=512, min=0, max=80)
- # gt_hist = torch.histc(common_area_2, bins=512, min=0, max=80)
-
- # pred_hist /= pred_hist.sum(dim=0, keepdim=True)
- # gt_hist /= gt_hist.sum(dim=0, keepdim=True)
-
- # # Compute cumulative histograms (CDF)
- # cdf_pred = torch.cumsum(pred_hist, dim=0)
- # cdf_gt = torch.cumsum(gt_hist, dim=0)
-
- # # Compute Earth Mover's Distance (EMD) between the CDFs
- # consistency_loss = torch.mean(torch.abs(cdf_pred - cdf_gt))
- consistency_loss = F.mse_loss(common_area_1, common_area_2)
-
- return consistency_loss
-
- else:
- raise NotImplementedError
\ No newline at end of file
diff --git a/zoedepth/trainers/loss_sample.py b/zoedepth/trainers/loss_sample.py
deleted file mode 100644
index b2cf3c146654b491300eaa3650b73387cb8632c9..0000000000000000000000000000000000000000
--- a/zoedepth/trainers/loss_sample.py
+++ /dev/null
@@ -1,173 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Zhenyu Li
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.cuda.amp as amp
-import numpy as np
-
-
-KEY_OUTPUT = 'metric_depth'
-
-
-def extract_key(prediction, key):
- if isinstance(prediction, dict):
- return prediction[key]
- return prediction
-
-
-# Main loss function used for ZoeDepth. Copy/paste from AdaBins repo (https://github.com/shariqfarooq123/AdaBins/blob/0952d91e9e762be310bb4cd055cbfe2448c0ce20/loss.py#L7)
-class SILogLoss(nn.Module):
- """SILog loss (pixel-wise)"""
- def __init__(self, beta=0.15):
- super(SILogLoss, self).__init__()
- self.name = 'SILog'
- self.beta = beta
-
- def forward(self, input, target, mask=None):
- input = extract_key(input, KEY_OUTPUT)
-
- if mask is not None:
- input_filtered = input[mask]
- target_filtered = target[mask]
-
- with amp.autocast(enabled=False): # amp causes NaNs in this loss function
- alpha = 1e-7
- g = torch.log(input_filtered + alpha) - torch.log(target_filtered + alpha)
- Dg = torch.var(g) + self.beta * torch.pow(torch.mean(g), 2)
- loss = 10 * torch.sqrt(Dg)
-
- if torch.isnan(loss):
- print("Nan SILog loss")
- print("input:", input.shape)
- print("target:", target.shape)
- print("G", torch.sum(torch.isnan(g)))
- print("Input min max", torch.min(input), torch.max(input))
- print("Target min max", torch.min(target), torch.max(target))
- print("Dg", torch.isnan(Dg))
- print("loss", torch.isnan(loss))
-
- return loss
-
-
-
-import torch
-import torch.nn.functional as F
-
-def gaussian(mu, sigma, labels):
- return torch.exp(-0.5*(mu-labels)** 2/ sigma** 2)/sigma
-
-def laplacian(mu, b, labels):
- # a = torch.abs(mu-labels)/b
- # print("1 isnan: {}".format(torch.isnan(a).any()))
- # print("1 isinf: {}".format(torch.isinf(a).any()))
- # a = torch.exp(-(torch.abs(mu-labels)/b))
- # print(a)
-
- # print("1 isnan: {}".format(torch.isnan(a).any()))
- # print("1 isinf: {}".format(torch.isinf(a).any()))
- return 0.5 * torch.exp(-(torch.abs(mu-labels)/b))/b
-
-def distribution(mu, sigma, labels, dist="gaussian"):
- return gaussian(mu, sigma, labels) if dist=="gaussian" else \
- laplacian(mu, sigma, labels)
-
-def bimodal_loss(mu0, mu1, sigma0, sigma1, w0, w1, labels, dist="gaussian"):
- # first_term = w0 * distribution(mu0, sigma0, labels, dist)
- # print(first_term)
- # print("f isnan: {}".format(torch.isnan(first_term).any()))
- # print("f isinf: {}".format(torch.isinf(first_term).any()))
- # second_term = w1 * distribution(mu1, sigma1, labels, dist)
- # print(second_term)
- # print("s isnan: {}".format(torch.isnan(second_term).any()))
- # print("s isinf: {}".format(torch.isinf(second_term).any()))
-
- loss = w0 * distribution(mu0, sigma0, labels, dist) + w1 * distribution(mu1, sigma1, labels, dist)
- # loss = torch.clamp(loss, min=1e-12)
- # print(loss)
- return - torch.log(loss)
-
-def unimodal_loss(mu, sigma, labels):
- return torch.abs(mu - labels)/sigma + torch.log(sigma)
-
-def smooth_l1_loss(preds, labels, reduce=None):
- return F.smooth_l1_loss(preds, labels, reduce=reduce)
-
-def l1_loss(preds, labels, reduce=None):
- return F.l1_loss(preds, labels, reduce=reduce)
-
-
-class DistributionLoss(nn.Module):
- def __init__(self, max_depth):
- super(DistributionLoss, self).__init__()
- self.name = 'DistributionLoss'
- self.max_depth = max_depth
-
- def forward(self, input, target, mask=None, dist='biLaplacian'):
-
-
- mu0 = input['mu0']
- mu1 = input['mu1']
- sigma0 = input['sigma0']
- sigma1 = input['sigma1']
- pi0 = input['pi0']
- pi1 = input['pi1']
-
- pred_mask = (pi0 / sigma0 > pi1 / sigma1).float()
- pred_depth = (mu0 * pred_mask + mu1 * (1. - pred_mask))
- pred_metric_depth = (1 - pred_depth) * self.max_depth
-
-
- if mask is not None:
- mu0 = mu0[mask]
- mu1 = mu1[mask]
- sigma0 = sigma0[mask]
- sigma1 = sigma1[mask]
- pi0 = pi0[mask]
- pi1 = pi1[mask]
-
- # real_input = real_depth[mask]
-
- real_input = mu0
- pred_metric_depth = pred_metric_depth[mask]
- record_target = target[mask]
-
-
- target_filtered = 1 - target[mask] / self.max_depth
- bi_loss = bimodal_loss(mu0, mu1, sigma0, sigma1, pi0, pi1, target_filtered, dist=dist).mean()
- # print(bi_loss)
-
- alpha = 1e-7
- beta = 0.15
- g = torch.log(real_input + alpha) - torch.log(record_target + alpha)
- Dg = torch.var(g) + beta * torch.pow(torch.mean(g), 2)
- sig_loss = 10 * torch.sqrt(Dg)
- # print(sig_loss)
-
- return bi_loss, sig_loss
-
-
-
-
\ No newline at end of file
diff --git a/zoedepth/trainers/zoedepth_custom_trainer.py b/zoedepth/trainers/zoedepth_custom_trainer.py
deleted file mode 100644
index e4cd353a1a304e4cd13b82c0fe539862c9e2d85d..0000000000000000000000000000000000000000
--- a/zoedepth/trainers/zoedepth_custom_trainer.py
+++ /dev/null
@@ -1,722 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Zhenyu Li
-
-# This file is partly inspired from ZoeDepth (https://github.com/isl-org/ZoeDepth/blob/main/zoedepth/trainers/zoedepth_trainer.py); author: Shariq Farooq Bhat
-
-import os
-
-import torch
-import torch.cuda.amp as amp
-import torch.nn as nn
-
-from zoedepth.trainers.loss_sample import SILogLoss, DistributionLoss
-from zoedepth.trainers.loss import SILogLoss as DenseSILogLoss
-from zoedepth.trainers.loss import BudgetConstraint, HistogramMatchingLoss, SSIM, ConsistencyLoss
-from zoedepth.utils.config import DATASETS_CONFIG
-from zoedepth.utils.misc import compute_metrics
-from zoedepth.data.preprocess import get_black_border
-
-from .base_trainer import BaseTrainer, is_rank_zero, colors, flatten
-from torchvision import transforms
-from PIL import Image
-import numpy as np
-
-import wandb
-import uuid
-from tqdm import tqdm
-from datetime import datetime as dt
-import torch.distributed as dist
-
-import copy
-from zoedepth.utils.misc import generatemask
-import torch.optim as optim
-
-class Trainer(BaseTrainer):
- def __init__(self, config, model, train_loader, test_loader=None, device=None):
- self.addf = config.get("addf", False)
- self.lazy_epoch = -1
- self.boostingdepth = config.get("boostingdepth", False)
- super().__init__(config, model, train_loader,
- test_loader=test_loader, device=device)
- self.device = device
- self.silog_loss = SILogLoss(beta=config.get("beta", 0.15))
- self.dense_silog_loss = DenseSILogLoss(beta=config.get("beta", 0.15))
- print("sigloss's beta is set to {}".format(config.get("beta", 0.15)))
- self.scaler = amp.GradScaler(enabled=self.config.use_amp)
- self.distribution_loss = DistributionLoss(max_depth=self.config.max_depth)
- self.sampled_training = config.get("sampled_training", False)
- self.sec_stage = config.get("sec_stage", False)
- self.multi_consistency = config.get("multi_consistency", False)
- self.use_blur = config.get("use_blur", False)
- self.dynamic = config.get("dynamic", False)
- if self.dynamic:
- self.dynamic_unupdate_rate = config.get("dynamic_unupdate_rate", 0.0)
- self.budget_loss = BudgetConstraint(loss_mu=0.0, flops_all=21552.5684, warm_up=True)
- self.use_scale_loss = config.get("use_scale_loss", False)
- if self.use_scale_loss:
- if config.get("scale_type", "ssim"):
- self.scale_loss = SSIM(window_size=config.get("window_size", int(11)))
- else:
- self.scale_loss = HistogramMatchingLoss(min_depth=self.config.min_depth, max_depth=self.config.max_depth)
- self.scale_target = config.get("scale_target", None)
- self.consistency_training = config.get("consistency_training", False)
- if self.consistency_training:
- self.consistency_target = config.get("consistency_target", None)
- self.consistency_loss = ConsistencyLoss(self.consistency_target, config.get("focus_flatten", False), config.get("w_p", 1.0))
- print("current weight for consistency loss is {}. focus_flatten is {}. w_p is {}".format(self.config.w_consistency, config.get("focus_flatten", False), config.get("w_p", 1.0)))
-
-
-
- def train_on_batch(self, batch, train_step, step_rate):
- """
- Expects a batch of images and depth as input
- batch["image"].shape : batch_size, c, h, w
- batch["depth"].shape : batch_size, 1, h, w
- """
-
- images, depths_gt = batch['image'].to(self.device), batch['depth'].to(self.device)
- image_raw = batch.get("image_raw", None)
- if image_raw is not None:
- image_raw = image_raw.to(self.device)
- sample_points = None
- if self.sampled_training:
- sample_points = batch['sample_points'].to(self.device)
- bbox = batch.get("bbox", None)
- if bbox is not None:
- bbox = bbox.to(self.device)
- bbox_raw = batch.get("bbox_raw", None)
- if bbox_raw is not None:
- bbox_raw = bbox_raw.to(self.device)
- depth_raw = batch.get("depth_raw", None)
- if depth_raw is not None:
- depth_raw = depth_raw.to(self.device)
- crop_area = batch.get("crop_area", None)
- if crop_area is not None:
- crop_area = crop_area.to(self.device)
- shift = batch.get("shift", None)
- if shift is not None:
- shift = shift.to(self.device)
-
-
- dataset = batch['dataset'][0]
-
- b, c, h, w = images.size()
- mask = batch["mask"].to(self.device).to(torch.bool)
- sample_mask = batch.get("sample_mask", None)
- if sample_mask is not None:
- sample_mask = sample_mask.to(self.device).to(torch.bool)
- mask_raw = batch.get("mask_raw", None)
- if mask_raw is not None:
- mask_raw = mask_raw.to(self.device).to(torch.bool)
-
- losses = {}
-
- with amp.autocast(enabled=self.config.use_amp):
- if self.sampled_training:
- output = self.model(images, sample_points, mode='train', image_raw=image_raw, bbox=bbox, depth_raw=depth_raw, crop_area=crop_area, shift=shift, bbox_raw=bbox_raw)
- else:
- output = self.model(images, None, mode='train', image_raw=image_raw, bbox=bbox, depth_raw=depth_raw, crop_area=crop_area, shift=shift, bbox_raw=bbox_raw)
-
-
- if self.boostingdepth:
- if self.lazy_epoch < self.epoch:
- output.update_learning_rate()
- self.lazy_epoch = self.epoch
- input_dict = dict()
- input_dict['data_gtfake'] = depths_gt
- output.set_input_train_gt(input_dict)
- output.optimize_parameters()
- pred_depths = output.fake_B
- pred = output.fake_B
- # print(torch.min(pred), torch.max(pred))
- losses = output.get_current_losses()
-
- else:
- pred_depths = output['metric_depth']
-
- if self.sampled_training:
- sampled_depth_gt = sample_points[:, :, -1].float().unsqueeze(dim=-1)
- sampled_depth_gt = sampled_depth_gt.permute(0, 2, 1)
-
- if self.config.get("representation", "") == 'biLaplacian':
- # only for sampled training for now
- l_dist, l_si = self.distribution_loss(output, sampled_depth_gt, mask=sample_mask)
- loss = self.config.w_dist * l_dist + self.config.w_si * l_si
- losses['distribution_loss'] = l_dist
- losses['sigloss'] = l_si
-
- if self.multi_consistency:
- coarse, fine = output['coarse_depth_pred'], output['fine_depth_pred']
- l_si_f = self.dense_silog_loss(
- fine, depths_gt, mask=mask, interpolate=True, return_interpolated=False)
- l_si_c = self.dense_silog_loss(
- coarse, depth_raw, mask=mask_raw, interpolate=True, return_interpolated=False)
-
- losses['sigloss_f'] = l_si_f
- losses['l_si_c'] = l_si_c
- loss += self.config.w_si * (l_si_f + l_si_c)
-
- else:
- if self.sampled_training:
- l_si = self.silog_loss(
- pred_depths, sampled_depth_gt, mask=sample_mask)
- loss = self.config.w_si * l_si
- losses[self.silog_loss.name] = l_si
-
- if self.multi_consistency:
- coarse, fine = output['coarse_depth_pred'], output['fine_depth_pred']
- l_si_f = self.dense_silog_loss(
- fine, depths_gt, mask=mask, interpolate=True, return_interpolated=False)
- l_si_c = self.dense_silog_loss(
- coarse, depth_raw, mask=mask_raw, interpolate=True, return_interpolated=False)
-
- losses['sigloss_f'] = l_si_f
- losses['l_si_c'] = l_si_c
- loss += self.config.w_si * (l_si_f + l_si_c)
-
- else:
- if self.multi_consistency:
- #### here here here
- pred_depths, coarse, fine = output['metric_depth'], output['coarse_depth_pred'], output['fine_depth_pred']
-
- if self.consistency_training:
- depths_gt = torch.split(depths_gt, 1, dim=1)
- depths_gt = torch.cat(depths_gt, dim=0)
- mask = torch.split(mask, 1, dim=-1)
- mask = torch.cat(mask, dim=0).permute(0, 3, 1, 2)
- mask_raw = torch.cat([mask_raw, mask_raw], dim=0)
- depth_raw = torch.cat([depth_raw, depth_raw], dim=0)
- temp_features = output.get('temp_features', None)
-
-
- l_si_1, pred = self.dense_silog_loss(
- pred_depths, depths_gt, mask=mask, interpolate=True, return_interpolated=True)
- l_si_f, pred_f = self.dense_silog_loss(
- fine, depths_gt, mask=mask, interpolate=True, return_interpolated=True)
- l_si_c = self.dense_silog_loss(
- coarse, depth_raw, mask=mask_raw, interpolate=True, return_interpolated=False)
-
- losses[self.silog_loss.name] = l_si_1
- losses['sigloss_f'] = l_si_f
- losses['l_si_c'] = l_si_c
- # loss = l_si_1 + l_si_f + l_si_c
- loss = l_si_1
-
- if self.consistency_training:
- try:
- # depths_gt? pred_f?
- l_consistency = self.consistency_loss(pred, shift, mask, temp_features, pred_f=depths_gt) # use the resized pred
- except RuntimeError as e:
- print(e)
- print("some runtime error here! Hack with 0")
- l_consistency = torch.Tensor([0]).squeeze()
-
- losses[self.consistency_loss.name] = l_consistency
- loss += l_consistency * self.config.w_consistency
-
- else:
-
- l_si, pred = self.dense_silog_loss(
- pred_depths, depths_gt, mask=mask, interpolate=True, return_interpolated=True)
-
- loss = self.config.w_si * l_si
- losses[self.silog_loss.name] = l_si
-
- if self.dynamic:
- if step_rate > self.dynamic_unupdate_rate:
- warm_up_rate = min(1.0, (step_rate - self.dynamic_unupdate_rate) / 0.02)
- flop_cost = self.budget_loss(output['all_cell_flops'], warm_up_rate=warm_up_rate)
- loss += self.config.w_flop * flop_cost
- losses['flop_loss'] = flop_cost
- else:
- flop_cost = self.budget_loss(output['all_cell_flops'], warm_up_rate=1)
- loss += 0 * flop_cost
- losses['flop_loss'] = flop_cost
-
- if self.use_scale_loss:
- if self.scale_target == 'coarse':
- h_loss = self.scale_loss(pred_depths, output['coarse_depth_pred_roi'], mask, interpolate=True)
- else:
- h_loss = self.scale_loss(pred_depths, depths_gt, mask, interpolate=True)
- loss += self.config.w_scale * h_loss
- losses['scale_loss'] = h_loss
-
-
- # self.scaler.scale(loss).backward()
-
- # if self.config.clip_grad > 0:
- # self.scaler.unscale_(self.optimizer)
- # nn.utils.clip_grad_norm_(
- # self.model.parameters(), self.config.clip_grad)
-
- # self.scaler.step(self.optimizer)
-
- # self.scaler.update()
- # self.optimizer.zero_grad()
-
-
- self.scaler.scale(loss).backward()
-
- if self.config.clip_grad > 0:
- self.scaler.unscale_(self.optimizer)
- nn.utils.clip_grad_norm_(
- self.model.parameters(), self.config.clip_grad)
-
- self.scaler.step(self.optimizer)
-
- self.scaler.update()
- self.optimizer.zero_grad()
-
-
- if self.should_log and (self.step % int(self.config.log_images_every * self.iters_per_epoch)) == 0:
- if self.config.get("debug", False):
- pred = nn.functional.interpolate(
- pred[0:1], depths_gt.shape[-2:], mode='bilinear', align_corners=True)[0]
- import matplotlib.pyplot as plt
- plt.imshow(pred.squeeze().detach().cpu().numpy())
- plt.savefig('debug.png')
- pass
- else:
- pred = nn.functional.interpolate(
- pred[0:1], depths_gt.shape[-2:], mode='bilinear', align_corners=True)[0]
- depths_gt[torch.logical_not(mask)] = DATASETS_CONFIG[dataset]['max_depth']
- if self.consistency_training:
- split_images = torch.split(images, 3, dim=1)
- images = torch.cat(split_images, dim=0)
- self.log_images(rgb={"Input": images[0, ...]}, depth={"GT": depths_gt[0], "PredictedMono": pred}, prefix="Train",
- min_depth=DATASETS_CONFIG[dataset]['min_depth'], max_depth=DATASETS_CONFIG[dataset]['max_depth'])
-
- return losses
-
- @torch.no_grad()
- def eval_infer(self, x, image_raw, bboxs=None, crop_area=None, dataset='u4k', bbox_raw=None):
- m = self.model.module if self.config.multigpu else self.model
-
- if dataset == 'u4k':
- base_h = 540
- base_w = 960
- elif dataset == 'gta':
- base_h = 270
- base_w = 480
- elif dataset == 'nyu':
- base_h = 120 * 2
- base_w = 160 * 2
- else:
- raise NotImplementedError
-
- if dataset == 'nyu':
- if self.sec_stage:
- images_crops = torch.split(x, 3, dim=1)
- bboxs_list = torch.split(bboxs, 1, dim=1)
- crop_areas = torch.split(crop_area, 1, dim=1)
-
- pred_depth_crops = []
- for i, (img, bbox, crop_area) in enumerate(zip(images_crops, bboxs_list, crop_areas)):
- with amp.autocast(enabled=self.config.use_amp):
- if i == 0:
- out_dict = m(img, mode='eval', image_raw=image_raw, bbox=bbox[0], crop_area=crop_area, bbox_raw=bbox_raw[:, i, :] if bbox_raw is not None else None)
- # whole_depth_pred = out_dict['coarse_depth_pred']
- pred_depth_crop = out_dict['metric_depth']
- else:
- pred_depth_crop = m(img, mode='eval', image_raw=image_raw, bbox=bbox[0], crop_area=crop_area, bbox_raw=bbox_raw[:, i, :] if bbox_raw is not None else None)['metric_depth']
-
- pred_depth_crop = nn.functional.interpolate(
- pred_depth_crop, (base_h, base_w), mode='bilinear', align_corners=True)
- pred_depth_crops.append(pred_depth_crop)
-
- x_start, y_start = [0, base_h], [0, base_w]
- pred_depth = torch.zeros((base_h*2, base_w*2)).cuda()
- inner_idx = 0
- for ii, x in enumerate(x_start):
- for jj, y in enumerate(y_start):
- if self.use_blur:
- pred_depth[x: x+base_h, y: y+base_w] = pred_depth_crops[inner_idx].squeeze() # do not care about boundry during validation
- else:
- pred_depth[x: x+base_h, y: y+base_w] = pred_depth_crops[inner_idx].squeeze()
- inner_idx += 1
- pred_depth = pred_depth.squeeze(dim=0)
-
- else:
-
- with amp.autocast(enabled=self.config.use_amp):
- pred_depth = m(x, mode='eval', image_raw=image_raw)['metric_depth']
-
- else:
- if self.sec_stage:
- images_crops = torch.split(x, 3, dim=1)
- bboxs_list = torch.split(bboxs, 1, dim=1)
- crop_areas = torch.split(crop_area, 1, dim=1)
-
- pred_depth_crops = []
- for i, (img, bbox, crop_area) in enumerate(zip(images_crops, bboxs_list, crop_areas)):
- with amp.autocast(enabled=self.config.use_amp):
- if i == 0:
- out_dict = m(img, mode='eval', image_raw=image_raw, bbox=bbox[0], crop_area=crop_area, bbox_raw=bbox_raw[:, i, :] if bbox_raw is not None else None)
- # whole_depth_pred = out_dict['coarse_depth_pred']
- pred_depth_crop = out_dict['metric_depth']
- else:
- pred_depth_crop = m(img, mode='eval', image_raw=image_raw, bbox=bbox[0], crop_area=crop_area, bbox_raw=bbox_raw[:, i, :] if bbox_raw is not None else None)['metric_depth']
-
- pred_depth_crop = nn.functional.interpolate(
- pred_depth_crop, (base_h, base_w), mode='bilinear', align_corners=True)
- pred_depth_crops.append(pred_depth_crop)
-
- x_start, y_start = [0, base_h], [0, base_w]
- pred_depth = torch.zeros((base_h*2, base_w*2)).cuda()
- inner_idx = 0
- for ii, x in enumerate(x_start):
- for jj, y in enumerate(y_start):
- if self.use_blur:
- pred_depth[x: x+base_h, y: y+base_w] = pred_depth_crops[inner_idx].squeeze() # do not care about boundry during validation
- else:
- pred_depth[x: x+base_h, y: y+base_w] = pred_depth_crops[inner_idx].squeeze()
- inner_idx += 1
- pred_depth = pred_depth.squeeze(dim=0)
-
- else:
-
- with amp.autocast(enabled=self.config.use_amp):
- pred_depth = m(x, mode='eval', image_raw=image_raw)['metric_depth']
-
- return pred_depth
-
- @torch.no_grad()
- def crop_aware_infer(self, x, image_raw):
- # if we are not avoiding the black border, we can just use the normal inference
- if not self.config.get("avoid_boundary", False):
- return self.eval_infer(x)
-
- # otherwise, we need to crop the image to avoid the black border
- # For now, this may be a bit slow due to converting to numpy and back
- # We assume no normalization is done on the input image
-
- # get the black border
- assert x.shape[0] == 1, "Only batch size 1 is supported for now"
- x_pil = transforms.ToPILImage()(x[0].cpu())
- x_np = np.array(x_pil, dtype=np.uint8)
- black_border_params = get_black_border(x_np)
- top, bottom, left, right = black_border_params.top, black_border_params.bottom, black_border_params.left, black_border_params.right
- x_np_cropped = x_np[top:bottom, left:right, :]
- x_cropped = transforms.ToTensor()(Image.fromarray(x_np_cropped))
-
- # run inference on the cropped image
- pred_depths_cropped = self.eval_infer(x_cropped.unsqueeze(0).to(self.device))
-
- # resize the prediction to x_np_cropped's size
- pred_depths_cropped = nn.functional.interpolate(
- pred_depths_cropped, size=(x_np_cropped.shape[0], x_np_cropped.shape[1]), mode="bilinear", align_corners=False)
-
-
- # pad the prediction back to the original size
- pred_depths = torch.zeros((1, 1, x_np.shape[0], x_np.shape[1]), device=pred_depths_cropped.device, dtype=pred_depths_cropped.dtype)
- pred_depths[:, :, top:bottom, left:right] = pred_depths_cropped
-
- return pred_depths
-
- def validate_on_batch(self, batch, val_step):
- images = batch['image'].to(self.device)
- depths_gt = batch['depth'].to(self.device)
- dataset = batch['dataset'][0]
- image_raw = batch['image_raw'].to(self.device)
- mask = batch["mask"].to(self.device)
- disp_gt_edges = batch['disp_gt_edges'].squeeze().numpy()
- bboxs = batch.get("bbox", None)
- if bboxs is not None:
- bboxs = bboxs.to(self.device)
- bbox_raw = batch.get("bbox_raw", None)
- if bbox_raw is not None:
- bbox_raw = bbox_raw.to(self.device)
- crop_area = batch.get("crop_area", None)
- if crop_area is not None:
- crop_area = crop_area.to(self.device)
-
- if 'has_valid_depth' in batch:
- if not batch['has_valid_depth']:
- return None, None
-
- depths_gt = depths_gt.squeeze().unsqueeze(0).unsqueeze(0)
- mask = mask.squeeze().unsqueeze(0).unsqueeze(0)
- # if dataset == 'nyu':
- # pred_depths = self.crop_aware_infer(images, image_raw)
- # else:
- # pred_depths = self.eval_infer(images, image_raw, bboxs, crop_area, dataset, bbox_raw)
- pred_depths = self.eval_infer(images, image_raw, bboxs, crop_area, dataset, bbox_raw)
- pred_depths = pred_depths.squeeze().unsqueeze(0).unsqueeze(0)
- # print(pred_depths.shape) # torch.Size([1, 1, 2160, 3840])
- # print(depths_gt.shape) # torch.Size([1, 1, 2160, 3840])
-
- with amp.autocast(enabled=self.config.use_amp):
- if self.sampled_training:
- l_depth = self.silog_loss(
- pred_depths, depths_gt, mask=mask.to(torch.bool))
- else:
- l_depth = self.dense_silog_loss(
- pred_depths, depths_gt, mask=mask.to(torch.bool), interpolate=True)
-
-
-
- metrics = compute_metrics(depths_gt, pred_depths, disp_gt_edges=disp_gt_edges, **self.config)
- losses = {f"{self.silog_loss.name}": l_depth.item()}
-
- if self.should_log and self.config.get("debug", False):
- print(metrics)
- if val_step in [21, 27] and self.should_log:
- if self.config.get("debug", False):
- pass
- else:
- if self.sec_stage:
- log_rgb = image_raw
- else:
- log_rgb = images
-
- scale_pred = nn.functional.interpolate(
- pred_depths[0:1], depths_gt.shape[-2:], mode='bilinear', align_corners=True)[0]
- depths_gt[torch.logical_not(mask)] = DATASETS_CONFIG[dataset]['max_depth']
- self.log_images(rgb={"Input": log_rgb[0]}, depth={"GT": depths_gt[0], "PredictedMono": scale_pred}, prefix="Test",
- min_depth=DATASETS_CONFIG[dataset]['min_depth'], max_depth=DATASETS_CONFIG[dataset]['max_depth'])
-
- return metrics, losses
-
-
-
- def train(self):
- print(f"Training {self.config.name}")
- if self.config.uid is None:
- self.config.uid = str(uuid.uuid4()).split('-')[-1]
- run_id = f"{dt.now().strftime('%d-%h_%H-%M')}-{self.config.uid}"
- self.config.run_id = run_id
- self.config.experiment_id = f"{self.config.wandb_start}_{self.config.name}{self.config.version_name}_{run_id}"
- self.should_write = ((not self.config.distributed)
- or self.config.rank == 0)
- self.should_log = self.should_write # and logging
- if self.should_log:
- if self.config.get("debug", False):
- pass
- else:
- tags = self.config.tags.split(
- ',') if self.config.tags != '' else None
- wandb.init(project=self.config.project, name=self.config.experiment_id, config=flatten(self.config), dir=self.config.root,
- tags=tags, notes=self.config.notes, settings=wandb.Settings(start_method="fork"))
-
- self.model.train()
- self.step = 0
- best_loss = np.inf
- validate_every = int(self.config.validate_every * self.iters_per_epoch)
-
-
- if self.config.prefetch:
-
- for i, batch in tqdm(enumerate(self.train_loader), desc=f"Prefetching...",
- total=self.iters_per_epoch) if is_rank_zero(self.config) else enumerate(self.train_loader):
- pass
-
- losses = {}
- def stringify_losses(L): return "; ".join(map(
- lambda kv: f"{colors.fg.purple}{kv[0]}{colors.reset}: {round(kv[1].item(),3):.4e}", L.items()))
-
- epoch_len = len(self.train_loader)
- total_step = epoch_len * self.config.epochs
-
- for epoch in range(self.config.epochs):
- if self.should_early_stop():
- break
-
- self.epoch = epoch
- # self.save_checkpoint(f"{self.config.experiment_id}_latest.pt") # debug
- ################################# Train loop ##########################################################
- if self.should_log:
- if self.config.get("debug", False):
- pass
- else:
- wandb.log({"Epoch": epoch}, step=self.step)
- pbar = tqdm(enumerate(self.train_loader), desc=f"Epoch: {epoch + 1}/{self.config.epochs}. Loop: Train",
- total=self.iters_per_epoch) if is_rank_zero(self.config) else enumerate(self.train_loader)
-
- # 1532146.125
- for i, batch in pbar:
- current_step = epoch_len * epoch + i
- step_rate = current_step / total_step
- # metrics, test_losses = self.validate()
- # print(metrics)
- if self.should_early_stop():
- print("Early stopping")
- break
- # print(f"Batch {self.step+1} on rank {self.config.rank}")
- losses = self.train_on_batch(batch, i, step_rate)
- # print(f"trained batch {self.step+1} on rank {self.config.rank}")
- if self.config.get("debug", False):
- log_info = ""
- for name, loss in losses.items():
- log_info += "{}: {}, ".format(name, loss)
- print(log_info)
-
- if self.boostingdepth:
- for k,v in losses.items():
- losses[k] = torch.tensor(v)
-
- self.raise_if_nan(losses)
- if is_rank_zero(self.config) and self.config.print_losses:
- pbar.set_description(
- f"Epoch: {epoch + 1}/{self.config.epochs}. Loop: Train. Losses: {stringify_losses(losses)}")
- self.scheduler.step()
-
- if self.should_log and self.step % 50 == 0:
- if self.config.get("debug", False):
- log_info = ""
- for name, loss in losses.items():
- log_info += "{}: {}, ".format(name, loss)
- print(log_info)
- else:
- wandb.log({f"Train/{name}": loss.item()
- for name, loss in losses.items()}, step=self.step)
- # current_lr = self.optimizer.param_groups[0]['lr']
- current_lr = self.scheduler.get_last_lr()[0]
- wandb.log({f"Train/LR": current_lr}, step=self.step)
- momentum = self.optimizer.param_groups[0]['betas'][0]
- wandb.log({f"Train/momentum": momentum}, step=self.step)
- wandb.log({f"Train/step_rate": step_rate}, step=self.step)
-
- self.step += 1
-
- ########################################################################################################
-
- if self.test_loader:
- if (self.step % validate_every) == 0:
- self.model.eval()
- if self.should_write:
- self.save_checkpoint(
- f"{self.config.experiment_id}_latest.pt")
-
- ################################# Validation loop ##################################################
- # validate on the entire validation set in every process but save only from rank 0, I know, inefficient, but avoids divergence of processes
- metrics, test_losses = self.validate()
- # print("Validated: {}".format(metrics))
- if self.should_log:
- if self.config.get("debug", False):
- log_info = ""
- for name, loss in test_losses.items():
- log_info += "{}: {}, ".format(name, loss)
- log_info = "\n"
- for name, val in metrics.items():
- log_info += "{}: {}, ".format(name, val)
- print(log_info)
-
- else:
- wandb.log(
- {f"Test/{name}": tloss for name, tloss in test_losses.items()}, step=self.step)
-
- wandb.log({f"Metrics/{k}": v for k,
- v in metrics.items()}, step=self.step)
-
- if (metrics[self.metric_criterion] < best_loss) and self.should_write:
- self.save_checkpoint(
- f"{self.config.experiment_id}_best.pt")
- best_loss = metrics[self.metric_criterion]
-
- self.model.train()
-
- if self.config.distributed:
- dist.barrier()
- # print(f"Validated: {metrics} on device {self.config.rank}")
-
- # print(f"Finished step {self.step} on device {self.config.rank}")
- #################################################################################################
-
- # Save / validate at the end
- self.step += 1 # log as final point
- self.model.eval()
- self.save_checkpoint(f"{self.config.experiment_id}_latest.pt")
- if self.test_loader:
-
- ################################# Validation loop ##################################################
- metrics, test_losses = self.validate()
- # print("Validated: {}".format(metrics))
- if self.should_log:
- if self.config.get("debug", False):
- log_info = ""
- for name, loss in test_losses.items():
- log_info += "{}: {}, ".format(name, loss)
- log_info = "\n"
- for name, val in metrics.items():
- log_info += "{}: {}, ".format(name, val)
- print(log_info)
-
- else:
- wandb.log({f"Test/{name}": tloss for name,
- tloss in test_losses.items()}, step=self.step)
- wandb.log({f"Metrics/{k}": v for k,
- v in metrics.items()}, step=self.step)
-
- if (metrics[self.metric_criterion] < best_loss) and self.should_write:
- self.save_checkpoint(
- f"{self.config.experiment_id}_best.pt")
- best_loss = metrics[self.metric_criterion]
-
- self.model.train()
-
-
- def init_optimizer(self):
- m = self.model.module if self.config.multigpu else self.model
-
- if self.config.same_lr:
- print("Using same LR")
- if hasattr(m, 'core'):
- m.core.unfreeze()
- params = self.model.parameters()
- else:
- print("Using diff LR")
- if not hasattr(m, 'get_lr_params'):
- raise NotImplementedError(
- f"Model {m.__class__.__name__} does not implement get_lr_params. Please implement it or use the same LR for all parameters.")
-
- params = m.get_lr_params(self.config.lr)
-
- # if self.addf:
- # return optim.Adam(params, lr=self.config.lr, betas=(0.5, 0.999))
- # else:
- # return optim.AdamW(params, lr=self.config.lr, weight_decay=self.config.wd)
-
- return optim.AdamW(params, lr=self.config.lr, weight_decay=self.config.wd)
-
-
-
- def save_checkpoint(self, filename):
- if not self.should_write:
- return
- root = self.config.save_dir
- if not os.path.isdir(root):
- os.makedirs(root)
-
- fpath = os.path.join(root, filename)
- m = self.model.module if self.config.multigpu else self.model
- torch.save(
- {
- "model": m.state_dict(),
- "optimizer": None, # TODO : Change to self.optimizer.state_dict() if resume support is needed, currently None to reduce file size
- "epoch": self.epoch
- }, fpath)
-
- if self.boostingdepth:
- fpath = os.path.join(root, "_fusion" + filename)
- m.fusion_network.save_networks(fpath)
diff --git a/zoedepth/trainers/zoedepth_nk_trainer.py b/zoedepth/trainers/zoedepth_nk_trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..d528ae126f1c51b2f25fd31f94a39591ceb2f43a
--- /dev/null
+++ b/zoedepth/trainers/zoedepth_nk_trainer.py
@@ -0,0 +1,143 @@
+# MIT License
+
+# Copyright (c) 2022 Intelligent Systems Lab Org
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# File author: Shariq Farooq Bhat
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+
+from zoedepth.trainers.loss import GradL1Loss, SILogLoss
+from zoedepth.utils.config import DATASETS_CONFIG
+from zoedepth.utils.misc import compute_metrics
+
+from .base_trainer import BaseTrainer
+
+
+class Trainer(BaseTrainer):
+ def __init__(self, config, model, train_loader, test_loader=None, device=None):
+ super().__init__(config, model, train_loader,
+ test_loader=test_loader, device=device)
+ self.device = device
+ self.silog_loss = SILogLoss()
+ self.grad_loss = GradL1Loss()
+ self.domain_classifier_loss = nn.CrossEntropyLoss()
+
+ self.scaler = amp.GradScaler(enabled=self.config.use_amp)
+
+ def train_on_batch(self, batch, train_step):
+ """
+ Expects a batch of images and depth as input
+ batch["image"].shape : batch_size, c, h, w
+ batch["depth"].shape : batch_size, 1, h, w
+
+ Assumes all images in a batch are from the same dataset
+ """
+
+ images, depths_gt = batch['image'].to(
+ self.device), batch['depth'].to(self.device)
+ # batch['dataset'] is a tensor strings all valued either 'nyu' or 'kitti'. labels nyu -> 0, kitti -> 1
+ dataset = batch['dataset'][0]
+ # Convert to 0s or 1s
+ domain_labels = torch.Tensor([dataset == 'kitti' for _ in range(
+ images.size(0))]).to(torch.long).to(self.device)
+
+ # m = self.model.module if self.config.multigpu else self.model
+
+ b, c, h, w = images.size()
+ mask = batch["mask"].to(self.device).to(torch.bool)
+
+ losses = {}
+
+ with amp.autocast(enabled=self.config.use_amp):
+ output = self.model(images)
+ pred_depths = output['metric_depth']
+ domain_logits = output['domain_logits']
+
+ l_si, pred = self.silog_loss(
+ pred_depths, depths_gt, mask=mask, interpolate=True, return_interpolated=True)
+ loss = self.config.w_si * l_si
+ losses[self.silog_loss.name] = l_si
+
+ if self.config.w_grad > 0:
+ l_grad = self.grad_loss(pred, depths_gt, mask=mask)
+ loss = loss + self.config.w_grad * l_grad
+ losses[self.grad_loss.name] = l_grad
+ else:
+ l_grad = torch.Tensor([0])
+
+ if self.config.w_domain > 0:
+ l_domain = self.domain_classifier_loss(
+ domain_logits, domain_labels)
+ loss = loss + self.config.w_domain * l_domain
+ losses["DomainLoss"] = l_domain
+ else:
+ l_domain = torch.Tensor([0.])
+
+ self.scaler.scale(loss).backward()
+
+ if self.config.clip_grad > 0:
+ self.scaler.unscale_(self.optimizer)
+ nn.utils.clip_grad_norm_(
+ self.model.parameters(), self.config.clip_grad)
+
+ self.scaler.step(self.optimizer)
+
+ if self.should_log and self.step > 1 and (self.step % int(self.config.log_images_every * self.iters_per_epoch)) == 0:
+ depths_gt[torch.logical_not(mask)] = -99
+ self.log_images(rgb={"Input": images[0, ...]}, depth={"GT": depths_gt[0], "PredictedMono": pred[0]}, prefix="Train",
+ min_depth=DATASETS_CONFIG[dataset]['min_depth'], max_depth=DATASETS_CONFIG[dataset]['max_depth'])
+
+ self.scaler.update()
+ self.optimizer.zero_grad(set_to_none=True)
+
+ return losses
+
+ def validate_on_batch(self, batch, val_step):
+ images = batch['image'].to(self.device)
+ depths_gt = batch['depth'].to(self.device)
+ dataset = batch['dataset'][0]
+ if 'has_valid_depth' in batch:
+ if not batch['has_valid_depth']:
+ return None, None
+
+ depths_gt = depths_gt.squeeze().unsqueeze(0).unsqueeze(0)
+ with amp.autocast(enabled=self.config.use_amp):
+ m = self.model.module if self.config.multigpu else self.model
+ pred_depths = m(images)["metric_depth"]
+ pred_depths = pred_depths.squeeze().unsqueeze(0).unsqueeze(0)
+
+ mask = torch.logical_and(
+ depths_gt > self.config.min_depth, depths_gt < self.config.max_depth)
+ with amp.autocast(enabled=self.config.use_amp):
+ l_depth = self.silog_loss(
+ pred_depths, depths_gt, mask=mask.to(torch.bool), interpolate=True)
+
+ metrics = compute_metrics(depths_gt, pred_depths, **self.config)
+ losses = {f"{self.silog_loss.name}": l_depth.item()}
+
+ if val_step == 1 and self.should_log:
+ depths_gt[torch.logical_not(mask)] = -99
+ self.log_images(rgb={"Input": images[0]}, depth={"GT": depths_gt[0], "PredictedMono": pred_depths[0]}, prefix="Test",
+ min_depth=DATASETS_CONFIG[dataset]['min_depth'], max_depth=DATASETS_CONFIG[dataset]['max_depth'])
+
+ return metrics, losses
diff --git a/zoedepth/trainers/zoedepth_trainer.py b/zoedepth/trainers/zoedepth_trainer.py
index 5fa5de7de1cc1ec388d5e3ee5b80b93ae3b56a1f..3ac1c24c0512c1c1b191670a7c24abb4fca47ba1 100644
--- a/zoedepth/trainers/zoedepth_trainer.py
+++ b/zoedepth/trainers/zoedepth_trainer.py
@@ -22,8 +22,6 @@
# File author: Shariq Farooq Bhat
-# This file may include modifications from author Zhenyu Li
-
import torch
import torch.cuda.amp as amp
import torch.nn as nn
@@ -64,6 +62,7 @@ class Trainer(BaseTrainer):
losses = {}
with amp.autocast(enabled=self.config.use_amp):
+
output = self.model(images)
pred_depths = output['metric_depth']
@@ -151,8 +150,6 @@ class Trainer(BaseTrainer):
depths_gt = batch['depth'].to(self.device)
dataset = batch['dataset'][0]
mask = batch["mask"].to(self.device)
- disp_gt_edges = batch['disp_gt_edges'].squeeze().numpy()
-
if 'has_valid_depth' in batch:
if not batch['has_valid_depth']:
return None, None
@@ -169,13 +166,12 @@ class Trainer(BaseTrainer):
l_depth = self.silog_loss(
pred_depths, depths_gt, mask=mask.to(torch.bool), interpolate=True)
- metrics = compute_metrics(depths_gt, pred_depths, disp_gt_edges=disp_gt_edges, **self.config)
+ metrics = compute_metrics(depths_gt, pred_depths, **self.config)
losses = {f"{self.silog_loss.name}": l_depth.item()}
- if self.should_log:
- if val_step in [1, 21, 41, 61]:
- depths_gt[torch.logical_not(mask)] = DATASETS_CONFIG[dataset]['max_depth']
- self.log_images(rgb={"Input": images[0]}, depth={"GT": depths_gt[0], "PredictedMono": pred_depths[0]}, prefix="Test",
- min_depth=DATASETS_CONFIG[dataset]['min_depth'], max_depth=DATASETS_CONFIG[dataset]['max_depth'])
+ if val_step == 1 and self.should_log:
+ depths_gt[torch.logical_not(mask)] = -99
+ self.log_images(rgb={"Input": images[0]}, depth={"GT": depths_gt[0], "PredictedMono": pred_depths[0]}, prefix="Test",
+ min_depth=DATASETS_CONFIG[dataset]['min_depth'], max_depth=DATASETS_CONFIG[dataset]['max_depth'])
return metrics, losses
diff --git a/zoedepth/utils/align/depth_alignment.py b/zoedepth/utils/align/depth_alignment.py
deleted file mode 100644
index 939eca0376eebe06f6c7f03756dfe2c4c82d89f1..0000000000000000000000000000000000000000
--- a/zoedepth/utils/align/depth_alignment.py
+++ /dev/null
@@ -1,374 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Shariq Farooq Bhat, Zhenyu Li
-
-import torch
-import numpy as np
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-from tqdm.auto import tqdm
-from torchvision.transforms import ToTensor, ToPILImage
-from typing import List, Tuple
-from PIL import Image
-# from models.monodepth.zoedepth import ZoeDepthLora
-# from zoedepth.utils.align.loss import SILogLoss, gradl1_loss, edge_aware_smoothness_per_pixel, ssim_loss
-from .loss import *
-import cv2
-from zoedepth.trainers.loss import *
-# from utils.misc import *
-
-
-
-@torch.no_grad()
-def scale_shift_linear(rendered_depth, predicted_depth, mask, fuse=True, return_params=False):
- """
- Optimize a scale and shift parameter in the least squares sense, such that rendered_depth and predicted_depth match.
- Formally, solves the following objective:
-
- min || (d * a + b) - d_hat ||
- a, b
-
- where d = 1 / predicted_depth, d_hat = 1 / rendered_depth
-
- :param rendered_depth: torch.Tensor (H, W)
- :param predicted_depth: torch.Tensor (H, W)
- :param mask: torch.Tensor (H, W) - 1: valid points of rendered_depth, 0: invalid points of rendered_depth (ignore)
- :param fuse: whether to fuse shifted/scaled predicted_depth with the rendered_depth
-
- :return: scale/shift corrected depth
- """
- if mask.sum() == 0:
- return predicted_depth
-
- # rendered_disparity = 1 / rendered_depth[mask].unsqueeze(-1)
- # predicted_disparity = 1 / predicted_depth[mask].unsqueeze(-1)
-
- rendered_disparity = rendered_depth[mask].unsqueeze(-1)
- predicted_disparity = predicted_depth[mask].unsqueeze(-1)
-
- X = torch.cat([predicted_disparity, torch.ones_like(predicted_disparity)], dim=1)
- XTX_inv = (X.T @ X).inverse()
- XTY = X.T @ rendered_disparity
- AB = XTX_inv @ XTY
-
- if return_params:
- return AB
-
- fixed_disparity = (predicted_depth) * AB[0] + AB[1]
- fixed_depth = fixed_disparity
-
- if fuse:
- fused_depth = torch.where(mask, rendered_depth, fixed_depth)
- return fused_depth
- else:
- return fixed_depth
-
-
-def np_scale_shift_linear(rendered_depth: np.ndarray, predicted_depth: np.ndarray, mask: np.ndarray, fuse: bool=True):
- """
- Optimize a scale and shift parameter in the least squares sense, such that rendered_depth and predicted_depth match.
- Formally, solves the following objective:
-
- min || (d * a + b) - d_hat ||
- a, b
-
- where d = predicted_depth, d_hat = rendered_depth
-
- :param rendered_depth: np.ndarray (H, W)
- :param predicted_depth: np.ndarray (H, W)
- :param mask: np.ndarray (H, W) - 1: valid points of rendered_depth, 0: invalid points of rendered_depth (ignore)
- :param fuse: whether to fuse shifted/scaled predicted_depth with the rendered_depth
-
- :return: scale/shift corrected depth
- """
- if mask.sum() == 0:
- return predicted_depth
-
- # rendered_disparity = 1 / rendered_depth[mask].reshape(-1, 1)
- # predicted_disparity = 1 / predicted_depth[mask].reshape(-1, 1)
-
- rendered_disparity = rendered_depth[mask].reshape(-1, 1)
- predicted_disparity = predicted_depth[mask].reshape(-1, 1)
-
- X = np.concatenate([predicted_disparity, np.ones_like(predicted_disparity)], axis=1)
- XTX_inv = np.linalg.inv(X.T @ X)
- XTY = X.T @ rendered_disparity
- AB = XTX_inv @ XTY
-
- fixed_disparity = (predicted_depth) * AB[0] + AB[1]
- fixed_depth = fixed_disparity
-
- if fuse:
- fused_depth = np.where(mask, rendered_depth, fixed_depth)
- return fused_depth
- else:
- return fixed_depth
-
-
-@torch.no_grad()
-def apply_depth_smoothing(depth, mask):
-
- def dilate(x, k=3):
- x = as_bchw_tensor(x.float(), 1)
- x = torch.nn.functional.conv2d(x.float(),
- torch.ones(1, 1, k, k).to(x.device),
- padding="same"
- )
- return x.squeeze() > 0
-
- def sobel(x):
- flipped_sobel_x = torch.tensor([
- [-1, 0, 1],
- [-2, 0, 2],
- [-1, 0, 1]
- ], dtype=torch.float32).to(x.device)
- flipped_sobel_x = torch.stack([flipped_sobel_x, flipped_sobel_x.t()]).unsqueeze(1)
-
- x_pad = torch.nn.functional.pad(x.float(), (1, 1, 1, 1), mode="replicate")
-
- x = torch.nn.functional.conv2d(
- x_pad,
- flipped_sobel_x,
- padding="valid"
- )
- dx, dy = x.unbind(dim=-3)
- # return torch.sqrt(dx**2 + dy**2).squeeze()
- # new content is created mostly in x direction, sharp edges in y direction are wanted (e.g. table --> wall)
- return dx
-
- depth = as_bchw_tensor(depth, 1)
- mask = as_bchw_tensor(mask, 1).float()
-
- edges = sobel(mask)
- dilated_edges = dilate(edges, k=21)
-
- depth_numpy = depth.squeeze().float().cpu().numpy()
- blur_bilateral = cv2.bilateralFilter(depth_numpy, 5, 140, 140)
- blur_gaussian = cv2.GaussianBlur(blur_bilateral, (5, 5), 0)
- blur_gaussian = torch.from_numpy(blur_gaussian).to(depth)
- # print("blur_gaussian", blur_gaussian.shape)
- # plt.imshow(blur_gaussian.cpu().squeeze().numpy())
- # plt.title("depth smoothed whole")
- # plt.show()
- depth_smooth = torch.where(dilated_edges, blur_gaussian, depth)
- return depth_smooth
-
-def get_dilated_only_mask(mask: torch.Tensor, k=7):
- x = as_bchw_tensor(mask.float(), 1)
- x = torch.nn.functional.conv2d(x, torch.ones(1, 1, k, k).to(mask.device),padding="same")
- dilated = x.squeeze() > 0
- dilated_only = dilated ^ mask
- return dilated_only
-
-def get_boundary_mask(mask: torch.Tensor, k=7):
- return get_dilated_only_mask(mask, k=k) | get_dilated_only_mask(~mask, k=k)
-
-
-@torch.no_grad()
-def ss_align_and_blur(rendered_depth: torch.Tensor, predicted_depth: torch.Tensor, mask: torch.Tensor, fuse: bool=True):
- aligned = scale_shift_linear(rendered_depth, predicted_depth, mask, fuse=fuse)
- aligned = apply_depth_smoothing(aligned, mask)
- return aligned
-
-
-def np_ss_align_and_blur(rendered_depth: np.ndarray, predicted_depth: np.ndarray, mask: np.ndarray, fuse: bool=True):
- aligned = np_scale_shift_linear(rendered_depth, predicted_depth, mask, fuse=fuse)
- aligned = apply_depth_smoothing(aligned, mask).cpu().numpy()
- return aligned
-
-
-
-def stitch(depth_src: torch.Tensor, depth_target: torch.Tensor, mask_src: torch.Tensor, smoothen=True, device='cuda:0'):
- depth_src = as_bchw_tensor(depth_src, 1, device=device)
- depth_target = as_bchw_tensor(depth_target, 1, device=device)
- mask_src = as_bchw_tensor(mask_src, 1, device=device)
-
- stitched = depth_src * mask_src.float() + depth_target * (~mask_src).float()
- # plt.imshow(stitched.cpu().squeeze().numpy())
- # plt.title("stitched before smoothing")
- # plt.show()
- # apply smoothing
- if smoothen:
- stitched = apply_depth_smoothing(stitched, mask_src).squeeze().float()
- return stitched
-
-
-
-
-
-
-
-
-
-def smoothness_loss(depth, mask=None):
- depth_grad_x = torch.abs(depth[:, :, :, :-1] - depth[:, :, :, 1:])
- depth_grad_y = torch.abs(depth[:, :, :-1, :] - depth[:, :, 1:, :])
- if mask is not None:
- return torch.mean(depth_grad_x[mask[:, :, :, :-1]]) + torch.mean(depth_grad_y[mask[:, :, :-1, :]])
- return torch.mean(depth_grad_x) + torch.mean(depth_grad_y)
-
-import torch.optim as optim
-from torch.optim import lr_scheduler
-def finetune_on_sample(model, image_pil, target_depth, mask=None,
- iters=10, lr=0.1, beta=0.5, w_boundary_grad=1, w_grad=0.1, gamma=0.99):
- model.train()
- model_device = next(model.parameters()).device
- x = as_bchw_tensor(image_pil, 3, device=model_device)
- target_depth = as_bchw_tensor(target_depth, 1, device=model_device)
- if mask is None:
- mask = target_depth > 0
- elif (not isinstance(mask, torch.Tensor)) or mask.shape != target_depth.shape:
- mask = as_bchw_tensor(mask, 1, device=model_device).to(torch.bool)
-
-
- history = []
- optimizer = torch.optim.Adam(model.parameters(), lr=lr)
- # scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma)
- scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=lr, steps_per_epoch=iters, epochs=1)
- # main_loss = nn.L1Loss()
- main_loss = SILogLoss(beta=beta)
-
- orig_y = model.infer(x, with_flip_aug=False).detach()
-
- # scale, shift = scale_shift_linear(target_depth, orig_y, mask, return_params=True)
-
- gl1 = gradl1_loss
- pbar = tqdm(range(iters), desc="Finetuning on sample")
- for i in pbar:
- optimizer.zero_grad()
- y = model.infer(x, with_flip_aug=False)
- # y = y * scale + shift
- stitched = y * (~mask).float() + (target_depth * (mask).float()).detach()
- # loss = F.mse_loss(y[mask], target_depth[mask])
- loss_si = main_loss(y[mask], target_depth[mask])
- # loss = loss_si \
- # + wgrad * ( gl1(y, stitched) \
- # + 2*gl1(y, orig_y) ) \
- # + wboundary_smoothness * smoothness_loss(y, mask=get_boundary_mask(mask))
- loss_grad = gl1(y, orig_y)
-
- bmask = get_boundary_mask(mask)
- loss_boundary_grad = laplacian_matching_loss(stitched, orig_y, bmask)
- loss = loss_si + w_boundary_grad * loss_boundary_grad + w_grad * loss_grad
-
- # check if loss is nan
- if torch.isnan(loss):
- print("Loss is nan, breaking")
- break
- loss.backward()
-
- optimizer.step()
- scheduler.step()
- # history.append(loss.item())
-
- pbar.set_postfix(loss=loss.item(), si=loss_si.item())
- model.eval()
- return model, history
-
-# def align_by_finetuning_lora(model: ZoeDepthLora, image, target_depth, mask=None, iters=10, lr=0.1, gamma=0.99, **kwargs):
-# # model.reset_lora()
-# model.set_only_lora_trainable()
-# model, history = finetune_on_sample(model, image, target_depth, mask=mask, iters=iters, lr=lr, gamma=gamma)
-# aligned_depth = model.infer(as_bchw_tensor(image, 3, device=next(model.parameters()).device))
-# return dict(model=model, history=history, aligned_depth=aligned_depth)
-
-
-
-
-
-
-import torch.nn as nn
-import torch.nn.functional as F
-# from utils.misc import as_bchw_tensor
-
-def as_bchw_tensor(input_tensor, num, device):
- input_tensor = torch.tensor(input_tensor).unsqueeze(dim=0).unsqueeze(dim=0).cuda()
- return input_tensor
-
-def optimize_depth_deformation(rendered_depth, pred_depth, mask, h=10, w=10, iters=100, init_lr=0.1, gamma=0.996,
- init_deformation=None,
- device='cuda:0'):
- rendered_depth = as_bchw_tensor(rendered_depth, 1, device=device)
- pred_depth = as_bchw_tensor(pred_depth, 1, device=device)
- mask = as_bchw_tensor(mask, 1, device=device).to(torch.bool)
- # initialize a grid of scalar values (with zeros) that will be optimized
- # to deform the depth map
- if init_deformation is None:
- deformation = torch.zeros((1,1,h,w), requires_grad=True, device=device)
- else:
- deformation = init_deformation
- deformation.requires_grad = True
- assert deformation.shape == (1,1,h,w)
-
- optimizer = torch.optim.Adam([deformation], lr=init_lr)
- # exponential LR schedule
- scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma)
- # optimize the deformation
- history = []
- grad_loss = GradL1Loss()
- for i in tqdm(range(iters)):
- scalar_deformation = torch.exp(deformation)
- scalar_deformation = F.interpolate(scalar_deformation, size=pred_depth.shape[-2:], mode='bilinear', align_corners=True)
- adjusted_depth = pred_depth * scalar_deformation
- loss = F.mse_loss(adjusted_depth[mask], rendered_depth[mask], reduction='none')
- loss_g = grad_loss(adjusted_depth, rendered_depth, mask)
- loss = loss.mean() + 0.1*loss_g
- # loss = loss.mean()
- optimizer.zero_grad()
- loss.backward()
- optimizer.step()
- scheduler.step()
- if i % 10 == 0:
- history.append(loss.item())
-
- scalar_deformation = torch.exp(deformation)
- scalar_deformation = F.interpolate(scalar_deformation, size=pred_depth.shape[-2:], mode='bilinear', align_corners=True)
- adjusted_depth = pred_depth * scalar_deformation
- # return dict(aligned_depth=adjusted_depth.detach().cpu().numpy().squeeze(),
- # history=history,
- # deformation=deformation)
- return adjusted_depth.detach().cpu().squeeze()
-
-def stage_wise_optimization(rendered_depth, pred_depth, mask,
- stages=[(4,4), (8,8), (16,16), (32,32)],
- iters=100, init_lr=0.1, gamma=0.996, device='cuda:1'):
-
- h_init, w_init = stages[0]
- init_deformation = torch.zeros((1,1,h_init,w_init), device=device)
- result = optimize_depth_deformation(rendered_depth, pred_depth, mask, h=h_init, w=w_init, iters=iters, init_lr=init_lr, gamma=gamma, init_deformation=init_deformation, device=device)
- init_deformation = result['deformation']
- history_stages = [result['history']]
- for h, w in stages[1:]:
- init_deformation = F.interpolate(init_deformation, size=(h,w), mode='bilinear', align_corners=True).detach()
- result = optimize_depth_deformation(rendered_depth, pred_depth, mask, h=h, w=w, iters=iters, init_lr=init_lr, gamma=gamma, init_deformation=init_deformation, device=device)
- init_deformation = result['deformation']
- history_stages.append(result['history'])
- init_lr *= gamma**2
-
- return dict(aligned_depth=result['aligned_depth'], history_stages=history_stages)
-
-
-
-
-
diff --git a/zoedepth/utils/align/loss.py b/zoedepth/utils/align/loss.py
deleted file mode 100644
index cdca588d7abb8417a43c28c850fa7c6e58a220ac..0000000000000000000000000000000000000000
--- a/zoedepth/utils/align/loss.py
+++ /dev/null
@@ -1,178 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Shariq Farooq Bhat, Zhenyu Li
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-
-def gaussian(window_size, sigma):
- gauss = torch.Tensor([np.exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
- return gauss/gauss.sum()
-
-def create_window(window_size, channel=1):
- _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
- _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
- window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
- return window
-
-def ssim(img1, img2, val_range, window_size=11, window=None, size_average=True, full=False):
- img1 = nn.functional.interpolate(img1, (256, 256), mode='bilinear', align_corners=True)
- img2 = nn.functional.interpolate(img2, (256, 256), mode='bilinear', align_corners=True)
- # h, w = 256, 256
- L = val_range
-
- padd = 0
- (_, channel, height, width) = img1.size()
- if window is None:
- real_size = min(window_size, height, width)
- window = create_window(real_size, channel=channel).to(img1.device)
-
- mu1 = F.conv2d(img1, window, padding=padd, groups=channel)
- mu2 = F.conv2d(img2, window, padding=padd, groups=channel)
-
- mu1_sq = mu1.pow(2)
- mu2_sq = mu2.pow(2)
- mu1_mu2 = mu1 * mu2
-
- sigma1_sq = F.conv2d(img1 * img1, window, padding=padd, groups=channel) - mu1_sq
- sigma2_sq = F.conv2d(img2 * img2, window, padding=padd, groups=channel) - mu2_sq
- sigma12 = F.conv2d(img1 * img2, window, padding=padd, groups=channel) - mu1_mu2
-
- C1 = (0.01 * L) ** 2
- C2 = (0.03 * L) ** 2
-
- v1 = 2.0 * sigma12 + C2
- v2 = sigma1_sq + sigma2_sq + C2
- cs = torch.mean(v1 / v2) # contrast sensitivity
-
- ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
-
- if size_average:
- ret = ssim_map.mean()
- else:
- ret = ssim_map.mean(1).mean(1).mean(1)
-
- if full:
- return ret, cs
-
- return ret
-
-
-
-class SSIMLoss(nn.Module):
- def __init__(self, min_depth=1e-3, max_depth=10):
- super(SSIMLoss, self).__init__()
- self.name = 'SSIM'
- self.min_depth = min_depth
- self.max_depth = max_depth
-
- def forward(self, input, target):
- loss = torch.clamp((1 - ssim(input, target, val_range=self.max_depth/self.min_depth)) * 0.5, 0, 1)
- return loss
-
-
-# Main loss function used for ZoeDepth. Copy/paste from AdaBins repo (https://github.com/shariqfarooq123/AdaBins/blob/0952d91e9e762be310bb4cd055cbfe2448c0ce20/loss.py#L7)
-class SILogLoss(nn.Module):
- """SILog loss (pixel-wise)"""
- def __init__(self, beta=0.15):
- super().__init__()
- self.name = 'SILog'
- self.beta = beta
-
- def forward(self, input, target):
- alpha = 1e-10
- g = torch.log(input + alpha) - torch.log(target + alpha)
-
- # n, c, h, w = g.shape
- # norm = 1/(h*w)
- # Dg = norm * torch.sum(g**2) - (0.85/(norm**2)) * (torch.sum(g))**2
-
- Dg = torch.var(g) + self.beta * torch.pow(torch.mean(g), 2)
-
- loss = 10 * torch.sqrt(Dg)
- return loss
-
-def gradient_y(img):
- gy = torch.cat( [F.conv2d(img[:, i, :, :].unsqueeze(0), torch.Tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]]).view((1, 1, 3, 3)).to(img.device), padding=1) for i in range(img.shape[1])], 1)
- return gy
-
-def gradient_x(img):
- gx = torch.cat( [F.conv2d(img[:, i, :, :].unsqueeze(0), torch.Tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]]).view((1, 1, 3, 3)).to(img.device), padding=1) for i in range(img.shape[1])], 1)
- return gx
-
-def laplacian(img):
- lap = torch.cat( [F.conv2d(img[:, i, :, :].unsqueeze(0), torch.Tensor([[0, 1, 0], [1, -4, 1], [0, 1, 0]]).view((1, 1, 3, 3)).to(img.device), padding=1) for i in range(img.shape[1])], 1)
- return lap
-
-def laplacian_matching_loss(img1, img2, mask=None):
- return torch.mean(torch.abs(laplacian(img1)[mask] - laplacian(img2)[mask]))
-
-class GradL1Loss(nn.Module):
- def __init__(self):
- super(GradL1Loss, self).__init__()
- self.name = 'GradL1'
-
- def forward(self, input, target, mask=None):
-
- grad_gt_x = gradient_x(target)
- grad_gt_y = gradient_y(target)
- grad_pred_x = gradient_x(input)
- grad_pred_y = gradient_y(input)
- loss = torch.mean(torch.abs(grad_pred_x[mask] - grad_gt_x[mask])) + torch.mean(torch.abs(grad_pred_y[mask] - grad_gt_y[mask]))
- return loss
-
-# Edge aware smoothness loss implementation is adapted from: https://github.com/anuragranj/cc
-def edge_aware_smoothness_per_pixel(img, pred):
- """ A measure of how closely the gradients of a predicted disparity/depth map match the
- gradients of the RGB image.
-
- Args:
- img (c x 3 x h x w tensor): RGB image
- pred (c x h x w tensor): predicted depth/disparity
-
- Returns:
- c x 1 tensor: measure of gradient matching (smoothness loss)
- """
-
-
-
- pred_gradients_x = gradient_x(pred)
- pred_gradients_y = gradient_y(pred)
-
- image_gradients_x = gradient_x(img)
- image_gradients_y = gradient_y(img)
-
- weights_x = torch.exp(-torch.mean(torch.abs(image_gradients_x), 1, keepdim=True))
- weights_y = torch.exp(-torch.mean(torch.abs(image_gradients_y), 1, keepdim=True))
-
- smoothness_x = torch.abs(pred_gradients_x) * weights_x
- smoothness_y = torch.abs(pred_gradients_y) * weights_y
-
- return torch.mean(smoothness_x) + torch.mean(smoothness_y)
-
-
-
-ssim_loss = SSIMLoss()
-gradl1_loss = GradL1Loss()
\ No newline at end of file
diff --git a/zoedepth/utils/align/mlp_alignment.py b/zoedepth/utils/align/mlp_alignment.py
deleted file mode 100644
index 2ea91df321e956883847cd912aef5bd7b4e934d6..0000000000000000000000000000000000000000
--- a/zoedepth/utils/align/mlp_alignment.py
+++ /dev/null
@@ -1,262 +0,0 @@
-# MIT License
-
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Shariq Farooq Bhat, Zhenyu Li
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from tqdm.auto import tqdm
-import torch.optim as optim
-import torch.optim.lr_scheduler
-from zoedepth.utils.align.loss import SILogLoss, gradl1_loss, edge_aware_smoothness_per_pixel
-# from utils.misc import *
-from .depth_alignment import apply_depth_smoothing, scale_shift_linear
-import cv2
-import numpy as np
-
-def as_bchw_tensor(input_tensor, num, device=None):
- if len(input_tensor.shape) == 2:
- input_tensor = torch.tensor(input_tensor).unsqueeze(dim=0).unsqueeze(dim=0)
- elif len(input_tensor.shape) == 3:
- input_tensor = torch.tensor(input_tensor).unsqueeze(dim=0)
- else:
- input_tensor = input_tensor
- if device is not None:
- input_tensor = input_tensor.to(device)
- return input_tensor
-
-
-def get_mlp(in_channels, out_channels):
- conv_config = dict(kernel_size=1, padding=0, stride=1)
- net = nn.Sequential(
- # nn.Conv2d(in_channels, 64, kernel_size=7, padding=3, stride=1),
- # nn.GELU(),
- nn.Conv2d(in_channels, 64, **conv_config),
- nn.GELU(),
- nn.Conv2d(64, 128, **conv_config),
- nn.GELU(),
- nn.Conv2d(128, out_channels, **conv_config),
- )
-
- # initialize last layer to predict zeroes
- # net[-1].weight.data.zero_()
- # net[-1].bias.data.zero_()
- return net
-
-def smoothness_loss(depth):
- depth_dx = depth[:, :, :-1, :-1] - depth[:, :, :-1, 1:]
- depth_dy = depth[:, :, :-1, :-1] - depth[:, :, 1:, :-1]
- depth_dx = depth_dx.abs().mean()
- depth_dy = depth_dy.abs().mean()
- return depth_dx + depth_dy
-
-def curvature_loss(depth):
- depth_dx = depth[:, :, :-1, :-1] - depth[:, :, :-1, 1:]
- depth_dy = depth[:, :, :-1, :-1] - depth[:, :, 1:, :-1]
- depth_dxx = depth_dx[:, :, :, :-1] - depth_dx[:, :, :, 1:]
- depth_dyy = depth_dy[:, :, :-1, :] - depth_dy[:, :, 1:, :]
- depth_dxy = depth_dx[:, :, :-1, :-1] - depth_dx[:, :, 1:, 1:]
- depth_dxx = depth_dxx.abs().mean()
- depth_dyy = depth_dyy.abs().mean()
- depth_dxy = depth_dxy.abs().mean()
- return depth_dxx + depth_dyy + depth_dxy
-
-def multi_scale_curvature_loss(depth, scales=[1, 2, 4]):
- loss = 0
- for s in scales:
- loss += curvature_loss(F.interpolate(depth, scale_factor=1/s, mode='bilinear', align_corners=False))
- return loss
-
-
-def tv_loss(x):
- """Total variation loss."""
- b, c, h, w = x.shape
- dh = torch.abs(x[:, :, 1:, :] - x[:, :, :-1, :])
- dw = torch.abs(x[:, :, :, 1:] - x[:, :, :, :-1])
- return torch.sum(dh) + torch.sum(dw)
-
-def scale_invariant_gradient_loss(pred, gt):
- alpha = 1e-10
- kernel_grad_x = torch.Tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]]).view((1, 1, 3, 3)).to(pred.device)
- kernel_grad_y = torch.Tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]]).view((1, 1, 3, 3)).to(pred.device)
-
- g = torch.log(pred + alpha) - torch.log(gt + alpha)
- g_x = F.conv2d(g, kernel_grad_x, padding=1)
- g_y = F.conv2d(g, kernel_grad_y, padding=1)
- # n, c, h, w = g.shape
- # norm = 1/(h*w)
- # Dg = norm * torch.sum(g**2) - (0.85/(norm**2)) * (torch.sum(g))**2
-
- Dgx = torch.var(g_x) + 0.5 * torch.pow(torch.mean(g_x), 2)
- Dgy = torch.var(g_y) + 0.5 * torch.pow(torch.mean(g_y), 2)
-
-
- loss = 10 * torch.sqrt(Dgx) + 10 * torch.sqrt(Dgy)
- return loss
-
-
-def positionalencoding2d(d_model, height, width):
- """
- :param d_model: dimension of the model
- :param height: height of the positions
- :param width: width of the positions
- :return: d_model*height*width position matrix
- """
- if d_model % 4 != 0:
- raise ValueError("Cannot use sin/cos positional encoding with "
- "odd dimension (got dim={:d})".format(d_model))
- pe = torch.zeros(d_model, height, width)
- # Each dimension use half of d_model
- d_model = int(d_model / 2)
- div_term = torch.exp(torch.arange(0., d_model, 2) *
- -(np.log(10000.0) / d_model))
- pos_w = torch.arange(0., width).unsqueeze(1)
- pos_h = torch.arange(0., height).unsqueeze(1)
- pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
- pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
- pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
- pe[d_model + 1::2, :, :] = torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
-
- return pe
-
-def gaussian_rff(d_model, height, width, sigma=10):
- assert d_model % 2 == 0
- B = torch.randn((d_model//2, 2)) * sigma
-
- x = torch.linspace(-1, 1, width)
- y = torch.linspace(-1, 1, height)
-
- x, y = torch.meshgrid(x, y)
- xy = torch.stack([x, y], dim=-1).view(-1, 2)
-
- xy = torch.matmul(B, xy.T).T
- xy = 2 * np.pi * xy.view(height, width, d_model//2)
- enc = torch.cat([torch.sin(xy), torch.cos(xy)], dim=-1)
- return enc.permute(2, 0, 1)
-
-
-def get_depth(pred, D):
- a, b, c = torch.split(pred, 1, dim=1)
- return 1e-7 + torch.relu(torch.exp(a) * D + (torch.sigmoid(c)-0.5)*torch.exp(b))
- # return 1e-4 + torch.exp(a) * D + b
- # return nn.Softplus()(a)
-
-
-def train_mlp(image, mask, dr, dp, lr=3e-2, num_iters=3000, device='cuda:0', pos_dim=32, loss_config=dict(beta=0.99),
- w_smooth=1, w_curvature=0.0, w_gl1=0.1, w_tv=0.1, w_shift_reg=0.1, **kwargs):
-
- mlp = get_mlp(pos_dim+4, 3)
- # mlp = get_mlp(4, 3)
- mlp = mlp.to(device)
-
- optimizer = optim.AdamW(mlp.parameters(), lr=lr)
- scheduler = optim.lr_scheduler.OneCycleLR(optimizer, lr, epochs=num_iters, steps_per_epoch=1)
-
- image = as_bchw_tensor(image, 3, device=device).detach()
- h, w = image.shape[-2:]
- pe = positionalencoding2d(pos_dim, h, w)
- pe = as_bchw_tensor(pe, pos_dim, device=device)
- D = as_bchw_tensor(dp, 1, device=device).detach()
- # pe = as_bchw_tensor(gaussian_rff(pos_dim, h, w, sigma=5), pos_dim, device=device)
- X = torch.cat([image, D, pe], dim=1) # bchw
- # X = torch.cat([image, D], dim=1) # bchw
-
- Y = as_bchw_tensor(dr, 1, device=device).detach()
- mask = as_bchw_tensor(mask, 1, device=device).detach()
- pbar = tqdm(range(num_iters), desc=f"Training")
- # beta_min, beta_max = 0.
- si_log = SILogLoss(**loss_config)
-
- for i in pbar:
- optimizer.zero_grad()
- # pred = dr.max().item() * torch.sigmoid(mlp(X))
- pred = mlp(X)
- a, b, c = torch.split(pred, 1, dim=1)
- pred = get_depth(pred, D.detach())
- loss_si = si_log(pred[mask], Y[mask])
- loss = loss_si + w_curvature * multi_scale_curvature_loss(pred) + w_gl1 * gradl1_loss(pred, D.detach()) + w_smooth * edge_aware_smoothness_per_pixel(image, pred)
- # loss_tv = w_tv * (tv_loss(a) + tv_loss(b) + tv_loss(c))
- # loss_gl1 = w_gl1 * gradl1_loss(pred, D.detach())
- # loss_gl1 = w_gl1 * scale_invariant_gradient_loss(pred, D.detach())
- # loss_shift_reg = w_shift_reg * torch.mean(b**2)
- # loss = loss_si + loss_gl1
- # loss = F.mse_loss(pred[mask], Y[mask])
- loss.backward()
- optimizer.step()
- scheduler.step()
- pbar.set_postfix(loss=loss.item(), si=loss_si.item())
-
- return mlp
-
-def predict_aligned(mlp, image, dp, pos_dim=32, **kwargs):
- device = next(mlp.parameters()).device
- image = as_bchw_tensor(image, 3, device=device)
- h, w = image.shape[-2:]
- pe = positionalencoding2d(pos_dim, h, w)
- pe = as_bchw_tensor(pe, pos_dim, device=device)
- D = as_bchw_tensor(dp, 1, device=device)
- # pe = as_bchw_tensor(gaussian_rff(pos_dim, h, w, sigma=5), pos_dim, device=device)
-
- X = torch.cat([image, D, pe], dim=1) # bchw
- # X = torch.cat([image, D], dim=1) # bchw
- pred = mlp(X)
- pred = get_depth(pred, D)
- return pred.detach()
-
-def align_by_mlp(image, mask, dr, dp, **kwargs):
- mlp = train_mlp(image, mask, dr, dp, **kwargs)
- pred = predict_aligned(mlp, image, dp, **kwargs)
- return pred
-
-from abc import ABC, abstractmethod
-
-
-# Abstract class for depth alignment. All depth alignment methods should inherit from this class.
-# The abstract class defines the interface for depth alignment.
-class DepthAligner(ABC):
- def __init__(self):
- super().__init__()
-
- @abstractmethod
- def align(self, depth_src, depth_target, valid_mask, *args, **kwargs):
- """
- Aligns the depth_src to the depth_target such that the aligned depth_src is as close as possible to the depth_target.
- """
- raise NotImplementedError
-
-class MLPAligner(DepthAligner):
- def __init__(self):
- super().__init__()
-
- def align(self, depth_src, depth_target, valid_mask, image, **kwargs):
- depth_src = as_bchw_tensor(depth_src, 1)
- depth_target = as_bchw_tensor(depth_target, 1)
- valid_mask = as_bchw_tensor(valid_mask, 1)
- depth_target = scale_shift_linear(depth_target, depth_src, valid_mask)
- aligned = align_by_mlp(image, valid_mask, depth_target, depth_src, **kwargs)
-
- depth_numpy = aligned.squeeze().float().cpu().numpy()
- blur_bilateral = cv2.bilateralFilter(depth_numpy, 5, 140, 140)
- blur_gaussian = cv2.GaussianBlur(blur_bilateral, (5, 5), 0)
- blur_gaussian = torch.from_numpy(blur_gaussian).to(aligned)
- return blur_gaussian.unsqueeze(0)
\ No newline at end of file
diff --git a/zoedepth/utils/arg_utils.py b/zoedepth/utils/arg_utils.py
index 8a2f3213c2f99c54c76886f6ffb4d9a71d363560..8a3004ec3679c0a40fd8961253733fb4343ad545 100644
--- a/zoedepth/utils/arg_utils.py
+++ b/zoedepth/utils/arg_utils.py
@@ -1,26 +1,4 @@
-# MIT License
-# Copyright (c) 2022 Intelligent Systems Lab Org
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# File author: Shariq Farooq Bhat
def infer_type(x): # hacky way to infer type from string args
if not isinstance(x, str):
diff --git a/zoedepth/utils/config.py b/zoedepth/utils/config.py
index b55b1c086e6d1b6711af254d77dffa0ad82d2ec3..363b0e186cd1247e8f5fc224e6f69f5fc8190c99 100644
--- a/zoedepth/utils/config.py
+++ b/zoedepth/utils/config.py
@@ -22,7 +22,7 @@
# File author: Shariq Farooq Bhat
-import json5
+import json
import os
from zoedepth.utils.easydict import EasyDict as edict
@@ -36,8 +36,7 @@ ROOT = pathlib.Path(__file__).parent.parent.resolve()
HOME_DIR = os.path.expanduser("~")
COMMON_CONFIG = {
- # "save_dir": os.path.expanduser("~/shortcuts/monodepth3_checkpoints"),
- "save_dir": "./nfs/monodepth3_checkpoints",
+ "save_dir": os.path.expanduser("~/shortcuts/monodepth3_checkpoints"),
"project": "ZoeDepth",
"tags": '',
"notes": "",
@@ -99,14 +98,14 @@ DATASETS_CONFIG = {
"avoid_boundary": False,
"min_depth": 1e-3, # originally 0.1
"max_depth": 10,
- "data_path": os.path.join("/ibex/ai/home/liz0l/codes/datasets/nyu/data_folder"),
- "gt_path": os.path.join("/ibex/ai/home/liz0l/codes/datasets/nyu/data_folder"),
- "filenames_file": "/ibex/ai/home/liz0l/codes/datasets/nyu/data_folder/nyu_train.txt",
+ "data_path": os.path.join(HOME_DIR, "shortcuts/datasets/nyu_depth_v2/sync/"),
+ "gt_path": os.path.join(HOME_DIR, "shortcuts/datasets/nyu_depth_v2/sync/"),
+ "filenames_file": "./train_test_inputs/nyudepthv2_train_files_with_gt.txt",
"input_height": 480,
"input_width": 640,
- "data_path_eval": os.path.join("/ibex/ai/home/liz0l/codes/datasets/nyu/data_folder"),
- "gt_path_eval": os.path.join("/ibex/ai/home/liz0l/codes/datasets/nyu/data_folder"),
- "filenames_file_eval": "/ibex/ai/home/liz0l/codes/datasets/nyu/data_folder/nyu_test.txt",
+ "data_path_eval": os.path.join(HOME_DIR, "shortcuts/datasets/nyu_depth_v2/official_splits/test/"),
+ "gt_path_eval": os.path.join(HOME_DIR, "shortcuts/datasets/nyu_depth_v2/official_splits/test/"),
+ "filenames_file_eval": "./train_test_inputs/nyudepthv2_test_files_with_gt.txt",
"min_depth_eval": 1e-3,
"max_depth_eval": 10,
"min_depth_diff": -10,
@@ -116,102 +115,7 @@ DATASETS_CONFIG = {
"degree": 1.0,
"do_kb_crop": False,
"garg_crop": False,
- "eigen_crop": False,
- },
- "u4k": {
- "dataset": "u4k",
- "min_depth": 1e-3, # originally 0.1
- "max_depth": 80,
- "data_path": os.path.join("/ibex/ai/home/liz0l/codes/datasets/u4k"),
- "filenames_train": "/ibex/ai/home/liz0l/codes/datasets/u4k/splits/train.txt",
- "input_height": 480, # ? will not be used (random crop)
- "input_width": 640, # ? will not be used (random crop)
- "filenames_val": "/ibex/ai/home/liz0l/codes/datasets/u4k/splits/val.txt",
- # "filenames_val": "/ibex/ai/home/liz0l/codes/datasets/u4k/splits/test.txt",
- "filenames_test": "/ibex/ai/home/liz0l/codes/datasets/u4k/splits/test.txt",
- "min_depth_eval": 1e-3,
- "max_depth_eval": 80,
- "min_depth_diff": -10,
- "max_depth_diff": 10,
-
- "do_random_rotate": True,
- "degree": 1.0,
- "do_kb_crop": False,
- "garg_crop": False,
- "eigen_crop": False,
-
- "num_sample_inout": 50000,
- # "num_sample_inout": 40000,
- "sampling_strategy": 'random',
- # "sampling_strategy": 'dda',
- "dilation_factor": 10,
-
- "use_rgb": False,
- "do_normalize": True, # do normalize in dataloader
- "do_input_resize": True
- },
- "mid": {
- "dataset": "mid",
- "min_depth": 1e-3, # originally 0.1
- "max_depth": 10,
- "data_path": os.path.join("/ibex/ai/home/liz0l/codes/datasets/middlebury"),
- "filenames_train": "/ibex/ai/home/liz0l/codes/datasets/middlebury/splits/train.txt",
- "input_height": 480, # ? will not be used (random crop)
- "input_width": 640, # ? will not be used (random crop)
- "filenames_val": "/ibex/ai/home/liz0l/codes/datasets/middlebury/splits/val.txt",
- "filenames_test": "/ibex/ai/home/liz0l/codes/datasets/middlebury/splits/test.txt",
- "min_depth_eval": 1e-3,
- "max_depth_eval": 10,
- "min_depth_diff": -10,
- "max_depth_diff": 10,
-
- "do_random_rotate": True,
- "degree": 1.0,
- "do_kb_crop": False,
- "garg_crop": False,
- "eigen_crop": False,
-
- "num_sample_inout": 50000,
- # "num_sample_inout": 40000,
- "sampling_strategy": 'random',
- # "sampling_strategy": 'dda',
- "dilation_factor": 10,
-
- "use_rgb": False,
- "do_normalize": True, # do normalize in dataloader
- "do_input_resize": True
- },
- "gta": {
- "dataset": "gta",
- "min_depth": 1e-3, # originally 0.1
- "max_depth": 80,
- "data_path": os.path.join("/ibex/ai/home/liz0l/codes/datasets/gta/GTAV_1080"),
- "filenames_train": "/ibex/ai/home/liz0l/codes/datasets/gta/GTAV_1080/train.txt",
- "input_height": 480, # ? will not be used (random crop)
- "input_width": 640, # ? will not be used (random crop)
- "filenames_val": "/ibex/ai/home/liz0l/codes/datasets/gta/GTAV_1080/val.txt",
- # "filenames_val": "/ibex/ai/home/liz0l/codes/datasets/u4k/splits/test.txt",
- "filenames_test": "/ibex/ai/home/liz0l/codes/datasets/gta/GTAV_1080/test.txt",
- "min_depth_eval": 1e-3,
- "max_depth_eval": 80,
- "min_depth_diff": -10,
- "max_depth_diff": 10,
-
- "do_random_rotate": True,
- "degree": 1.0,
- "do_kb_crop": False,
- "garg_crop": False,
- "eigen_crop": False,
-
- "num_sample_inout": 50000,
- # "num_sample_inout": 40000,
- "sampling_strategy": 'random',
- # "sampling_strategy": 'dda',
- "dilation_factor": 10,
-
- "use_rgb": False,
- "do_normalize": True, # do normalize in dataloader
- "do_input_resize": True
+ "eigen_crop": True
},
"ibims": {
"dataset": "ibims",
@@ -332,8 +236,7 @@ ALL_EVAL_DATASETS = ALL_INDOOR + ALL_OUTDOOR
COMMON_TRAINING_CONFIG = {
"dataset": "nyu",
"distributed": True,
- "workers": 24,
- # "workers": 6,
+ "workers": 16,
"clip_grad": 0.1,
"use_shared_dict": False,
"shared_dict": None,
@@ -345,11 +248,8 @@ COMMON_TRAINING_CONFIG = {
"translate_prob": 0.2,
"max_translation": 100,
- "validate_every": 1,
- # "validate_every": 1,
- # "log_images_every": 0.01,
- "log_images_every": 0.2,
- # "log_images_every": 0.05,
+ "validate_every": 0.25,
+ "log_images_every": 0.1,
"prefetch": False,
}
@@ -403,7 +303,7 @@ def parse_list(config, key, dtype=int):
), f"{key} should be a list of values dtype {dtype}. Given {config[key]} of type {type(config[key])} with values of type {[type(e) for e in config[key]]}."
-def get_model_config(model_name, model_version=None, model_cfg_path=""):
+def get_model_config(model_name, model_version=None):
"""Find and parse the .json config file for the model.
Args:
@@ -413,33 +313,26 @@ def get_model_config(model_name, model_version=None, model_cfg_path=""):
Returns:
easydict: the config dictionary for the model.
"""
-
- if model_cfg_path != "":
- config_file = model_cfg_path
- else:
- config_fname = f"config_{model_name}_{model_version}.json" if model_version is not None else f"config_{model_name}.json"
- config_file = os.path.join(ROOT, "models", model_name, config_fname)
+ config_fname = f"config_{model_name}_{model_version}.json" if model_version is not None else f"config_{model_name}.json"
+ config_file = os.path.join(ROOT, "models", model_name, config_fname)
if not os.path.exists(config_file):
return None
-
- # with open(config_file, "r") as f:
- # config = edict(json.load(f))
- # try to be more friendly
- with open(config_file, 'r') as f:
- config = edict(json5.load(f))
+
+ with open(config_file, "r") as f:
+ config = edict(json.load(f))
# handle dictionary inheritance
# only training config is supported for inheritance
if "inherit" in config.train and config.train.inherit is not None:
- inherit_config = get_model_config(config.train["inherit"], model_cfg_path=model_cfg_path).train
+ inherit_config = get_model_config(config.train["inherit"]).train
for key, value in inherit_config.items():
if key not in config.train:
config.train[key] = value
return edict(config)
-def update_model_config(config, mode, model_name, model_version=None, strict=False, model_cfg_path=""):
- model_config = get_model_config(model_name, model_version, model_cfg_path=model_cfg_path)
+def update_model_config(config, mode, model_name, model_version=None, strict=False):
+ model_config = get_model_config(model_name, model_version)
if model_config is not None:
config = {**config, **
flatten({**model_config.model, **model_config[mode]})}
@@ -458,7 +351,7 @@ KEYS_TYPE_BOOL = ["use_amp", "distributed", "use_shared_dict", "same_lr", "aug",
"prefetch", "cycle_momentum"] # Casting is not necessary as their int casted values in config are 0 or 1
-def get_config(model_name, mode='train', dataset=None, model_cfg_path=None, **overwrite_kwargs):
+def get_config(model_name, mode='train', dataset=None, **overwrite_kwargs):
"""Main entry point to get the config for the model.
Args:
@@ -478,24 +371,24 @@ def get_config(model_name, mode='train', dataset=None, model_cfg_path=None, **ov
easydict: The config dictionary for the model.
"""
- check_choices("Model", model_name, ["zoedepth", "zoedepth_nk", "zoedepth_custom"])
+
+ check_choices("Model", model_name, ["zoedepth", "zoedepth_nk"])
check_choices("Mode", mode, ["train", "infer", "eval"])
if mode == "train":
- check_choices("Dataset", dataset, ["nyu", "kitti", "mix", 'u4k', 'mid', 'gta', None])
+ check_choices("Dataset", dataset, ["nyu", "kitti", "mix", None])
-
config = flatten({**COMMON_CONFIG, **COMMON_TRAINING_CONFIG})
- config = update_model_config(config, mode, model_name, model_cfg_path=model_cfg_path)
+ config = update_model_config(config, mode, model_name)
# update with model version specific config
version_name = overwrite_kwargs.get("version_name", config["version_name"])
- config = update_model_config(config, mode, model_name, version_name, model_cfg_path=model_cfg_path)
+ config = update_model_config(config, mode, model_name, version_name)
# update with config version if specified
config_version = overwrite_kwargs.get("config_version", None)
if config_version is not None:
print("Overwriting config with config_version", config_version)
- config = update_model_config(config, mode, model_name, config_version, model_cfg_path=model_cfg_path)
+ config = update_model_config(config, mode, model_name, config_version)
# update with overwrite_kwargs
# Combined args are useful for hyperparameter search
@@ -542,69 +435,3 @@ def get_config(model_name, mode='train', dataset=None, model_cfg_path=None, **ov
def change_dataset(config, new_dataset):
config.update(DATASETS_CONFIG[new_dataset])
return config
-
-
-def get_config_user(model_name, mode='infer', model_cfg_path=None, **overwrite_kwargs):
- """Main entry point to get the config for the model.
-
- Args:
- model_name (str): name of the desired model.
- mode (str, optional): "train" or "infer". Defaults to 'train'.
- dataset (str, optional): If specified, the corresponding dataset configuration is loaded as well. Defaults to None.
-
- Keyword Args: key-value pairs of arguments to overwrite the default config.
-
- The order of precedence for overwriting the config is (Higher precedence first):
- # 1. overwrite_kwargs
- # 2. "config_version": Config file version if specified in overwrite_kwargs. The corresponding config loaded is config_{model_name}_{config_version}.json
- # 3. "version_name": Default Model version specific config specified in overwrite_kwargs. The corresponding config loaded is config_{model_name}_{version_name}.json
- # 4. common_config: Default config for all models specified in COMMON_CONFIG
-
- Returns:
- easydict: The config dictionary for the model.
- """
-
- check_choices("Model", model_name, ["zoedepth", "zoedepth_nk", "zoedepth_custom"])
- check_choices("Mode", mode, ["train", "infer", "eval"])
-
- config = flatten({**COMMON_CONFIG, **COMMON_TRAINING_CONFIG})
- config = update_model_config(config, mode, model_name, model_cfg_path=model_cfg_path)
-
- # update with model version specific config
- version_name = overwrite_kwargs.get("version_name", config["version_name"])
- config = update_model_config(config, mode, model_name, version_name, model_cfg_path=model_cfg_path)
-
- # update with config version if specified
- config_version = overwrite_kwargs.get("config_version", None)
- if config_version is not None:
- print("Overwriting config with config_version", config_version)
- config = update_model_config(config, mode, model_name, config_version, model_cfg_path=model_cfg_path)
-
- # update with overwrite_kwargs
- # Combined args are useful for hyperparameter search
- overwrite_kwargs = split_combined_args(overwrite_kwargs)
- config = {**config, **overwrite_kwargs}
-
- # Casting to bool # TODO: Not necessary. Remove and test
- for key in KEYS_TYPE_BOOL:
- if key in config:
- config[key] = bool(config[key])
-
- # Model specific post processing of config
- parse_list(config, "n_attractors")
-
- # adjust n_bins for each bin configuration if bin_conf is given and n_bins is passed in overwrite_kwargs
- if 'bin_conf' in config and 'n_bins' in overwrite_kwargs:
- bin_conf = config['bin_conf'] # list of dicts
- n_bins = overwrite_kwargs['n_bins']
- new_bin_conf = []
- for conf in bin_conf:
- conf['n_bins'] = n_bins
- new_bin_conf.append(conf)
- config['bin_conf'] = new_bin_conf
-
- config['model'] = model_name
- typed_config = {k: infer_type(v) for k, v in config.items()}
- # add hostname to config
- config['hostname'] = platform.node()
- return edict(typed_config)
diff --git a/zoedepth/utils/geometry.py b/zoedepth/utils/geometry.py
index ac34589041e866949c8f65eac85abae4e17ad89f..e3da8c75b5a8e39b4b58a4dcd827b84d79b9115c 100644
--- a/zoedepth/utils/geometry.py
+++ b/zoedepth/utils/geometry.py
@@ -24,21 +24,21 @@
import numpy as np
-def get_intrinsics(H,W,fov=55):
+def get_intrinsics(H,W):
"""
Intrinsics for a pinhole camera model.
Assume fov of 55 degrees and central principal point.
"""
- f = 0.5 * W / np.tan(0.5 * fov * np.pi / 180.0)
+ f = 0.5 * W / np.tan(0.5 * 55 * np.pi / 180.0)
cx = 0.5 * W
cy = 0.5 * H
return np.array([[f, 0, cx],
[0, f, cy],
[0, 0, 1]])
-def depth_to_points(depth, R=None, t=None, fov=55):
+def depth_to_points(depth, R=None, t=None):
- K = get_intrinsics(depth.shape[1], depth.shape[2], fov=fov)
+ K = get_intrinsics(depth.shape[1], depth.shape[2])
Kinv = np.linalg.inv(K)
if R is None:
R = np.eye(3)
diff --git a/zoedepth/utils/misc.py b/zoedepth/utils/misc.py
index 452ec23fbbcf4564d1eeb9309d0dcc4760469626..4bbe403d3669829eecdf658458c76aa5e87e2b33 100644
--- a/zoedepth/utils/misc.py
+++ b/zoedepth/utils/misc.py
@@ -42,9 +42,7 @@ import torch.nn as nn
import torch.utils.data.distributed
from PIL import Image
from torchvision.transforms import ToTensor
-import cv2
-import matplotlib
class RunningAverage:
def __init__(self):
@@ -96,7 +94,7 @@ class RunningAverageDict:
return {key: value.get_value() for key, value in self._dict.items()}
-def colorize(value, vmin=None, vmax=None, cmap='turbo_r', invalid_val=-99, invalid_mask=None, background_color=(128, 128, 128, 255), gamma_corrected=False, value_transform=None):
+def colorize(value, vmin=None, vmax=None, cmap='gray_r', invalid_val=-99, invalid_mask=None, background_color=(128, 128, 128, 255), gamma_corrected=False, value_transform=None):
"""Converts a depth map to a color image.
Args:
@@ -201,88 +199,7 @@ def compute_errors(gt, pred):
silog=silog, sq_rel=sq_rel)
-def shift_2d_replace(data, dx, dy, constant=False):
- shifted_data = np.roll(data, dx, axis=1)
- if dx < 0:
- shifted_data[:, dx:] = constant
- elif dx > 0:
- shifted_data[:, 0:dx] = constant
-
- shifted_data = np.roll(shifted_data, dy, axis=0)
- if dy < 0:
- shifted_data[dy:, :] = constant
- elif dy > 0:
- shifted_data[0:dy, :] = constant
- return shifted_data
-
-def soft_edge_error(pred, gt, radius=1):
- abs_diff=[]
- for i in range(-radius, radius + 1):
- for j in range(-radius, radius + 1):
- abs_diff.append(np.abs(shift_2d_replace(gt, i, j, 0) - pred))
- return np.minimum.reduce(abs_diff)
-
-def get_boundaries(disp, th=1., dilation=10):
- edges_y = np.logical_or(np.pad(np.abs(disp[1:, :] - disp[:-1, :]) > th, ((1, 0), (0, 0))),
- np.pad(np.abs(disp[:-1, :] - disp[1:, :]) > th, ((0, 1), (0, 0))))
- edges_x = np.logical_or(np.pad(np.abs(disp[:, 1:] - disp[:, :-1]) > th, ((0, 0), (1, 0))),
- np.pad(np.abs(disp[:, :-1] - disp[:,1:]) > th, ((0, 0), (0, 1))))
- edges = np.logical_or(edges_y, edges_x).astype(np.float32)
-
- if dilation > 0:
- kernel = np.ones((dilation, dilation), np.uint8)
- edges = cv2.dilate(edges, kernel, iterations=1)
-
- return edges
-
-
-@torch.no_grad()
-def scale_shift_linear(rendered_depth, predicted_depth, mask, fuse=True, return_params=False):
- """
- Optimize a scale and shift parameter in the least squares sense, such that rendered_depth and predicted_depth match.
- Formally, solves the following objective:
-
- min || (d * a + b) - d_hat ||
- a, b
-
- where d = 1 / predicted_depth, d_hat = 1 / rendered_depth
-
- :param rendered_depth: torch.Tensor (H, W)
- :param predicted_depth: torch.Tensor (H, W)
- :param mask: torch.Tensor (H, W) - True: valid points of rendered_depth, False: invalid points of rendered_depth (ignore)
- :param fuse: whether to fuse shifted/scaled predicted_depth with the rendered_depth
-
- :return: scale/shift corrected depth
- """
- if mask.sum() == 0:
- return predicted_depth
-
- # rendered_disparity = 1 / rendered_depth[mask].unsqueeze(-1)
- # predicted_disparity = 1 / predicted_depth[mask].unsqueeze(-1)
-
- rendered_disparity = rendered_depth[mask].unsqueeze(-1)
- predicted_disparity = predicted_depth[mask].unsqueeze(-1)
-
- X = torch.cat([predicted_disparity, torch.ones_like(predicted_disparity)], dim=1)
- XTX_inv = (X.T @ X).inverse()
- XTY = X.T @ rendered_disparity
- AB = XTX_inv @ XTY
-
- if return_params:
- return AB
-
- fixed_disparity = (predicted_depth) * AB[0] + AB[1]
- fixed_depth = fixed_disparity
-
- if fuse:
- fused_depth = torch.where(mask, rendered_depth, fixed_depth)
- return fused_depth
- else:
- return fixed_depth
-
-
-
-def compute_metrics(gt, pred, interpolate=True, garg_crop=False, eigen_crop=True, dataset='nyu', min_depth_eval=0.1, max_depth_eval=10, disp_gt_edges=None, pred_depths=None, **kwargs):
+def compute_metrics(gt, pred, interpolate=True, garg_crop=False, eigen_crop=True, dataset='nyu', min_depth_eval=0.1, max_depth_eval=10, **kwargs):
"""Compute metrics of predicted depth maps. Applies cropping and masking as necessary or specified via arguments. Refer to compute_errors for more details on metrics.
"""
if 'config' in kwargs:
@@ -294,7 +211,7 @@ def compute_metrics(gt, pred, interpolate=True, garg_crop=False, eigen_crop=True
if gt.shape[-2:] != pred.shape[-2:] and interpolate:
pred = nn.functional.interpolate(
- pred.unsqueeze(dim=0).unsqueeze(dim=0), gt.shape[-2:], mode='bilinear', align_corners=True).squeeze()
+ pred, gt.shape[-2:], mode='bilinear', align_corners=True)
pred = pred.squeeze().cpu().numpy()
pred[pred < min_depth_eval] = min_depth_eval
@@ -306,7 +223,6 @@ def compute_metrics(gt, pred, interpolate=True, garg_crop=False, eigen_crop=True
valid_mask = np.logical_and(
gt_depth > min_depth_eval, gt_depth < max_depth_eval)
- eval_mask = np.ones(valid_mask.shape)
if garg_crop or eigen_crop:
gt_height, gt_width = gt_depth.shape
eval_mask = np.zeros(valid_mask.shape)
@@ -326,36 +242,8 @@ def compute_metrics(gt, pred, interpolate=True, garg_crop=False, eigen_crop=True
else:
eval_mask = np.ones(valid_mask.shape)
valid_mask = np.logical_and(valid_mask, eval_mask)
+ return compute_errors(gt_depth[valid_mask], pred[valid_mask])
- # if dataset == 'nyu':
- # # pred = scale_shift_linear(torch.tensor(pred_depths), torch.tensor(pred), torch.tensor(valid_mask), fuse=False).numpy()
- # pred = scale_shift_linear(torch.tensor(gt), torch.tensor(pred), torch.tensor(valid_mask), fuse=False).numpy()
-
- metrics = compute_errors(gt_depth[valid_mask], pred[valid_mask])
-
- mask = valid_mask.squeeze() # squeeze
- gt = gt_depth
- pred = pred
- see_depth = 0
- if disp_gt_edges is None:
- print("Maybe we need edge maps from origin disp!")
- edges = get_boundaries(gt, th=0.08, dilation=0)
- else:
- edges = disp_gt_edges
-
- mask = np.logical_and(mask, edges)
- import matplotlib.pyplot as plt
- if mask.sum() > 0:
- see_depth = soft_edge_error(pred, gt)[mask].mean()
- metrics['see'] = see_depth
-
- return metrics
-
-# 'abs_rel': 0.07546425755890458, 'rmse': 0.2714709522322233, base zoe
-# 'abs_rel': 0.04409278385819647, 'rmse': 0.18093922881791188, base zoe+opt
-
-# patchfusion + pred scale-shift: 'abs_rel': 0.09078774519765959, 'rmse': 0.31991247948976803,
-# gt scale-shift abs_rel': 0.06316796072476771, 'rmse': 0.24189620860353886,
#################################### Model uilts ################################################
@@ -384,8 +272,6 @@ def parallelize(config, model, find_unused_parameters=True):
model = model.cuda(config.gpu)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[config.gpu], output_device=config.gpu,
find_unused_parameters=find_unused_parameters)
- # model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[config.gpu], output_device=config.gpu,
- # find_unused_parameters=True)
elif config.gpu is None:
# Use DP
@@ -479,20 +365,4 @@ def save_raw_16bit(depth, fpath="raw.png"):
depth = depth.astype(np.uint16)
depth = Image.fromarray(depth)
depth.save(fpath)
- print("Saved raw depth to", fpath)
-
-
-def generatemask(size, k_size=-1, sigma=-1, h_factor=0.03, w_factor=0.02):
- # Generates a Guassian mask
- mask = np.zeros(size, dtype=np.float32)
- if sigma == -1:
- sigma = int(size[0]/16)
- if k_size == -1:
- k_size = int(2 * np.ceil(2 * int(size[0]/16)) + 1)
- # mask[int(0.02*size[0]):size[0] - int(0.02*size[0]), int(0.015*size[1]): size[1] - int(0.015*size[1])] = 1
- mask[int(h_factor*size[0]):size[0] - int(h_factor*size[0]), int(w_factor*size[1]): size[1] - int(w_factor*size[1])] = 1
- mask = cv2.GaussianBlur(mask, (int(k_size), int(k_size)), sigma)
- mask = (mask - mask.min()) / (mask.max() - mask.min())
- mask = mask.astype(np.float32)
- return mask
-
+ print("Saved raw depth to", fpath)
\ No newline at end of file