Spaces:
Running
on
Zero
Running
on
Zero
Upload 299 files
Browse filesThis view is limited to 50 files because it contains too many changes.
See raw diff
- controlnet_aux-0.0.7.dist-info/INSTALLER +1 -0
- controlnet_aux-0.0.7.dist-info/LICENSE.txt +201 -0
- controlnet_aux-0.0.7.dist-info/METADATA +148 -0
- controlnet_aux-0.0.7.dist-info/RECORD +298 -0
- controlnet_aux-0.0.7.dist-info/REQUESTED +0 -0
- controlnet_aux-0.0.7.dist-info/WHEEL +5 -0
- controlnet_aux-0.0.7.dist-info/top_level.txt +1 -0
- controlnet_aux/.DS_Store +0 -0
- controlnet_aux/__init__.py +18 -0
- controlnet_aux/__pycache__/__init__.cpython-310.pyc +0 -0
- controlnet_aux/__pycache__/processor.cpython-310.pyc +0 -0
- controlnet_aux/__pycache__/util.cpython-310.pyc +0 -0
- controlnet_aux/canny/__init__.py +36 -0
- controlnet_aux/canny/__pycache__/__init__.cpython-310.pyc +0 -0
- controlnet_aux/dwpose/__init__.py +91 -0
- controlnet_aux/dwpose/__pycache__/__init__.cpython-310.pyc +0 -0
- controlnet_aux/dwpose/__pycache__/util.cpython-310.pyc +0 -0
- controlnet_aux/dwpose/__pycache__/wholebody.cpython-310.pyc +0 -0
- controlnet_aux/dwpose/util.py +303 -0
- controlnet_aux/dwpose/wholebody.py +121 -0
- controlnet_aux/hed/__init__.py +129 -0
- controlnet_aux/hed/__pycache__/__init__.cpython-310.pyc +0 -0
- controlnet_aux/leres/__init__.py +118 -0
- controlnet_aux/leres/__pycache__/__init__.cpython-310.pyc +0 -0
- controlnet_aux/leres/leres/Resnet.py +199 -0
- controlnet_aux/leres/leres/Resnext_torch.py +237 -0
- controlnet_aux/leres/leres/__init__.py +0 -0
- controlnet_aux/leres/leres/__pycache__/Resnet.cpython-310.pyc +0 -0
- controlnet_aux/leres/leres/__pycache__/Resnext_torch.cpython-310.pyc +0 -0
- controlnet_aux/leres/leres/__pycache__/__init__.cpython-310.pyc +0 -0
- controlnet_aux/leres/leres/__pycache__/depthmap.cpython-310.pyc +0 -0
- controlnet_aux/leres/leres/__pycache__/multi_depth_model_woauxi.cpython-310.pyc +0 -0
- controlnet_aux/leres/leres/__pycache__/net_tools.cpython-310.pyc +0 -0
- controlnet_aux/leres/leres/__pycache__/network_auxi.cpython-310.pyc +0 -0
- controlnet_aux/leres/leres/depthmap.py +548 -0
- controlnet_aux/leres/leres/multi_depth_model_woauxi.py +35 -0
- controlnet_aux/leres/leres/net_tools.py +54 -0
- controlnet_aux/leres/leres/network_auxi.py +417 -0
- controlnet_aux/leres/pix2pix/__init__.py +0 -0
- controlnet_aux/leres/pix2pix/__pycache__/__init__.cpython-310.pyc +0 -0
- controlnet_aux/leres/pix2pix/models/__init__.py +67 -0
- controlnet_aux/leres/pix2pix/models/__pycache__/__init__.cpython-310.pyc +0 -0
- controlnet_aux/leres/pix2pix/models/__pycache__/base_model.cpython-310.pyc +0 -0
- controlnet_aux/leres/pix2pix/models/__pycache__/base_model_hg.cpython-310.pyc +0 -0
- controlnet_aux/leres/pix2pix/models/__pycache__/networks.cpython-310.pyc +0 -0
- controlnet_aux/leres/pix2pix/models/__pycache__/pix2pix4depth_model.cpython-310.pyc +0 -0
- controlnet_aux/leres/pix2pix/models/base_model.py +244 -0
- controlnet_aux/leres/pix2pix/models/base_model_hg.py +58 -0
- controlnet_aux/leres/pix2pix/models/networks.py +623 -0
- controlnet_aux/leres/pix2pix/models/pix2pix4depth_model.py +155 -0
controlnet_aux-0.0.7.dist-info/INSTALLER
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
pip
|
controlnet_aux-0.0.7.dist-info/LICENSE.txt
ADDED
@@ -0,0 +1,201 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
Apache License
|
2 |
+
Version 2.0, January 2004
|
3 |
+
http://www.apache.org/licenses/
|
4 |
+
|
5 |
+
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
6 |
+
|
7 |
+
1. Definitions.
|
8 |
+
|
9 |
+
"License" shall mean the terms and conditions for use, reproduction,
|
10 |
+
and distribution as defined by Sections 1 through 9 of this document.
|
11 |
+
|
12 |
+
"Licensor" shall mean the copyright owner or entity authorized by
|
13 |
+
the copyright owner that is granting the License.
|
14 |
+
|
15 |
+
"Legal Entity" shall mean the union of the acting entity and all
|
16 |
+
other entities that control, are controlled by, or are under common
|
17 |
+
control with that entity. For the purposes of this definition,
|
18 |
+
"control" means (i) the power, direct or indirect, to cause the
|
19 |
+
direction or management of such entity, whether by contract or
|
20 |
+
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
21 |
+
outstanding shares, or (iii) beneficial ownership of such entity.
|
22 |
+
|
23 |
+
"You" (or "Your") shall mean an individual or Legal Entity
|
24 |
+
exercising permissions granted by this License.
|
25 |
+
|
26 |
+
"Source" form shall mean the preferred form for making modifications,
|
27 |
+
including but not limited to software source code, documentation
|
28 |
+
source, and configuration files.
|
29 |
+
|
30 |
+
"Object" form shall mean any form resulting from mechanical
|
31 |
+
transformation or translation of a Source form, including but
|
32 |
+
not limited to compiled object code, generated documentation,
|
33 |
+
and conversions to other media types.
|
34 |
+
|
35 |
+
"Work" shall mean the work of authorship, whether in Source or
|
36 |
+
Object form, made available under the License, as indicated by a
|
37 |
+
copyright notice that is included in or attached to the work
|
38 |
+
(an example is provided in the Appendix below).
|
39 |
+
|
40 |
+
"Derivative Works" shall mean any work, whether in Source or Object
|
41 |
+
form, that is based on (or derived from) the Work and for which the
|
42 |
+
editorial revisions, annotations, elaborations, or other modifications
|
43 |
+
represent, as a whole, an original work of authorship. For the purposes
|
44 |
+
of this License, Derivative Works shall not include works that remain
|
45 |
+
separable from, or merely link (or bind by name) to the interfaces of,
|
46 |
+
the Work and Derivative Works thereof.
|
47 |
+
|
48 |
+
"Contribution" shall mean any work of authorship, including
|
49 |
+
the original version of the Work and any modifications or additions
|
50 |
+
to that Work or Derivative Works thereof, that is intentionally
|
51 |
+
submitted to Licensor for inclusion in the Work by the copyright owner
|
52 |
+
or by an individual or Legal Entity authorized to submit on behalf of
|
53 |
+
the copyright owner. For the purposes of this definition, "submitted"
|
54 |
+
means any form of electronic, verbal, or written communication sent
|
55 |
+
to the Licensor or its representatives, including but not limited to
|
56 |
+
communication on electronic mailing lists, source code control systems,
|
57 |
+
and issue tracking systems that are managed by, or on behalf of, the
|
58 |
+
Licensor for the purpose of discussing and improving the Work, but
|
59 |
+
excluding communication that is conspicuously marked or otherwise
|
60 |
+
designated in writing by the copyright owner as "Not a Contribution."
|
61 |
+
|
62 |
+
"Contributor" shall mean Licensor and any individual or Legal Entity
|
63 |
+
on behalf of whom a Contribution has been received by Licensor and
|
64 |
+
subsequently incorporated within the Work.
|
65 |
+
|
66 |
+
2. Grant of Copyright License. Subject to the terms and conditions of
|
67 |
+
this License, each Contributor hereby grants to You a perpetual,
|
68 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
69 |
+
copyright license to reproduce, prepare Derivative Works of,
|
70 |
+
publicly display, publicly perform, sublicense, and distribute the
|
71 |
+
Work and such Derivative Works in Source or Object form.
|
72 |
+
|
73 |
+
3. Grant of Patent License. Subject to the terms and conditions of
|
74 |
+
this License, each Contributor hereby grants to You a perpetual,
|
75 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
76 |
+
(except as stated in this section) patent license to make, have made,
|
77 |
+
use, offer to sell, sell, import, and otherwise transfer the Work,
|
78 |
+
where such license applies only to those patent claims licensable
|
79 |
+
by such Contributor that are necessarily infringed by their
|
80 |
+
Contribution(s) alone or by combination of their Contribution(s)
|
81 |
+
with the Work to which such Contribution(s) was submitted. If You
|
82 |
+
institute patent litigation against any entity (including a
|
83 |
+
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
84 |
+
or a Contribution incorporated within the Work constitutes direct
|
85 |
+
or contributory patent infringement, then any patent licenses
|
86 |
+
granted to You under this License for that Work shall terminate
|
87 |
+
as of the date such litigation is filed.
|
88 |
+
|
89 |
+
4. Redistribution. You may reproduce and distribute copies of the
|
90 |
+
Work or Derivative Works thereof in any medium, with or without
|
91 |
+
modifications, and in Source or Object form, provided that You
|
92 |
+
meet the following conditions:
|
93 |
+
|
94 |
+
(a) You must give any other recipients of the Work or
|
95 |
+
Derivative Works a copy of this License; and
|
96 |
+
|
97 |
+
(b) You must cause any modified files to carry prominent notices
|
98 |
+
stating that You changed the files; and
|
99 |
+
|
100 |
+
(c) You must retain, in the Source form of any Derivative Works
|
101 |
+
that You distribute, all copyright, patent, trademark, and
|
102 |
+
attribution notices from the Source form of the Work,
|
103 |
+
excluding those notices that do not pertain to any part of
|
104 |
+
the Derivative Works; and
|
105 |
+
|
106 |
+
(d) If the Work includes a "NOTICE" text file as part of its
|
107 |
+
distribution, then any Derivative Works that You distribute must
|
108 |
+
include a readable copy of the attribution notices contained
|
109 |
+
within such NOTICE file, excluding those notices that do not
|
110 |
+
pertain to any part of the Derivative Works, in at least one
|
111 |
+
of the following places: within a NOTICE text file distributed
|
112 |
+
as part of the Derivative Works; within the Source form or
|
113 |
+
documentation, if provided along with the Derivative Works; or,
|
114 |
+
within a display generated by the Derivative Works, if and
|
115 |
+
wherever such third-party notices normally appear. The contents
|
116 |
+
of the NOTICE file are for informational purposes only and
|
117 |
+
do not modify the License. You may add Your own attribution
|
118 |
+
notices within Derivative Works that You distribute, alongside
|
119 |
+
or as an addendum to the NOTICE text from the Work, provided
|
120 |
+
that such additional attribution notices cannot be construed
|
121 |
+
as modifying the License.
|
122 |
+
|
123 |
+
You may add Your own copyright statement to Your modifications and
|
124 |
+
may provide additional or different license terms and conditions
|
125 |
+
for use, reproduction, or distribution of Your modifications, or
|
126 |
+
for any such Derivative Works as a whole, provided Your use,
|
127 |
+
reproduction, and distribution of the Work otherwise complies with
|
128 |
+
the conditions stated in this License.
|
129 |
+
|
130 |
+
5. Submission of Contributions. Unless You explicitly state otherwise,
|
131 |
+
any Contribution intentionally submitted for inclusion in the Work
|
132 |
+
by You to the Licensor shall be under the terms and conditions of
|
133 |
+
this License, without any additional terms or conditions.
|
134 |
+
Notwithstanding the above, nothing herein shall supersede or modify
|
135 |
+
the terms of any separate license agreement you may have executed
|
136 |
+
with Licensor regarding such Contributions.
|
137 |
+
|
138 |
+
6. Trademarks. This License does not grant permission to use the trade
|
139 |
+
names, trademarks, service marks, or product names of the Licensor,
|
140 |
+
except as required for reasonable and customary use in describing the
|
141 |
+
origin of the Work and reproducing the content of the NOTICE file.
|
142 |
+
|
143 |
+
7. Disclaimer of Warranty. Unless required by applicable law or
|
144 |
+
agreed to in writing, Licensor provides the Work (and each
|
145 |
+
Contributor provides its Contributions) on an "AS IS" BASIS,
|
146 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
147 |
+
implied, including, without limitation, any warranties or conditions
|
148 |
+
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
149 |
+
PARTICULAR PURPOSE. You are solely responsible for determining the
|
150 |
+
appropriateness of using or redistributing the Work and assume any
|
151 |
+
risks associated with Your exercise of permissions under this License.
|
152 |
+
|
153 |
+
8. Limitation of Liability. In no event and under no legal theory,
|
154 |
+
whether in tort (including negligence), contract, or otherwise,
|
155 |
+
unless required by applicable law (such as deliberate and grossly
|
156 |
+
negligent acts) or agreed to in writing, shall any Contributor be
|
157 |
+
liable to You for damages, including any direct, indirect, special,
|
158 |
+
incidental, or consequential damages of any character arising as a
|
159 |
+
result of this License or out of the use or inability to use the
|
160 |
+
Work (including but not limited to damages for loss of goodwill,
|
161 |
+
work stoppage, computer failure or malfunction, or any and all
|
162 |
+
other commercial damages or losses), even if such Contributor
|
163 |
+
has been advised of the possibility of such damages.
|
164 |
+
|
165 |
+
9. Accepting Warranty or Additional Liability. While redistributing
|
166 |
+
the Work or Derivative Works thereof, You may choose to offer,
|
167 |
+
and charge a fee for, acceptance of support, warranty, indemnity,
|
168 |
+
or other liability obligations and/or rights consistent with this
|
169 |
+
License. However, in accepting such obligations, You may act only
|
170 |
+
on Your own behalf and on Your sole responsibility, not on behalf
|
171 |
+
of any other Contributor, and only if You agree to indemnify,
|
172 |
+
defend, and hold each Contributor harmless for any liability
|
173 |
+
incurred by, or claims asserted against, such Contributor by reason
|
174 |
+
of your accepting any such warranty or additional liability.
|
175 |
+
|
176 |
+
END OF TERMS AND CONDITIONS
|
177 |
+
|
178 |
+
APPENDIX: How to apply the Apache License to your work.
|
179 |
+
|
180 |
+
To apply the Apache License to your work, attach the following
|
181 |
+
boilerplate notice, with the fields enclosed by brackets "[]"
|
182 |
+
replaced with your own identifying information. (Don't include
|
183 |
+
the brackets!) The text should be enclosed in the appropriate
|
184 |
+
comment syntax for the file format. We also recommend that a
|
185 |
+
file or class name and description of purpose be included on the
|
186 |
+
same "printed page" as the copyright notice for easier
|
187 |
+
identification within third-party archives.
|
188 |
+
|
189 |
+
Copyright [yyyy] [name of copyright owner]
|
190 |
+
|
191 |
+
Licensed under the Apache License, Version 2.0 (the "License");
|
192 |
+
you may not use this file except in compliance with the License.
|
193 |
+
You may obtain a copy of the License at
|
194 |
+
|
195 |
+
http://www.apache.org/licenses/LICENSE-2.0
|
196 |
+
|
197 |
+
Unless required by applicable law or agreed to in writing, software
|
198 |
+
distributed under the License is distributed on an "AS IS" BASIS,
|
199 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
200 |
+
See the License for the specific language governing permissions and
|
201 |
+
limitations under the License.
|
controlnet_aux-0.0.7.dist-info/METADATA
ADDED
@@ -0,0 +1,148 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
Metadata-Version: 2.1
|
2 |
+
Name: controlnet_aux
|
3 |
+
Version: 0.0.7
|
4 |
+
Summary: Auxillary models for controlnet
|
5 |
+
Home-page: https://github.com/patrickvonplaten/controlnet_aux
|
6 |
+
Author: The HuggingFace team
|
7 |
+
Author-email: patrick@huggingface.co
|
8 |
+
License: Apache
|
9 |
+
Keywords: deep learning
|
10 |
+
Classifier: Development Status :: 5 - Production/Stable
|
11 |
+
Classifier: Intended Audience :: Developers
|
12 |
+
Classifier: Intended Audience :: Education
|
13 |
+
Classifier: Intended Audience :: Science/Research
|
14 |
+
Classifier: License :: OSI Approved :: Apache Software License
|
15 |
+
Classifier: Operating System :: OS Independent
|
16 |
+
Classifier: Programming Language :: Python :: 3
|
17 |
+
Classifier: Programming Language :: Python :: 3.7
|
18 |
+
Classifier: Programming Language :: Python :: 3.8
|
19 |
+
Classifier: Programming Language :: Python :: 3.9
|
20 |
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
21 |
+
Requires-Python: >=3.7.0
|
22 |
+
Description-Content-Type: text/markdown
|
23 |
+
License-File: LICENSE.txt
|
24 |
+
Requires-Dist: torch
|
25 |
+
Requires-Dist: importlib-metadata
|
26 |
+
Requires-Dist: huggingface-hub
|
27 |
+
Requires-Dist: scipy
|
28 |
+
Requires-Dist: opencv-python
|
29 |
+
Requires-Dist: filelock
|
30 |
+
Requires-Dist: numpy
|
31 |
+
Requires-Dist: Pillow
|
32 |
+
Requires-Dist: einops
|
33 |
+
Requires-Dist: torchvision
|
34 |
+
Requires-Dist: timm
|
35 |
+
Requires-Dist: scikit-image
|
36 |
+
|
37 |
+
# ControlNet auxiliary models
|
38 |
+
|
39 |
+
This is a PyPi installable package of [lllyasviel's ControlNet Annotators](https://github.com/lllyasviel/ControlNet/tree/main/annotator)
|
40 |
+
|
41 |
+
The code is copy-pasted from the respective folders in https://github.com/lllyasviel/ControlNet/tree/main/annotator and connected to [the 🤗 Hub](https://huggingface.co/lllyasviel/Annotators).
|
42 |
+
|
43 |
+
All credit & copyright goes to https://github.com/lllyasviel .
|
44 |
+
|
45 |
+
## Install
|
46 |
+
|
47 |
+
```
|
48 |
+
pip install controlnet-aux==0.0.6
|
49 |
+
```
|
50 |
+
|
51 |
+
To support DWPose which is dependent on MMDetection, MMCV and MMPose
|
52 |
+
```
|
53 |
+
pip install -U openmim
|
54 |
+
mim install mmengine
|
55 |
+
mim install "mmcv>=2.0.1"
|
56 |
+
mim install "mmdet>=3.1.0"
|
57 |
+
mim install "mmpose>=1.1.0"
|
58 |
+
```
|
59 |
+
## Usage
|
60 |
+
|
61 |
+
|
62 |
+
You can use the processor class, which can load each of the auxiliary models with the following code
|
63 |
+
```python
|
64 |
+
import requests
|
65 |
+
from PIL import Image
|
66 |
+
from io import BytesIO
|
67 |
+
|
68 |
+
from controlnet_aux.processor import Processor
|
69 |
+
|
70 |
+
# load image
|
71 |
+
url = "https://huggingface.co/lllyasviel/sd-controlnet-openpose/resolve/main/images/pose.png"
|
72 |
+
|
73 |
+
response = requests.get(url)
|
74 |
+
img = Image.open(BytesIO(response.content)).convert("RGB").resize((512, 512))
|
75 |
+
|
76 |
+
# load processor from processor_id
|
77 |
+
# options are:
|
78 |
+
# ["canny", "depth_leres", "depth_leres++", "depth_midas", "depth_zoe", "lineart_anime",
|
79 |
+
# "lineart_coarse", "lineart_realistic", "mediapipe_face", "mlsd", "normal_bae", "normal_midas",
|
80 |
+
# "openpose", "openpose_face", "openpose_faceonly", "openpose_full", "openpose_hand",
|
81 |
+
# "scribble_hed, "scribble_pidinet", "shuffle", "softedge_hed", "softedge_hedsafe",
|
82 |
+
# "softedge_pidinet", "softedge_pidsafe", "dwpose"]
|
83 |
+
processor_id = 'scribble_hed'
|
84 |
+
processor = Processor(processor_id)
|
85 |
+
|
86 |
+
processed_image = processor(img, to_pil=True)
|
87 |
+
```
|
88 |
+
|
89 |
+
Each model can be loaded individually by importing and instantiating them as follows
|
90 |
+
```python
|
91 |
+
from PIL import Image
|
92 |
+
import requests
|
93 |
+
from io import BytesIO
|
94 |
+
from controlnet_aux import HEDdetector, MidasDetector, MLSDdetector, OpenposeDetector, PidiNetDetector, NormalBaeDetector, LineartDetector, LineartAnimeDetector, CannyDetector, ContentShuffleDetector, ZoeDetector, MediapipeFaceDetector, SamDetector, LeresDetector, DWposeDetector
|
95 |
+
|
96 |
+
# load image
|
97 |
+
url = "https://huggingface.co/lllyasviel/sd-controlnet-openpose/resolve/main/images/pose.png"
|
98 |
+
|
99 |
+
response = requests.get(url)
|
100 |
+
img = Image.open(BytesIO(response.content)).convert("RGB").resize((512, 512))
|
101 |
+
|
102 |
+
# load checkpoints
|
103 |
+
hed = HEDdetector.from_pretrained("lllyasviel/Annotators")
|
104 |
+
midas = MidasDetector.from_pretrained("lllyasviel/Annotators")
|
105 |
+
mlsd = MLSDdetector.from_pretrained("lllyasviel/Annotators")
|
106 |
+
open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
|
107 |
+
pidi = PidiNetDetector.from_pretrained("lllyasviel/Annotators")
|
108 |
+
normal_bae = NormalBaeDetector.from_pretrained("lllyasviel/Annotators")
|
109 |
+
lineart = LineartDetector.from_pretrained("lllyasviel/Annotators")
|
110 |
+
lineart_anime = LineartAnimeDetector.from_pretrained("lllyasviel/Annotators")
|
111 |
+
zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
|
112 |
+
sam = SamDetector.from_pretrained("ybelkada/segment-anything", subfolder="checkpoints")
|
113 |
+
mobile_sam = SamDetector.from_pretrained("dhkim2810/MobileSAM", model_type="vit_t", filename="mobile_sam.pt")
|
114 |
+
leres = LeresDetector.from_pretrained("lllyasviel/Annotators")
|
115 |
+
|
116 |
+
# specify configs, ckpts and device, or it will be downloaded automatically and use cpu by default
|
117 |
+
# det_config: ./src/controlnet_aux/dwpose/yolox_config/yolox_l_8xb8-300e_coco.py
|
118 |
+
# det_ckpt: https://download.openmmlab.com/mmdetection/v2.0/yolox/yolox_l_8x8_300e_coco/yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth
|
119 |
+
# pose_config: ./src/controlnet_aux/dwpose/dwpose_config/dwpose-l_384x288.py
|
120 |
+
# pose_ckpt: https://huggingface.co/wanghaofan/dw-ll_ucoco_384/resolve/main/dw-ll_ucoco_384.pth
|
121 |
+
import torch
|
122 |
+
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
123 |
+
dwpose = DWposeDetector(det_config=det_config, det_ckpt=det_ckpt, pose_config=pose_config, pose_ckpt=pose_ckpt, device=device)
|
124 |
+
|
125 |
+
# instantiate
|
126 |
+
canny = CannyDetector()
|
127 |
+
content = ContentShuffleDetector()
|
128 |
+
face_detector = MediapipeFaceDetector()
|
129 |
+
|
130 |
+
|
131 |
+
# process
|
132 |
+
processed_image_hed = hed(img)
|
133 |
+
processed_image_midas = midas(img)
|
134 |
+
processed_image_mlsd = mlsd(img)
|
135 |
+
processed_image_open_pose = open_pose(img, hand_and_face=True)
|
136 |
+
processed_image_pidi = pidi(img, safe=True)
|
137 |
+
processed_image_normal_bae = normal_bae(img)
|
138 |
+
processed_image_lineart = lineart(img, coarse=True)
|
139 |
+
processed_image_lineart_anime = lineart_anime(img)
|
140 |
+
processed_image_zoe = zoe(img)
|
141 |
+
processed_image_sam = sam(img)
|
142 |
+
processed_image_leres = leres(img)
|
143 |
+
|
144 |
+
processed_image_canny = canny(img)
|
145 |
+
processed_image_content = content(img)
|
146 |
+
processed_image_mediapipe_face = face_detector(img)
|
147 |
+
processed_image_dwpose = dwpose(img)
|
148 |
+
```
|
controlnet_aux-0.0.7.dist-info/RECORD
ADDED
@@ -0,0 +1,298 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
controlnet_aux-0.0.7.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
|
2 |
+
controlnet_aux-0.0.7.dist-info/LICENSE.txt,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
|
3 |
+
controlnet_aux-0.0.7.dist-info/METADATA,sha256=jXnomjJx-vD8tEzCdmwrSFSvUAGJ41bXh5febeFFrYw,5919
|
4 |
+
controlnet_aux-0.0.7.dist-info/RECORD,,
|
5 |
+
controlnet_aux-0.0.7.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
6 |
+
controlnet_aux-0.0.7.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
|
7 |
+
controlnet_aux-0.0.7.dist-info/top_level.txt,sha256=Oq49heF2JZSmikDijGkL5KPHbV9jBc3XD6MKYNVuHt0,15
|
8 |
+
controlnet_aux/__init__.py,sha256=Ul5ghIQH3gNS_mff11WkPgf8iB1s5KRiqCxQIYfcXBQ,583
|
9 |
+
controlnet_aux/__pycache__/__init__.cpython-310.pyc,,
|
10 |
+
controlnet_aux/__pycache__/processor.cpython-310.pyc,,
|
11 |
+
controlnet_aux/__pycache__/util.cpython-310.pyc,,
|
12 |
+
controlnet_aux/canny/__init__.py,sha256=8N4a_posVCZ6NX-mDosEwgd3vVYiKaT7rDnxesKqc6s,1353
|
13 |
+
controlnet_aux/canny/__pycache__/__init__.cpython-310.pyc,,
|
14 |
+
controlnet_aux/dwpose/__init__.py,sha256=xNGJM8sObZTMukpBvBt9Npp024595nkPpKp_FS_mlr8,2996
|
15 |
+
controlnet_aux/dwpose/__pycache__/__init__.cpython-310.pyc,,
|
16 |
+
controlnet_aux/dwpose/__pycache__/util.cpython-310.pyc,,
|
17 |
+
controlnet_aux/dwpose/__pycache__/wholebody.cpython-310.pyc,,
|
18 |
+
controlnet_aux/dwpose/util.py,sha256=zg0l31t77mdMnwCCOTJhst7y1Itv_LBReKrCy0CleMM,10660
|
19 |
+
controlnet_aux/dwpose/wholebody.py,sha256=SYmaQ-40XPKyszXeorC8pzy5PdyFnZTb23v0ZMl6Y8g,4576
|
20 |
+
controlnet_aux/hed/__init__.py,sha256=9v8svaB7PpLY3-1m-C1qY-5mLyzTLxh8PnNmPzxP5C8,5732
|
21 |
+
controlnet_aux/hed/__pycache__/__init__.cpython-310.pyc,,
|
22 |
+
controlnet_aux/leres/__init__.py,sha256=VW_do0VMp5r7AdgMgX1bIPwRXF51ipHseLQ9nBM2O7M,4371
|
23 |
+
controlnet_aux/leres/__pycache__/__init__.cpython-310.pyc,,
|
24 |
+
controlnet_aux/leres/leres/Resnet.py,sha256=toTIPZ0pM69Q8r-rSHFVI78BHrszpZkBZcUIsHSYXLo,6241
|
25 |
+
controlnet_aux/leres/leres/Resnext_torch.py,sha256=aJoXU83Ub85-a7EBKfsS0RcMHdilj_IzdFS2ZHkL2j8,8547
|
26 |
+
controlnet_aux/leres/leres/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
27 |
+
controlnet_aux/leres/leres/__pycache__/Resnet.cpython-310.pyc,,
|
28 |
+
controlnet_aux/leres/leres/__pycache__/Resnext_torch.cpython-310.pyc,,
|
29 |
+
controlnet_aux/leres/leres/__pycache__/__init__.cpython-310.pyc,,
|
30 |
+
controlnet_aux/leres/leres/__pycache__/depthmap.cpython-310.pyc,,
|
31 |
+
controlnet_aux/leres/leres/__pycache__/multi_depth_model_woauxi.cpython-310.pyc,,
|
32 |
+
controlnet_aux/leres/leres/__pycache__/net_tools.cpython-310.pyc,,
|
33 |
+
controlnet_aux/leres/leres/__pycache__/network_auxi.cpython-310.pyc,,
|
34 |
+
controlnet_aux/leres/leres/depthmap.py,sha256=1Tu1LyuLQtDe8tQpMSpvJkpKXM6pjRpk6UCKiAuA-Pk,22911
|
35 |
+
controlnet_aux/leres/leres/multi_depth_model_woauxi.py,sha256=KaS-QOYlq_jcC0PUBWZqE2q8-3NyCi3RHUxl-GkKOmg,1122
|
36 |
+
controlnet_aux/leres/leres/net_tools.py,sha256=lqvW-5OMqxlzxtuRHfNQ-0dzb3RCnkO0X4hdJO5gV1Q,2029
|
37 |
+
controlnet_aux/leres/leres/network_auxi.py,sha256=xCfaumwIjvP1IqcUyhF1ON7x1TMO2hVGHa-s3fiC1hY,16887
|
38 |
+
controlnet_aux/leres/pix2pix/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
39 |
+
controlnet_aux/leres/pix2pix/__pycache__/__init__.cpython-310.pyc,,
|
40 |
+
controlnet_aux/leres/pix2pix/models/__init__.py,sha256=4XKo1HO-R6mWOZeClCnB_AiGblbFtQ-D4K8CMfeH6VM,3095
|
41 |
+
controlnet_aux/leres/pix2pix/models/__pycache__/__init__.cpython-310.pyc,,
|
42 |
+
controlnet_aux/leres/pix2pix/models/__pycache__/base_model.cpython-310.pyc,,
|
43 |
+
controlnet_aux/leres/pix2pix/models/__pycache__/base_model_hg.cpython-310.pyc,,
|
44 |
+
controlnet_aux/leres/pix2pix/models/__pycache__/networks.cpython-310.pyc,,
|
45 |
+
controlnet_aux/leres/pix2pix/models/__pycache__/pix2pix4depth_model.cpython-310.pyc,,
|
46 |
+
controlnet_aux/leres/pix2pix/models/base_model.py,sha256=4fACbbRqRMk19vqUaQVxQmbfZFBFv7Ub19ZxIF3yPdc,10714
|
47 |
+
controlnet_aux/leres/pix2pix/models/base_model_hg.py,sha256=N9Gk59RoZimRDOHmhVsif6-CU5oYZBdKesPjt5RFmpc,1660
|
48 |
+
controlnet_aux/leres/pix2pix/models/networks.py,sha256=PJtJJKCVmbeEvSfPjnRIPjwAfNjVC7EeoHplYXOT-qc,28960
|
49 |
+
controlnet_aux/leres/pix2pix/models/pix2pix4depth_model.py,sha256=CH6iPV8rwyTxJxe_vPi0s0KG2HgF8ZKG9_cUEIyz1SY,7404
|
50 |
+
controlnet_aux/leres/pix2pix/options/__init__.py,sha256=M68EwLxPx7qB5z8chmqDL7pE0i3yjxOMWKbzncq5BLQ,136
|
51 |
+
controlnet_aux/leres/pix2pix/options/__pycache__/__init__.cpython-310.pyc,,
|
52 |
+
controlnet_aux/leres/pix2pix/options/__pycache__/base_options.cpython-310.pyc,,
|
53 |
+
controlnet_aux/leres/pix2pix/options/__pycache__/test_options.cpython-310.pyc,,
|
54 |
+
controlnet_aux/leres/pix2pix/options/base_options.py,sha256=yzv-j_Yrm2FGx6c070Mxt5UpJf8ixqFFO5PD_Vj3u68,9364
|
55 |
+
controlnet_aux/leres/pix2pix/options/test_options.py,sha256=C1kz8xIRlkFMUROYqJwZMNWNJ1AY0VQwOXyD-yMfyPs,1062
|
56 |
+
controlnet_aux/leres/pix2pix/util/__init__.py,sha256=dhGRbMeN26T0VlEGm21RpXh4sM4kmtbP_J511irCMQY,83
|
57 |
+
controlnet_aux/leres/pix2pix/util/__pycache__/__init__.cpython-310.pyc,,
|
58 |
+
controlnet_aux/leres/pix2pix/util/__pycache__/util.cpython-310.pyc,,
|
59 |
+
controlnet_aux/leres/pix2pix/util/util.py,sha256=OJxsfHVjldBhJerlgEMP0V8o2akioacQMq_4gmwxDRo,3110
|
60 |
+
controlnet_aux/lineart/__init__.py,sha256=ZIvfYsDChjPPqGUriyAosl5myVANk9qLvpBkHGl_3W4,5805
|
61 |
+
controlnet_aux/lineart/__pycache__/__init__.cpython-310.pyc,,
|
62 |
+
controlnet_aux/lineart_anime/__init__.py,sha256=WdHZ1fPldpNZ46xX0kqu8oJNRC9R-8g94PJWYva9Jlk,8204
|
63 |
+
controlnet_aux/lineart_anime/__pycache__/__init__.cpython-310.pyc,,
|
64 |
+
controlnet_aux/mediapipe_face/__init__.py,sha256=wBSM4asWVr8IzZ_ICx1rb4HJBNOxQen-ZjxUZ2pgKxw,1944
|
65 |
+
controlnet_aux/mediapipe_face/__pycache__/__init__.cpython-310.pyc,,
|
66 |
+
controlnet_aux/mediapipe_face/__pycache__/mediapipe_face_common.cpython-310.pyc,,
|
67 |
+
controlnet_aux/mediapipe_face/mediapipe_face_common.py,sha256=WlwPrTPgY2-mzsAcZ7mxNmFFL5SXTgW7eZmttU9Rfw4,7118
|
68 |
+
controlnet_aux/midas/__init__.py,sha256=hpHljQZyeTOVkC5_q8FiQu_jkE4lriaVITFdjDAUC70,3581
|
69 |
+
controlnet_aux/midas/__pycache__/__init__.cpython-310.pyc,,
|
70 |
+
controlnet_aux/midas/__pycache__/api.cpython-310.pyc,,
|
71 |
+
controlnet_aux/midas/__pycache__/utils.cpython-310.pyc,,
|
72 |
+
controlnet_aux/midas/api.py,sha256=i6Nf_Dhk2nRQB-MrqznhcsfXgIw0bltg1lqnRduNJ4o,5276
|
73 |
+
controlnet_aux/midas/midas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
74 |
+
controlnet_aux/midas/midas/__pycache__/__init__.cpython-310.pyc,,
|
75 |
+
controlnet_aux/midas/midas/__pycache__/base_model.cpython-310.pyc,,
|
76 |
+
controlnet_aux/midas/midas/__pycache__/blocks.cpython-310.pyc,,
|
77 |
+
controlnet_aux/midas/midas/__pycache__/dpt_depth.cpython-310.pyc,,
|
78 |
+
controlnet_aux/midas/midas/__pycache__/midas_net.cpython-310.pyc,,
|
79 |
+
controlnet_aux/midas/midas/__pycache__/midas_net_custom.cpython-310.pyc,,
|
80 |
+
controlnet_aux/midas/midas/__pycache__/transforms.cpython-310.pyc,,
|
81 |
+
controlnet_aux/midas/midas/__pycache__/vit.cpython-310.pyc,,
|
82 |
+
controlnet_aux/midas/midas/base_model.py,sha256=pQ3HlIKvvNAjo8QB0X0h5fi1wMUckGDvhJp1vRgWSWs,367
|
83 |
+
controlnet_aux/midas/midas/blocks.py,sha256=GxuQUbGMd4HbT1guTD53lrUe19rxflbBJepLXaJpszI,9242
|
84 |
+
controlnet_aux/midas/midas/dpt_depth.py,sha256=qQp4mwUvlnLB2-UQbCZxgC-aCF4RSIRnhE-ufxt9bDo,3154
|
85 |
+
controlnet_aux/midas/midas/midas_net.py,sha256=BEg_Q3fpMLNBSDon5qWEEMaDxoahd3dkWNJONd7sO5g,2709
|
86 |
+
controlnet_aux/midas/midas/midas_net_custom.py,sha256=DDqnMJQpERhghaz0sVQIU0RPHAV1uXFoAQHoZnyntss,5207
|
87 |
+
controlnet_aux/midas/midas/transforms.py,sha256=qz0gt3BlYGMy5qm6yWpE9g8ikjp4SNYrV7VpCVmpiD0,7869
|
88 |
+
controlnet_aux/midas/midas/vit.py,sha256=z73EcJCCuc_z-8uPwE1gRY4OGFr7418_Ys3LvOg0ILw,14625
|
89 |
+
controlnet_aux/midas/utils.py,sha256=Gc8v39unu2rYPURXHAGjHI2XP4S8eu6f31OuOsxo2wg,4582
|
90 |
+
controlnet_aux/mlsd/__init__.py,sha256=TCdI_UA7M2mz0EAAhwZmt_gXQ7ZPO1WE2xmuJTSQoAo,2816
|
91 |
+
controlnet_aux/mlsd/__pycache__/__init__.cpython-310.pyc,,
|
92 |
+
controlnet_aux/mlsd/__pycache__/utils.cpython-310.pyc,,
|
93 |
+
controlnet_aux/mlsd/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
94 |
+
controlnet_aux/mlsd/models/__pycache__/__init__.cpython-310.pyc,,
|
95 |
+
controlnet_aux/mlsd/models/__pycache__/mbv2_mlsd_large.cpython-310.pyc,,
|
96 |
+
controlnet_aux/mlsd/models/__pycache__/mbv2_mlsd_tiny.cpython-310.pyc,,
|
97 |
+
controlnet_aux/mlsd/models/mbv2_mlsd_large.py,sha256=iob6RKlcnYyyuKpTg32te98AkDaHLZ3lO6VbQD-ojdA,9678
|
98 |
+
controlnet_aux/mlsd/models/mbv2_mlsd_tiny.py,sha256=wgvWp-hppEwOY9eEkfsSyiyGak9s53-N_dIeruCZpOs,9180
|
99 |
+
controlnet_aux/mlsd/utils.py,sha256=CxyvNMBevuZbNA9nYOPnIAxdirAMg7XsvDdSBHyEqvo,24189
|
100 |
+
controlnet_aux/normalbae/__init__.py,sha256=NsHuxs6hv5wRkxNn9EaAyt0-CvQs8nRVrt1bXk9sjYE,3647
|
101 |
+
controlnet_aux/normalbae/__pycache__/__init__.cpython-310.pyc,,
|
102 |
+
controlnet_aux/normalbae/nets/NNET.py,sha256=byClNlH8c4SXnKvL8JFHKxzpOR_5uI4xzhj9dvAFr4U,597
|
103 |
+
controlnet_aux/normalbae/nets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
104 |
+
controlnet_aux/normalbae/nets/__pycache__/NNET.cpython-310.pyc,,
|
105 |
+
controlnet_aux/normalbae/nets/__pycache__/__init__.cpython-310.pyc,,
|
106 |
+
controlnet_aux/normalbae/nets/__pycache__/baseline.cpython-310.pyc,,
|
107 |
+
controlnet_aux/normalbae/nets/baseline.py,sha256=5L4kUGzh3npGZEHY2UvSxaNrxnRko5bRRA6jFy_f0sg,2980
|
108 |
+
controlnet_aux/normalbae/nets/submodules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
109 |
+
controlnet_aux/normalbae/nets/submodules/__pycache__/__init__.cpython-310.pyc,,
|
110 |
+
controlnet_aux/normalbae/nets/submodules/__pycache__/decoder.cpython-310.pyc,,
|
111 |
+
controlnet_aux/normalbae/nets/submodules/__pycache__/encoder.cpython-310.pyc,,
|
112 |
+
controlnet_aux/normalbae/nets/submodules/__pycache__/submodules.cpython-310.pyc,,
|
113 |
+
controlnet_aux/normalbae/nets/submodules/decoder.py,sha256=G3IIcYUhPxYhb8t357K33-5VailhKOao20GE4xwzUag,10480
|
114 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
115 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__pycache__/__init__.cpython-310.pyc,,
|
116 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__pycache__/caffe2_benchmark.cpython-310.pyc,,
|
117 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__pycache__/caffe2_validate.cpython-310.pyc,,
|
118 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__pycache__/hubconf.cpython-310.pyc,,
|
119 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__pycache__/onnx_export.cpython-310.pyc,,
|
120 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__pycache__/onnx_optimize.cpython-310.pyc,,
|
121 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__pycache__/onnx_to_caffe.cpython-310.pyc,,
|
122 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__pycache__/onnx_validate.cpython-310.pyc,,
|
123 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__pycache__/setup.cpython-310.pyc,,
|
124 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__pycache__/utils.cpython-310.pyc,,
|
125 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/__pycache__/validate.cpython-310.pyc,,
|
126 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/caffe2_benchmark.py,sha256=Oyl6kY8b9uYucJRADDtH-g5b9zra3RThj9sM_X1EZk8,2428
|
127 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/caffe2_validate.py,sha256=_AWt3nwO0YRUdhnIHLouDPdJa9dKE9dL99W7iEc7CVA,5993
|
128 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/__init__.py,sha256=Hc3iclTcGUhdkLL40qbM1q7xA24adCshWPpYaPqnUwk,206
|
129 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/__pycache__/__init__.cpython-310.pyc,,
|
130 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/__pycache__/config.cpython-310.pyc,,
|
131 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/__pycache__/conv2d_layers.cpython-310.pyc,,
|
132 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/__pycache__/efficientnet_builder.cpython-310.pyc,,
|
133 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/__pycache__/gen_efficientnet.cpython-310.pyc,,
|
134 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/__pycache__/helpers.cpython-310.pyc,,
|
135 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/__pycache__/mobilenetv3.cpython-310.pyc,,
|
136 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/__pycache__/model_factory.cpython-310.pyc,,
|
137 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/__pycache__/version.cpython-310.pyc,,
|
138 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/activations/__init__.py,sha256=Kxytt-sfEHpo-q4fN4nXqVjMLoh_Oi0Z3LGQ4t7xUA0,4170
|
139 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/activations/__pycache__/__init__.cpython-310.pyc,,
|
140 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/activations/__pycache__/activations.cpython-310.pyc,,
|
141 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/activations/__pycache__/activations_jit.cpython-310.pyc,,
|
142 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/activations/__pycache__/activations_me.cpython-310.pyc,,
|
143 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/activations/activations.py,sha256=m0KlKEZY5yNTDey-rZuNDmAONMsU547cwZx25X8dAeU,2690
|
144 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/activations/activations_jit.py,sha256=qGNmkNhseyrm8rG8VYv17Wx7SguHnqp48qfwoAy6A3I,2294
|
145 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/activations/activations_me.py,sha256=HzYPuHidaIq_L4K9PVATILvwQX5wwucUDw6H2QrMCTs,4549
|
146 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/config.py,sha256=yZmDSsKJFvPg4rWRczzxnzpOcDCMNcZ2YMUNAhpLj9M,3350
|
147 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/conv2d_layers.py,sha256=Gua8LqudrhgiqwzbVzQvOQQ269PKWuBThR1Pmriedhw,12093
|
148 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/efficientnet_builder.py,sha256=R2z-oO2zedCh7L6RFVYWX2vYxZsVklEya7EiuIG2sEw,26514
|
149 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/gen_efficientnet.py,sha256=NTrOBon1JucOWUH0yVkLbdiOHQY-0Tn01abmxQjF6Fk,59925
|
150 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/helpers.py,sha256=FVZf_KKCkSs7-a8bpbdT6Fad1kA4VmKX-BZ6tOiQTxY,2833
|
151 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/mobilenetv3.py,sha256=pc8g44XG_34ZNe2sWuKm6WUXRlPUdbHexGAcq9Wxsok,15009
|
152 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/model_factory.py,sha256=6GZnIMpsRQASZIarzNKbrjWar6xNxrI6Qoc3cW0xEeg,707
|
153 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/geffnet/version.py,sha256=ZTipNH6wyC6qsq1H_lZsvJyVruUa1nu3y4EysMdkpuU,22
|
154 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/hubconf.py,sha256=mB6ILmwUbvIzXntX1PTDForTuzKlnECiotfzY8YhneY,2730
|
155 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/onnx_export.py,sha256=xtBK_Kj33oGhvCHVBP-Neu3YV_mUrjXQ7_wnVtD-UQg,5821
|
156 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/onnx_optimize.py,sha256=4fzu495Pc1tGJDy_r1qH_o8oHkbIPS_7EQohpxfMZCA,2932
|
157 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/onnx_to_caffe.py,sha256=oA1-LhRcjURncJw7NU9e6z_HVTbDgSXY9FWyt-aa4rU,843
|
158 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/onnx_validate.py,sha256=OdZp8LTyOdZtoqvsBeFWrurKdJ1nDcYNo8aXi3VzzJw,4905
|
159 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/setup.py,sha256=L0WUqiHJv6dbzHeABlHXyRQvSDaMEGljWBjHYYHSKQw,1737
|
160 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/utils.py,sha256=1nqxoclxwtED2dQAo9FDtGc10EJaHTaji_Jm9ket0MI,1297
|
161 |
+
controlnet_aux/normalbae/nets/submodules/efficientnet_repo/validate.py,sha256=s9s91Yo1uRJPHFtNNOr41f75zKBm-nk2cCopUdarkjQ,6632
|
162 |
+
controlnet_aux/normalbae/nets/submodules/encoder.py,sha256=-WtqlVog5jQF9yI1lG6mbMrT9s2yghfjncm75OIPG8k,1060
|
163 |
+
controlnet_aux/normalbae/nets/submodules/submodules.py,sha256=vR6p8DkmmyQtjkdyO7Hapjd24cor0BpoPhvsFfrcd78,5703
|
164 |
+
controlnet_aux/open_pose/__init__.py,sha256=MCUwq46m4hPL3sgs9bfdAjqw7hd41iDtf7vOIQIiliA,9477
|
165 |
+
controlnet_aux/open_pose/__pycache__/__init__.cpython-310.pyc,,
|
166 |
+
controlnet_aux/open_pose/__pycache__/body.cpython-310.pyc,,
|
167 |
+
controlnet_aux/open_pose/__pycache__/face.cpython-310.pyc,,
|
168 |
+
controlnet_aux/open_pose/__pycache__/hand.cpython-310.pyc,,
|
169 |
+
controlnet_aux/open_pose/__pycache__/model.cpython-310.pyc,,
|
170 |
+
controlnet_aux/open_pose/__pycache__/util.cpython-310.pyc,,
|
171 |
+
controlnet_aux/open_pose/body.py,sha256=GovIZo3USiTrFh8zJGEQX691AVONEqWolOVQb8RujsU,12428
|
172 |
+
controlnet_aux/open_pose/face.py,sha256=BgyCSVVar3dQB4OX6AP8A0I63dAMIr7tY-O75mEOeLU,13491
|
173 |
+
controlnet_aux/open_pose/hand.py,sha256=2_GO-vqXgCR1xLHP5xxTyqqyaeCPykQtaCkW_BpUMxY,3210
|
174 |
+
controlnet_aux/open_pose/model.py,sha256=sJXhIRcb6xBmULyulxFMAX2THf8nM8nA0umUoOl8xaE,8743
|
175 |
+
controlnet_aux/open_pose/util.py,sha256=p1K-x8unCk-oWDZYXvEpHSSAwYcHSuiGT9ommS31qS0,13742
|
176 |
+
controlnet_aux/pidi/__init__.py,sha256=_1L6KYghjIycok5PdNAu5fXSSu4CrLB0KQKTIRvM5JA,3131
|
177 |
+
controlnet_aux/pidi/__pycache__/__init__.cpython-310.pyc,,
|
178 |
+
controlnet_aux/pidi/__pycache__/model.cpython-310.pyc,,
|
179 |
+
controlnet_aux/pidi/model.py,sha256=vjLYSwFxSJnoMmY-N0mRTrWI9adSXTPvk_wnor-dYo4,21813
|
180 |
+
controlnet_aux/processor.py,sha256=w_1rWD6LtouYGvS005dO1xMb47ABxGSY3pSvVCOmKC0,6259
|
181 |
+
controlnet_aux/segment_anything/__init__.py,sha256=QErgZBr42AtyojiXrCyPP718UyZxhyzLa4Ee8Ijmw1A,3364
|
182 |
+
controlnet_aux/segment_anything/__pycache__/__init__.cpython-310.pyc,,
|
183 |
+
controlnet_aux/segment_anything/__pycache__/automatic_mask_generator.cpython-310.pyc,,
|
184 |
+
controlnet_aux/segment_anything/__pycache__/build_sam.cpython-310.pyc,,
|
185 |
+
controlnet_aux/segment_anything/__pycache__/predictor.cpython-310.pyc,,
|
186 |
+
controlnet_aux/segment_anything/automatic_mask_generator.py,sha256=ccaY0rbL9aYur6XxIQb987_2dhAq3xrtIoffOzmQtOA,15148
|
187 |
+
controlnet_aux/segment_anything/build_sam.py,sha256=vBlel82yAoa3MSQh95DLuUNlnKS_W6u2Db68ZADPMu0,4695
|
188 |
+
controlnet_aux/segment_anything/modeling/__init__.py,sha256=b50JHHIgNqM-rCrnZ-6ipf35QMb0elcQ-bBEpkBn8lI,419
|
189 |
+
controlnet_aux/segment_anything/modeling/__pycache__/__init__.cpython-310.pyc,,
|
190 |
+
controlnet_aux/segment_anything/modeling/__pycache__/common.cpython-310.pyc,,
|
191 |
+
controlnet_aux/segment_anything/modeling/__pycache__/image_encoder.cpython-310.pyc,,
|
192 |
+
controlnet_aux/segment_anything/modeling/__pycache__/mask_decoder.cpython-310.pyc,,
|
193 |
+
controlnet_aux/segment_anything/modeling/__pycache__/prompt_encoder.cpython-310.pyc,,
|
194 |
+
controlnet_aux/segment_anything/modeling/__pycache__/sam.cpython-310.pyc,,
|
195 |
+
controlnet_aux/segment_anything/modeling/__pycache__/tiny_vit_sam.cpython-310.pyc,,
|
196 |
+
controlnet_aux/segment_anything/modeling/__pycache__/transformer.cpython-310.pyc,,
|
197 |
+
controlnet_aux/segment_anything/modeling/common.py,sha256=WaU7-l0KTURt9_K7V9LpWAixMGAVkrk0YA-2Y9JWzhc,1479
|
198 |
+
controlnet_aux/segment_anything/modeling/image_encoder.py,sha256=FmWtopAnIsmozVZshA4yTWI-8rEvfZrp84skrtDjmIY,14420
|
199 |
+
controlnet_aux/segment_anything/modeling/mask_decoder.py,sha256=aiRV6ZBRXGegYRh8ksI2uHmxifIe9uv84WJ8rpprL5M,6615
|
200 |
+
controlnet_aux/segment_anything/modeling/prompt_encoder.py,sha256=vY8vWQPwZH9Vwhxgi8uzNO7CkB6-dqpBBhscUT30oG8,8594
|
201 |
+
controlnet_aux/segment_anything/modeling/sam.py,sha256=ir7679WhwO3wbJdJvXRqkwAtkps5D4yVzuSYZXECI40,7283
|
202 |
+
controlnet_aux/segment_anything/modeling/tiny_vit_sam.py,sha256=pnWwdNi5gFyc9qCnP-c_1WAOCRzApzFrNyCgHgbKjb4,24707
|
203 |
+
controlnet_aux/segment_anything/modeling/transformer.py,sha256=VmqPbbi2CDmN6TlUi9D0hbIXUv2ivdo2kA5kRO-TjRE,8397
|
204 |
+
controlnet_aux/segment_anything/predictor.py,sha256=8yzEm8YHkrRq8aEPKnLsJuUslb8vpeAUlckq1bCjGm0,11633
|
205 |
+
controlnet_aux/segment_anything/utils/__init__.py,sha256=NL2AacVHZOe41zp4kF2-ZGcUCi9zFwh1Eo9spNjN0Ko,197
|
206 |
+
controlnet_aux/segment_anything/utils/__pycache__/__init__.cpython-310.pyc,,
|
207 |
+
controlnet_aux/segment_anything/utils/__pycache__/amg.cpython-310.pyc,,
|
208 |
+
controlnet_aux/segment_anything/utils/__pycache__/onnx.cpython-310.pyc,,
|
209 |
+
controlnet_aux/segment_anything/utils/__pycache__/transforms.cpython-310.pyc,,
|
210 |
+
controlnet_aux/segment_anything/utils/amg.py,sha256=XBD43WP6y2Qv8JD2FCazWzxsOgfnKDlZXLBNdIKLHTk,12712
|
211 |
+
controlnet_aux/segment_anything/utils/onnx.py,sha256=tEEKwdyae7OgmCdSD6CEVPmzl2CHdeL-Tr-3yNuwJmY,5812
|
212 |
+
controlnet_aux/segment_anything/utils/transforms.py,sha256=JqZJ8ryopKJLTVZhdIeppUxCeeVZsrgTyd449YpgEHw,3972
|
213 |
+
controlnet_aux/shuffle/__init__.py,sha256=MdYBzt9T_1nZSAzWJclt7AVNSYOo5pp4IynMflWooxo,3330
|
214 |
+
controlnet_aux/shuffle/__pycache__/__init__.cpython-310.pyc,,
|
215 |
+
controlnet_aux/util.py,sha256=3ncv5FAW0C7qY7DXPSOKRjneazIzgx6vMA30QxLeMng,5434
|
216 |
+
controlnet_aux/zoe/__init__.py,sha256=_Gu8Ngdhpfy6poY_-O89bW3u9_esEvR4G0te3cY9Jhk,2913
|
217 |
+
controlnet_aux/zoe/__pycache__/__init__.cpython-310.pyc,,
|
218 |
+
controlnet_aux/zoe/zoedepth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
219 |
+
controlnet_aux/zoe/zoedepth/__pycache__/__init__.cpython-310.pyc,,
|
220 |
+
controlnet_aux/zoe/zoedepth/models/__init__.py,sha256=03AOka2O87LVtNxMUP36AtSUBXQH0bj-UrVg3fD5DRA,1154
|
221 |
+
controlnet_aux/zoe/zoedepth/models/__pycache__/__init__.cpython-310.pyc,,
|
222 |
+
controlnet_aux/zoe/zoedepth/models/__pycache__/builder.cpython-310.pyc,,
|
223 |
+
controlnet_aux/zoe/zoedepth/models/__pycache__/depth_model.cpython-310.pyc,,
|
224 |
+
controlnet_aux/zoe/zoedepth/models/__pycache__/model_io.cpython-310.pyc,,
|
225 |
+
controlnet_aux/zoe/zoedepth/models/base_models/__init__.py,sha256=03AOka2O87LVtNxMUP36AtSUBXQH0bj-UrVg3fD5DRA,1154
|
226 |
+
controlnet_aux/zoe/zoedepth/models/base_models/__pycache__/__init__.cpython-310.pyc,,
|
227 |
+
controlnet_aux/zoe/zoedepth/models/base_models/__pycache__/midas.cpython-310.pyc,,
|
228 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas.py,sha256=2qxx3yjewKyY7G6Wbez5RdbkFa2BbeABQ87-zv0ZstY,15248
|
229 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
230 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/__pycache__/__init__.cpython-310.pyc,,
|
231 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/__pycache__/hubconf.cpython-310.pyc,,
|
232 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/hubconf.py,sha256=Ft6H6I52XsOzNNrPIa4hmA8K0a8mIVngva-7kq61Tyo,13757
|
233 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
234 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/__pycache__/__init__.cpython-310.pyc,,
|
235 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/__pycache__/base_model.cpython-310.pyc,,
|
236 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/__pycache__/blocks.cpython-310.pyc,,
|
237 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/__pycache__/dpt_depth.cpython-310.pyc,,
|
238 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/__pycache__/midas_net.cpython-310.pyc,,
|
239 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/__pycache__/midas_net_custom.cpython-310.pyc,,
|
240 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/__pycache__/model_loader.cpython-310.pyc,,
|
241 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/__pycache__/transforms.cpython-310.pyc,,
|
242 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
243 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/__pycache__/__init__.cpython-310.pyc,,
|
244 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/__pycache__/beit.cpython-310.pyc,,
|
245 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/__pycache__/levit.cpython-310.pyc,,
|
246 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/__pycache__/next_vit.cpython-310.pyc,,
|
247 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/__pycache__/swin.cpython-310.pyc,,
|
248 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/__pycache__/swin2.cpython-310.pyc,,
|
249 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/__pycache__/swin_common.cpython-310.pyc,,
|
250 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/__pycache__/utils.cpython-310.pyc,,
|
251 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/__pycache__/vit.cpython-310.pyc,,
|
252 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/beit.py,sha256=zn6u7vz9nd3uZliDK4IBBmA7mlyrN0lP0uzhT9Yll7s,6909
|
253 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/levit.py,sha256=eC5r45FZgbokNI3XbGk_wU7mAvOcTXynoDTVLUoC1hM,3436
|
254 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/next_vit.py,sha256=wvNsvbXK_lLK34_xlCqFJIqmBEBdcueZ3Jg9YRJDMS8,1045
|
255 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/swin.py,sha256=EZ2fEAFeUEHxiscBr60033r0Iu6zEw65H27m6_CtXVI,333
|
256 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/swin2.py,sha256=oBaPyT8x5H9kH6lVSTgBHePmUApF6U2M8tP0cDCMakY,931
|
257 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/swin_common.py,sha256=4n_WOAFy7TQ8EDpa5VuHUWOzOYSkYISKotKL8upssuA,1576
|
258 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/utils.py,sha256=09TxrevO_QsYy6fVOiN8X_Qjjkdxl9u7RNXJLo4wHPs,7284
|
259 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/vit.py,sha256=Z-xSs9mo54Nsa8TboOh3etzSlkaH8wv9ArQjtqzkxHc,6899
|
260 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/base_model.py,sha256=pQ3HlIKvvNAjo8QB0X0h5fi1wMUckGDvhJp1vRgWSWs,367
|
261 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/blocks.py,sha256=g63_blM815Pz5T74XlrJ4qE3qGtHAMSajpjA8mSXWDk,12792
|
262 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/dpt_depth.py,sha256=c-_J8FrB6wwAkr2J933sGryGh4IedIvbUjbH-DGZVms,6099
|
263 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/midas_net.py,sha256=BEg_Q3fpMLNBSDon5qWEEMaDxoahd3dkWNJONd7sO5g,2709
|
264 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/midas_net_custom.py,sha256=DDqnMJQpERhghaz0sVQIU0RPHAV1uXFoAQHoZnyntss,5207
|
265 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/model_loader.py,sha256=RGRatKeaaIxkWXyteVpRjNrpeZEddMxLZNKFthfR69w,8552
|
266 |
+
controlnet_aux/zoe/zoedepth/models/base_models/midas_repo/midas/transforms.py,sha256=qz0gt3BlYGMy5qm6yWpE9g8ikjp4SNYrV7VpCVmpiD0,7869
|
267 |
+
controlnet_aux/zoe/zoedepth/models/builder.py,sha256=cnFFNBLC5x_5yfdKHGJfGh30btOdTN61a5aY6S9SVlA,2390
|
268 |
+
controlnet_aux/zoe/zoedepth/models/depth_model.py,sha256=EiuMLMKygdjNCT26o45MmS6TkvKsoLh7DfT9-KHhzrM,7362
|
269 |
+
controlnet_aux/zoe/zoedepth/models/layers/__init__.py,sha256=UmFTzx4rn06v3lvPX_Vm6BEkELgTcJUmsw07nroNchU,1153
|
270 |
+
controlnet_aux/zoe/zoedepth/models/layers/__pycache__/__init__.cpython-310.pyc,,
|
271 |
+
controlnet_aux/zoe/zoedepth/models/layers/__pycache__/attractor.cpython-310.pyc,,
|
272 |
+
controlnet_aux/zoe/zoedepth/models/layers/__pycache__/dist_layers.cpython-310.pyc,,
|
273 |
+
controlnet_aux/zoe/zoedepth/models/layers/__pycache__/localbins_layers.cpython-310.pyc,,
|
274 |
+
controlnet_aux/zoe/zoedepth/models/layers/__pycache__/patch_transformer.cpython-310.pyc,,
|
275 |
+
controlnet_aux/zoe/zoedepth/models/layers/attractor.py,sha256=BCvFrRabNe59borFWzswwNxFNE07P0h6vQeuLK5hf9U,8693
|
276 |
+
controlnet_aux/zoe/zoedepth/models/layers/dist_layers.py,sha256=Pn-kHTIAOeqYnV8HjoCD4twAV-YHZheRdu1mgdd4njQ,4838
|
277 |
+
controlnet_aux/zoe/zoedepth/models/layers/localbins_layers.py,sha256=qvRKNj8ISW071AhuuAaiVH6sA5hB3l_zSBKg7_LHzgw,6733
|
278 |
+
controlnet_aux/zoe/zoedepth/models/layers/patch_transformer.py,sha256=Abnsjih9fmETt9rHjpLL3ksezWye3VnnYaFo64yBNT4,4163
|
279 |
+
controlnet_aux/zoe/zoedepth/models/model_io.py,sha256=ub_VLEFmc-rQoP89wn_IzZG6RMtxJP2VbWdjgaMPGSw,3438
|
280 |
+
controlnet_aux/zoe/zoedepth/models/zoedepth/__init__.py,sha256=LT6C97t3t19rMc0Bl9zuvW-y2WVOo9beXBeDST3ndlM,1270
|
281 |
+
controlnet_aux/zoe/zoedepth/models/zoedepth/__pycache__/__init__.cpython-310.pyc,,
|
282 |
+
controlnet_aux/zoe/zoedepth/models/zoedepth/__pycache__/zoedepth_v1.cpython-310.pyc,,
|
283 |
+
controlnet_aux/zoe/zoedepth/models/zoedepth/config_zoedepth.json,sha256=Po3PsHF_mQd3_rrP3cYWHHOpCaF-fEiGJDh1z7G5SdE,1587
|
284 |
+
controlnet_aux/zoe/zoedepth/models/zoedepth/config_zoedepth_kitti.json,sha256=j0hzr9tCQjaFK-rESbDQW1e1I_THvW4g5q3lNf1buG8,556
|
285 |
+
controlnet_aux/zoe/zoedepth/models/zoedepth/zoedepth_v1.py,sha256=Own8PG2LGJkfQTHhPn-02HHv551fyHqMijw2ui7b6r0,12630
|
286 |
+
controlnet_aux/zoe/zoedepth/models/zoedepth_nk/__init__.py,sha256=xAwaR0Y0EEK8GRuZe-jD2s0tL56sO5aC9ZzvTAMDe5A,1276
|
287 |
+
controlnet_aux/zoe/zoedepth/models/zoedepth_nk/__pycache__/__init__.cpython-310.pyc,,
|
288 |
+
controlnet_aux/zoe/zoedepth/models/zoedepth_nk/__pycache__/zoedepth_nk_v1.cpython-310.pyc,,
|
289 |
+
controlnet_aux/zoe/zoedepth/models/zoedepth_nk/config_zoedepth_nk.json,sha256=64JBDGPvxnNoGlMJfhc9UXv6loAPE_MsrCrP1BoMMrc,1968
|
290 |
+
controlnet_aux/zoe/zoedepth/models/zoedepth_nk/zoedepth_nk_v1.py,sha256=iI8QlqEyq1RT2phk8gsSz-m5tjMxCSOM2mjf7IGRLZE,16264
|
291 |
+
controlnet_aux/zoe/zoedepth/utils/__init__.py,sha256=03AOka2O87LVtNxMUP36AtSUBXQH0bj-UrVg3fD5DRA,1154
|
292 |
+
controlnet_aux/zoe/zoedepth/utils/__pycache__/__init__.cpython-310.pyc,,
|
293 |
+
controlnet_aux/zoe/zoedepth/utils/__pycache__/arg_utils.cpython-310.pyc,,
|
294 |
+
controlnet_aux/zoe/zoedepth/utils/__pycache__/config.cpython-310.pyc,,
|
295 |
+
controlnet_aux/zoe/zoedepth/utils/arg_utils.py,sha256=_Rxk7MrLSJ2pnlWh2uBkje2BFbt3YovrLeEEbkN3XQI,624
|
296 |
+
controlnet_aux/zoe/zoedepth/utils/config.py,sha256=YVwbJMlnty9S81gaoN9sPVhNyA4PVTk9LnLdqtsonSI,16310
|
297 |
+
controlnet_aux/zoe/zoedepth/utils/easydict/__init__.py,sha256=iBUqkjdwdFFYCobeoRxCTMUZ4qFHIlzaGVz0kQkbejc,3426
|
298 |
+
controlnet_aux/zoe/zoedepth/utils/easydict/__pycache__/__init__.cpython-310.pyc,,
|
controlnet_aux-0.0.7.dist-info/REQUESTED
ADDED
File without changes
|
controlnet_aux-0.0.7.dist-info/WHEEL
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
Wheel-Version: 1.0
|
2 |
+
Generator: bdist_wheel (0.43.0)
|
3 |
+
Root-Is-Purelib: true
|
4 |
+
Tag: py3-none-any
|
5 |
+
|
controlnet_aux-0.0.7.dist-info/top_level.txt
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
controlnet_aux
|
controlnet_aux/.DS_Store
ADDED
Binary file (6.15 kB). View file
|
|
controlnet_aux/__init__.py
ADDED
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
__version__ = "0.0.7"
|
2 |
+
|
3 |
+
from .hed import HEDdetector
|
4 |
+
# from .leres import LeresDetector
|
5 |
+
# from .lineart import LineartDetector
|
6 |
+
# from .lineart_anime import LineartAnimeDetector
|
7 |
+
# from .midas import MidasDetector
|
8 |
+
# from .mlsd import MLSDdetector
|
9 |
+
from .normalbae import NormalBaeDetector
|
10 |
+
# from .open_pose import OpenposeDetector
|
11 |
+
# from .pidi import PidiNetDetector
|
12 |
+
# from .zoe import ZoeDetector
|
13 |
+
|
14 |
+
# from .canny import CannyDetector
|
15 |
+
# from .mediapipe_face import MediapipeFaceDetector
|
16 |
+
# from .segment_anything import SamDetector
|
17 |
+
# from .shuffle import ContentShuffleDetector
|
18 |
+
# from .dwpose import DWposeDetector
|
controlnet_aux/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (942 Bytes). View file
|
|
controlnet_aux/__pycache__/processor.cpython-310.pyc
ADDED
Binary file (4.46 kB). View file
|
|
controlnet_aux/__pycache__/util.cpython-310.pyc
ADDED
Binary file (7.82 kB). View file
|
|
controlnet_aux/canny/__init__.py
ADDED
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import warnings
|
2 |
+
import cv2
|
3 |
+
import numpy as np
|
4 |
+
from PIL import Image
|
5 |
+
from ..util import HWC3, resize_image
|
6 |
+
|
7 |
+
class CannyDetector:
|
8 |
+
def __call__(self, input_image=None, low_threshold=100, high_threshold=200, detect_resolution=512, image_resolution=512, output_type=None, **kwargs):
|
9 |
+
if "img" in kwargs:
|
10 |
+
warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning)
|
11 |
+
input_image = kwargs.pop("img")
|
12 |
+
|
13 |
+
if input_image is None:
|
14 |
+
raise ValueError("input_image must be defined.")
|
15 |
+
|
16 |
+
if not isinstance(input_image, np.ndarray):
|
17 |
+
input_image = np.array(input_image, dtype=np.uint8)
|
18 |
+
output_type = output_type or "pil"
|
19 |
+
else:
|
20 |
+
output_type = output_type or "np"
|
21 |
+
|
22 |
+
input_image = HWC3(input_image)
|
23 |
+
input_image = resize_image(input_image, detect_resolution)
|
24 |
+
|
25 |
+
detected_map = cv2.Canny(input_image, low_threshold, high_threshold)
|
26 |
+
detected_map = HWC3(detected_map)
|
27 |
+
|
28 |
+
img = resize_image(input_image, image_resolution)
|
29 |
+
H, W, C = img.shape
|
30 |
+
|
31 |
+
detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
|
32 |
+
|
33 |
+
if output_type == "pil":
|
34 |
+
detected_map = Image.fromarray(detected_map)
|
35 |
+
|
36 |
+
return detected_map
|
controlnet_aux/canny/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (1.33 kB). View file
|
|
controlnet_aux/dwpose/__init__.py
ADDED
@@ -0,0 +1,91 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Openpose
|
2 |
+
# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose
|
3 |
+
# 2nd Edited by https://github.com/Hzzone/pytorch-openpose
|
4 |
+
# 3rd Edited by ControlNet
|
5 |
+
# 4th Edited by ControlNet (added face and correct hands)
|
6 |
+
|
7 |
+
import os
|
8 |
+
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
|
9 |
+
|
10 |
+
import cv2
|
11 |
+
import torch
|
12 |
+
import numpy as np
|
13 |
+
from PIL import Image
|
14 |
+
|
15 |
+
from ..util import HWC3, resize_image
|
16 |
+
from . import util
|
17 |
+
|
18 |
+
|
19 |
+
def draw_pose(pose, H, W):
|
20 |
+
bodies = pose['bodies']
|
21 |
+
faces = pose['faces']
|
22 |
+
hands = pose['hands']
|
23 |
+
candidate = bodies['candidate']
|
24 |
+
subset = bodies['subset']
|
25 |
+
|
26 |
+
canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
|
27 |
+
canvas = util.draw_bodypose(canvas, candidate, subset)
|
28 |
+
canvas = util.draw_handpose(canvas, hands)
|
29 |
+
canvas = util.draw_facepose(canvas, faces)
|
30 |
+
|
31 |
+
return canvas
|
32 |
+
|
33 |
+
class DWposeDetector:
|
34 |
+
def __init__(self, det_config=None, det_ckpt=None, pose_config=None, pose_ckpt=None, device="cpu"):
|
35 |
+
from .wholebody import Wholebody
|
36 |
+
|
37 |
+
self.pose_estimation = Wholebody(det_config, det_ckpt, pose_config, pose_ckpt, device)
|
38 |
+
|
39 |
+
def to(self, device):
|
40 |
+
self.pose_estimation.to(device)
|
41 |
+
return self
|
42 |
+
|
43 |
+
def __call__(self, input_image, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs):
|
44 |
+
|
45 |
+
input_image = cv2.cvtColor(np.array(input_image, dtype=np.uint8), cv2.COLOR_RGB2BGR)
|
46 |
+
|
47 |
+
input_image = HWC3(input_image)
|
48 |
+
input_image = resize_image(input_image, detect_resolution)
|
49 |
+
H, W, C = input_image.shape
|
50 |
+
|
51 |
+
with torch.no_grad():
|
52 |
+
candidate, subset = self.pose_estimation(input_image)
|
53 |
+
nums, keys, locs = candidate.shape
|
54 |
+
candidate[..., 0] /= float(W)
|
55 |
+
candidate[..., 1] /= float(H)
|
56 |
+
body = candidate[:,:18].copy()
|
57 |
+
body = body.reshape(nums*18, locs)
|
58 |
+
score = subset[:,:18]
|
59 |
+
|
60 |
+
for i in range(len(score)):
|
61 |
+
for j in range(len(score[i])):
|
62 |
+
if score[i][j] > 0.3:
|
63 |
+
score[i][j] = int(18*i+j)
|
64 |
+
else:
|
65 |
+
score[i][j] = -1
|
66 |
+
|
67 |
+
un_visible = subset<0.3
|
68 |
+
candidate[un_visible] = -1
|
69 |
+
|
70 |
+
foot = candidate[:,18:24]
|
71 |
+
|
72 |
+
faces = candidate[:,24:92]
|
73 |
+
|
74 |
+
hands = candidate[:,92:113]
|
75 |
+
hands = np.vstack([hands, candidate[:,113:]])
|
76 |
+
|
77 |
+
bodies = dict(candidate=body, subset=score)
|
78 |
+
pose = dict(bodies=bodies, hands=hands, faces=faces)
|
79 |
+
|
80 |
+
detected_map = draw_pose(pose, H, W)
|
81 |
+
detected_map = HWC3(detected_map)
|
82 |
+
|
83 |
+
img = resize_image(input_image, image_resolution)
|
84 |
+
H, W, C = img.shape
|
85 |
+
|
86 |
+
detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
|
87 |
+
|
88 |
+
if output_type == "pil":
|
89 |
+
detected_map = Image.fromarray(detected_map)
|
90 |
+
|
91 |
+
return detected_map
|
controlnet_aux/dwpose/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (2.66 kB). View file
|
|
controlnet_aux/dwpose/__pycache__/util.cpython-310.pyc
ADDED
Binary file (7.82 kB). View file
|
|
controlnet_aux/dwpose/__pycache__/wholebody.cpython-310.pyc
ADDED
Binary file (3.63 kB). View file
|
|
controlnet_aux/dwpose/util.py
ADDED
@@ -0,0 +1,303 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import math
|
2 |
+
import numpy as np
|
3 |
+
import cv2
|
4 |
+
|
5 |
+
|
6 |
+
eps = 0.01
|
7 |
+
|
8 |
+
|
9 |
+
def smart_resize(x, s):
|
10 |
+
Ht, Wt = s
|
11 |
+
if x.ndim == 2:
|
12 |
+
Ho, Wo = x.shape
|
13 |
+
Co = 1
|
14 |
+
else:
|
15 |
+
Ho, Wo, Co = x.shape
|
16 |
+
if Co == 3 or Co == 1:
|
17 |
+
k = float(Ht + Wt) / float(Ho + Wo)
|
18 |
+
return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
|
19 |
+
else:
|
20 |
+
return np.stack([smart_resize(x[:, :, i], s) for i in range(Co)], axis=2)
|
21 |
+
|
22 |
+
|
23 |
+
def smart_resize_k(x, fx, fy):
|
24 |
+
if x.ndim == 2:
|
25 |
+
Ho, Wo = x.shape
|
26 |
+
Co = 1
|
27 |
+
else:
|
28 |
+
Ho, Wo, Co = x.shape
|
29 |
+
Ht, Wt = Ho * fy, Wo * fx
|
30 |
+
if Co == 3 or Co == 1:
|
31 |
+
k = float(Ht + Wt) / float(Ho + Wo)
|
32 |
+
return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
|
33 |
+
else:
|
34 |
+
return np.stack([smart_resize_k(x[:, :, i], fx, fy) for i in range(Co)], axis=2)
|
35 |
+
|
36 |
+
|
37 |
+
def padRightDownCorner(img, stride, padValue):
|
38 |
+
h = img.shape[0]
|
39 |
+
w = img.shape[1]
|
40 |
+
|
41 |
+
pad = 4 * [None]
|
42 |
+
pad[0] = 0 # up
|
43 |
+
pad[1] = 0 # left
|
44 |
+
pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down
|
45 |
+
pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right
|
46 |
+
|
47 |
+
img_padded = img
|
48 |
+
pad_up = np.tile(img_padded[0:1, :, :]*0 + padValue, (pad[0], 1, 1))
|
49 |
+
img_padded = np.concatenate((pad_up, img_padded), axis=0)
|
50 |
+
pad_left = np.tile(img_padded[:, 0:1, :]*0 + padValue, (1, pad[1], 1))
|
51 |
+
img_padded = np.concatenate((pad_left, img_padded), axis=1)
|
52 |
+
pad_down = np.tile(img_padded[-2:-1, :, :]*0 + padValue, (pad[2], 1, 1))
|
53 |
+
img_padded = np.concatenate((img_padded, pad_down), axis=0)
|
54 |
+
pad_right = np.tile(img_padded[:, -2:-1, :]*0 + padValue, (1, pad[3], 1))
|
55 |
+
img_padded = np.concatenate((img_padded, pad_right), axis=1)
|
56 |
+
|
57 |
+
return img_padded, pad
|
58 |
+
|
59 |
+
|
60 |
+
def transfer(model, model_weights):
|
61 |
+
transfered_model_weights = {}
|
62 |
+
for weights_name in model.state_dict().keys():
|
63 |
+
transfered_model_weights[weights_name] = model_weights['.'.join(weights_name.split('.')[1:])]
|
64 |
+
return transfered_model_weights
|
65 |
+
|
66 |
+
|
67 |
+
def draw_bodypose(canvas, candidate, subset):
|
68 |
+
H, W, C = canvas.shape
|
69 |
+
candidate = np.array(candidate)
|
70 |
+
subset = np.array(subset)
|
71 |
+
|
72 |
+
stickwidth = 4
|
73 |
+
|
74 |
+
limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
|
75 |
+
[10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
|
76 |
+
[1, 16], [16, 18], [3, 17], [6, 18]]
|
77 |
+
|
78 |
+
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
|
79 |
+
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
|
80 |
+
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
|
81 |
+
|
82 |
+
for i in range(17):
|
83 |
+
for n in range(len(subset)):
|
84 |
+
index = subset[n][np.array(limbSeq[i]) - 1]
|
85 |
+
if -1 in index:
|
86 |
+
continue
|
87 |
+
Y = candidate[index.astype(int), 0] * float(W)
|
88 |
+
X = candidate[index.astype(int), 1] * float(H)
|
89 |
+
mX = np.mean(X)
|
90 |
+
mY = np.mean(Y)
|
91 |
+
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
|
92 |
+
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
|
93 |
+
polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
|
94 |
+
cv2.fillConvexPoly(canvas, polygon, colors[i])
|
95 |
+
|
96 |
+
canvas = (canvas * 0.6).astype(np.uint8)
|
97 |
+
|
98 |
+
for i in range(18):
|
99 |
+
for n in range(len(subset)):
|
100 |
+
index = int(subset[n][i])
|
101 |
+
if index == -1:
|
102 |
+
continue
|
103 |
+
x, y = candidate[index][0:2]
|
104 |
+
x = int(x * W)
|
105 |
+
y = int(y * H)
|
106 |
+
cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
|
107 |
+
|
108 |
+
return canvas
|
109 |
+
|
110 |
+
|
111 |
+
def draw_handpose(canvas, all_hand_peaks):
|
112 |
+
import matplotlib
|
113 |
+
|
114 |
+
H, W, C = canvas.shape
|
115 |
+
|
116 |
+
edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
|
117 |
+
[10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
|
118 |
+
|
119 |
+
# (person_number*2, 21, 2)
|
120 |
+
for i in range(len(all_hand_peaks)):
|
121 |
+
peaks = all_hand_peaks[i]
|
122 |
+
peaks = np.array(peaks)
|
123 |
+
|
124 |
+
for ie, e in enumerate(edges):
|
125 |
+
|
126 |
+
x1, y1 = peaks[e[0]]
|
127 |
+
x2, y2 = peaks[e[1]]
|
128 |
+
|
129 |
+
x1 = int(x1 * W)
|
130 |
+
y1 = int(y1 * H)
|
131 |
+
x2 = int(x2 * W)
|
132 |
+
y2 = int(y2 * H)
|
133 |
+
if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
|
134 |
+
cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255, thickness=2)
|
135 |
+
|
136 |
+
for _, keyponit in enumerate(peaks):
|
137 |
+
x, y = keyponit
|
138 |
+
|
139 |
+
x = int(x * W)
|
140 |
+
y = int(y * H)
|
141 |
+
if x > eps and y > eps:
|
142 |
+
cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
|
143 |
+
return canvas
|
144 |
+
|
145 |
+
|
146 |
+
def draw_facepose(canvas, all_lmks):
|
147 |
+
H, W, C = canvas.shape
|
148 |
+
for lmks in all_lmks:
|
149 |
+
lmks = np.array(lmks)
|
150 |
+
for lmk in lmks:
|
151 |
+
x, y = lmk
|
152 |
+
x = int(x * W)
|
153 |
+
y = int(y * H)
|
154 |
+
if x > eps and y > eps:
|
155 |
+
cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1)
|
156 |
+
return canvas
|
157 |
+
|
158 |
+
|
159 |
+
# detect hand according to body pose keypoints
|
160 |
+
# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp
|
161 |
+
def handDetect(candidate, subset, oriImg):
|
162 |
+
# right hand: wrist 4, elbow 3, shoulder 2
|
163 |
+
# left hand: wrist 7, elbow 6, shoulder 5
|
164 |
+
ratioWristElbow = 0.33
|
165 |
+
detect_result = []
|
166 |
+
image_height, image_width = oriImg.shape[0:2]
|
167 |
+
for person in subset.astype(int):
|
168 |
+
# if any of three not detected
|
169 |
+
has_left = np.sum(person[[5, 6, 7]] == -1) == 0
|
170 |
+
has_right = np.sum(person[[2, 3, 4]] == -1) == 0
|
171 |
+
if not (has_left or has_right):
|
172 |
+
continue
|
173 |
+
hands = []
|
174 |
+
#left hand
|
175 |
+
if has_left:
|
176 |
+
left_shoulder_index, left_elbow_index, left_wrist_index = person[[5, 6, 7]]
|
177 |
+
x1, y1 = candidate[left_shoulder_index][:2]
|
178 |
+
x2, y2 = candidate[left_elbow_index][:2]
|
179 |
+
x3, y3 = candidate[left_wrist_index][:2]
|
180 |
+
hands.append([x1, y1, x2, y2, x3, y3, True])
|
181 |
+
# right hand
|
182 |
+
if has_right:
|
183 |
+
right_shoulder_index, right_elbow_index, right_wrist_index = person[[2, 3, 4]]
|
184 |
+
x1, y1 = candidate[right_shoulder_index][:2]
|
185 |
+
x2, y2 = candidate[right_elbow_index][:2]
|
186 |
+
x3, y3 = candidate[right_wrist_index][:2]
|
187 |
+
hands.append([x1, y1, x2, y2, x3, y3, False])
|
188 |
+
|
189 |
+
for x1, y1, x2, y2, x3, y3, is_left in hands:
|
190 |
+
# pos_hand = pos_wrist + ratio * (pos_wrist - pos_elbox) = (1 + ratio) * pos_wrist - ratio * pos_elbox
|
191 |
+
# handRectangle.x = posePtr[wrist*3] + ratioWristElbow * (posePtr[wrist*3] - posePtr[elbow*3]);
|
192 |
+
# handRectangle.y = posePtr[wrist*3+1] + ratioWristElbow * (posePtr[wrist*3+1] - posePtr[elbow*3+1]);
|
193 |
+
# const auto distanceWristElbow = getDistance(poseKeypoints, person, wrist, elbow);
|
194 |
+
# const auto distanceElbowShoulder = getDistance(poseKeypoints, person, elbow, shoulder);
|
195 |
+
# handRectangle.width = 1.5f * fastMax(distanceWristElbow, 0.9f * distanceElbowShoulder);
|
196 |
+
x = x3 + ratioWristElbow * (x3 - x2)
|
197 |
+
y = y3 + ratioWristElbow * (y3 - y2)
|
198 |
+
distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2)
|
199 |
+
distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
|
200 |
+
width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder)
|
201 |
+
# x-y refers to the center --> offset to topLeft point
|
202 |
+
# handRectangle.x -= handRectangle.width / 2.f;
|
203 |
+
# handRectangle.y -= handRectangle.height / 2.f;
|
204 |
+
x -= width / 2
|
205 |
+
y -= width / 2 # width = height
|
206 |
+
# overflow the image
|
207 |
+
if x < 0: x = 0
|
208 |
+
if y < 0: y = 0
|
209 |
+
width1 = width
|
210 |
+
width2 = width
|
211 |
+
if x + width > image_width: width1 = image_width - x
|
212 |
+
if y + width > image_height: width2 = image_height - y
|
213 |
+
width = min(width1, width2)
|
214 |
+
# the max hand box value is 20 pixels
|
215 |
+
if width >= 20:
|
216 |
+
detect_result.append([int(x), int(y), int(width), is_left])
|
217 |
+
|
218 |
+
'''
|
219 |
+
return value: [[x, y, w, True if left hand else False]].
|
220 |
+
width=height since the network require squared input.
|
221 |
+
x, y is the coordinate of top left
|
222 |
+
'''
|
223 |
+
return detect_result
|
224 |
+
|
225 |
+
|
226 |
+
# Written by Lvmin
|
227 |
+
def faceDetect(candidate, subset, oriImg):
|
228 |
+
# left right eye ear 14 15 16 17
|
229 |
+
detect_result = []
|
230 |
+
image_height, image_width = oriImg.shape[0:2]
|
231 |
+
for person in subset.astype(int):
|
232 |
+
has_head = person[0] > -1
|
233 |
+
if not has_head:
|
234 |
+
continue
|
235 |
+
|
236 |
+
has_left_eye = person[14] > -1
|
237 |
+
has_right_eye = person[15] > -1
|
238 |
+
has_left_ear = person[16] > -1
|
239 |
+
has_right_ear = person[17] > -1
|
240 |
+
|
241 |
+
if not (has_left_eye or has_right_eye or has_left_ear or has_right_ear):
|
242 |
+
continue
|
243 |
+
|
244 |
+
head, left_eye, right_eye, left_ear, right_ear = person[[0, 14, 15, 16, 17]]
|
245 |
+
|
246 |
+
width = 0.0
|
247 |
+
x0, y0 = candidate[head][:2]
|
248 |
+
|
249 |
+
if has_left_eye:
|
250 |
+
x1, y1 = candidate[left_eye][:2]
|
251 |
+
d = max(abs(x0 - x1), abs(y0 - y1))
|
252 |
+
width = max(width, d * 3.0)
|
253 |
+
|
254 |
+
if has_right_eye:
|
255 |
+
x1, y1 = candidate[right_eye][:2]
|
256 |
+
d = max(abs(x0 - x1), abs(y0 - y1))
|
257 |
+
width = max(width, d * 3.0)
|
258 |
+
|
259 |
+
if has_left_ear:
|
260 |
+
x1, y1 = candidate[left_ear][:2]
|
261 |
+
d = max(abs(x0 - x1), abs(y0 - y1))
|
262 |
+
width = max(width, d * 1.5)
|
263 |
+
|
264 |
+
if has_right_ear:
|
265 |
+
x1, y1 = candidate[right_ear][:2]
|
266 |
+
d = max(abs(x0 - x1), abs(y0 - y1))
|
267 |
+
width = max(width, d * 1.5)
|
268 |
+
|
269 |
+
x, y = x0, y0
|
270 |
+
|
271 |
+
x -= width
|
272 |
+
y -= width
|
273 |
+
|
274 |
+
if x < 0:
|
275 |
+
x = 0
|
276 |
+
|
277 |
+
if y < 0:
|
278 |
+
y = 0
|
279 |
+
|
280 |
+
width1 = width * 2
|
281 |
+
width2 = width * 2
|
282 |
+
|
283 |
+
if x + width > image_width:
|
284 |
+
width1 = image_width - x
|
285 |
+
|
286 |
+
if y + width > image_height:
|
287 |
+
width2 = image_height - y
|
288 |
+
|
289 |
+
width = min(width1, width2)
|
290 |
+
|
291 |
+
if width >= 20:
|
292 |
+
detect_result.append([int(x), int(y), int(width)])
|
293 |
+
|
294 |
+
return detect_result
|
295 |
+
|
296 |
+
|
297 |
+
# get max index of 2d array
|
298 |
+
def npmax(array):
|
299 |
+
arrayindex = array.argmax(1)
|
300 |
+
arrayvalue = array.max(1)
|
301 |
+
i = arrayvalue.argmax()
|
302 |
+
j = arrayindex[i]
|
303 |
+
return i, j
|
controlnet_aux/dwpose/wholebody.py
ADDED
@@ -0,0 +1,121 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) OpenMMLab. All rights reserved.
|
2 |
+
import os
|
3 |
+
import numpy as np
|
4 |
+
import warnings
|
5 |
+
|
6 |
+
try:
|
7 |
+
import mmcv
|
8 |
+
except ImportError:
|
9 |
+
warnings.warn(
|
10 |
+
"The module 'mmcv' is not installed. The package will have limited functionality. Please install it using the command: mim install 'mmcv>=2.0.1'"
|
11 |
+
)
|
12 |
+
|
13 |
+
try:
|
14 |
+
from mmpose.apis import inference_topdown
|
15 |
+
from mmpose.apis import init_model as init_pose_estimator
|
16 |
+
from mmpose.evaluation.functional import nms
|
17 |
+
from mmpose.utils import adapt_mmdet_pipeline
|
18 |
+
from mmpose.structures import merge_data_samples
|
19 |
+
except ImportError:
|
20 |
+
warnings.warn(
|
21 |
+
"The module 'mmpose' is not installed. The package will have limited functionality. Please install it using the command: mim install 'mmpose>=1.1.0'"
|
22 |
+
)
|
23 |
+
|
24 |
+
try:
|
25 |
+
from mmdet.apis import inference_detector, init_detector
|
26 |
+
except ImportError:
|
27 |
+
warnings.warn(
|
28 |
+
"The module 'mmdet' is not installed. The package will have limited functionality. Please install it using the command: mim install 'mmdet>=3.1.0'"
|
29 |
+
)
|
30 |
+
|
31 |
+
|
32 |
+
class Wholebody:
|
33 |
+
def __init__(self,
|
34 |
+
det_config=None, det_ckpt=None,
|
35 |
+
pose_config=None, pose_ckpt=None,
|
36 |
+
device="cpu"):
|
37 |
+
|
38 |
+
if det_config is None:
|
39 |
+
det_config = os.path.join(os.path.dirname(__file__), "yolox_config/yolox_l_8xb8-300e_coco.py")
|
40 |
+
|
41 |
+
if pose_config is None:
|
42 |
+
pose_config = os.path.join(os.path.dirname(__file__), "dwpose_config/dwpose-l_384x288.py")
|
43 |
+
|
44 |
+
if det_ckpt is None:
|
45 |
+
det_ckpt = 'https://download.openmmlab.com/mmdetection/v2.0/yolox/yolox_l_8x8_300e_coco/yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth'
|
46 |
+
|
47 |
+
if pose_ckpt is None:
|
48 |
+
pose_ckpt = "https://huggingface.co/wanghaofan/dw-ll_ucoco_384/resolve/main/dw-ll_ucoco_384.pth"
|
49 |
+
|
50 |
+
# build detector
|
51 |
+
self.detector = init_detector(det_config, det_ckpt, device=device)
|
52 |
+
self.detector.cfg = adapt_mmdet_pipeline(self.detector.cfg)
|
53 |
+
|
54 |
+
# build pose estimator
|
55 |
+
self.pose_estimator = init_pose_estimator(
|
56 |
+
pose_config,
|
57 |
+
pose_ckpt,
|
58 |
+
device=device)
|
59 |
+
|
60 |
+
def to(self, device):
|
61 |
+
self.detector.to(device)
|
62 |
+
self.pose_estimator.to(device)
|
63 |
+
return self
|
64 |
+
|
65 |
+
def __call__(self, oriImg):
|
66 |
+
# predict bbox
|
67 |
+
det_result = inference_detector(self.detector, oriImg)
|
68 |
+
pred_instance = det_result.pred_instances.cpu().numpy()
|
69 |
+
bboxes = np.concatenate(
|
70 |
+
(pred_instance.bboxes, pred_instance.scores[:, None]), axis=1)
|
71 |
+
bboxes = bboxes[np.logical_and(pred_instance.labels == 0,
|
72 |
+
pred_instance.scores > 0.5)]
|
73 |
+
|
74 |
+
# set NMS threshold
|
75 |
+
bboxes = bboxes[nms(bboxes, 0.7), :4]
|
76 |
+
|
77 |
+
# predict keypoints
|
78 |
+
if len(bboxes) == 0:
|
79 |
+
pose_results = inference_topdown(self.pose_estimator, oriImg)
|
80 |
+
else:
|
81 |
+
pose_results = inference_topdown(self.pose_estimator, oriImg, bboxes)
|
82 |
+
preds = merge_data_samples(pose_results)
|
83 |
+
preds = preds.pred_instances
|
84 |
+
|
85 |
+
# preds = pose_results[0].pred_instances
|
86 |
+
keypoints = preds.get('transformed_keypoints',
|
87 |
+
preds.keypoints)
|
88 |
+
if 'keypoint_scores' in preds:
|
89 |
+
scores = preds.keypoint_scores
|
90 |
+
else:
|
91 |
+
scores = np.ones(keypoints.shape[:-1])
|
92 |
+
|
93 |
+
if 'keypoints_visible' in preds:
|
94 |
+
visible = preds.keypoints_visible
|
95 |
+
else:
|
96 |
+
visible = np.ones(keypoints.shape[:-1])
|
97 |
+
keypoints_info = np.concatenate(
|
98 |
+
(keypoints, scores[..., None], visible[..., None]),
|
99 |
+
axis=-1)
|
100 |
+
# compute neck joint
|
101 |
+
neck = np.mean(keypoints_info[:, [5, 6]], axis=1)
|
102 |
+
# neck score when visualizing pred
|
103 |
+
neck[:, 2:4] = np.logical_and(
|
104 |
+
keypoints_info[:, 5, 2:4] > 0.3,
|
105 |
+
keypoints_info[:, 6, 2:4] > 0.3).astype(int)
|
106 |
+
new_keypoints_info = np.insert(
|
107 |
+
keypoints_info, 17, neck, axis=1)
|
108 |
+
mmpose_idx = [
|
109 |
+
17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3
|
110 |
+
]
|
111 |
+
openpose_idx = [
|
112 |
+
1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17
|
113 |
+
]
|
114 |
+
new_keypoints_info[:, openpose_idx] = \
|
115 |
+
new_keypoints_info[:, mmpose_idx]
|
116 |
+
keypoints_info = new_keypoints_info
|
117 |
+
|
118 |
+
keypoints, scores, visible = keypoints_info[
|
119 |
+
..., :2], keypoints_info[..., 2], keypoints_info[..., 3]
|
120 |
+
|
121 |
+
return keypoints, scores
|
controlnet_aux/hed/__init__.py
ADDED
@@ -0,0 +1,129 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# This is an improved version and model of HED edge detection with Apache License, Version 2.0.
|
2 |
+
# Please use this implementation in your products
|
3 |
+
# This implementation may produce slightly different results from Saining Xie's official implementations,
|
4 |
+
# but it generates smoother edges and is more suitable for ControlNet as well as other image-to-image translations.
|
5 |
+
# Different from official models and other implementations, this is an RGB-input model (rather than BGR)
|
6 |
+
# and in this way it works better for gradio's RGB protocol
|
7 |
+
|
8 |
+
import os
|
9 |
+
import warnings
|
10 |
+
|
11 |
+
import cv2
|
12 |
+
import numpy as np
|
13 |
+
import torch
|
14 |
+
from einops import rearrange
|
15 |
+
from huggingface_hub import hf_hub_download
|
16 |
+
from PIL import Image
|
17 |
+
|
18 |
+
from ..util import HWC3, nms, resize_image, safe_step
|
19 |
+
|
20 |
+
|
21 |
+
class DoubleConvBlock(torch.nn.Module):
|
22 |
+
def __init__(self, input_channel, output_channel, layer_number):
|
23 |
+
super().__init__()
|
24 |
+
self.convs = torch.nn.Sequential()
|
25 |
+
self.convs.append(torch.nn.Conv2d(in_channels=input_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1))
|
26 |
+
for i in range(1, layer_number):
|
27 |
+
self.convs.append(torch.nn.Conv2d(in_channels=output_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1))
|
28 |
+
self.projection = torch.nn.Conv2d(in_channels=output_channel, out_channels=1, kernel_size=(1, 1), stride=(1, 1), padding=0)
|
29 |
+
|
30 |
+
def __call__(self, x, down_sampling=False):
|
31 |
+
h = x
|
32 |
+
if down_sampling:
|
33 |
+
h = torch.nn.functional.max_pool2d(h, kernel_size=(2, 2), stride=(2, 2))
|
34 |
+
for conv in self.convs:
|
35 |
+
h = conv(h)
|
36 |
+
h = torch.nn.functional.relu(h)
|
37 |
+
return h, self.projection(h)
|
38 |
+
|
39 |
+
|
40 |
+
class ControlNetHED_Apache2(torch.nn.Module):
|
41 |
+
def __init__(self):
|
42 |
+
super().__init__()
|
43 |
+
self.norm = torch.nn.Parameter(torch.zeros(size=(1, 3, 1, 1)))
|
44 |
+
self.block1 = DoubleConvBlock(input_channel=3, output_channel=64, layer_number=2)
|
45 |
+
self.block2 = DoubleConvBlock(input_channel=64, output_channel=128, layer_number=2)
|
46 |
+
self.block3 = DoubleConvBlock(input_channel=128, output_channel=256, layer_number=3)
|
47 |
+
self.block4 = DoubleConvBlock(input_channel=256, output_channel=512, layer_number=3)
|
48 |
+
self.block5 = DoubleConvBlock(input_channel=512, output_channel=512, layer_number=3)
|
49 |
+
|
50 |
+
def __call__(self, x):
|
51 |
+
h = x - self.norm
|
52 |
+
h, projection1 = self.block1(h)
|
53 |
+
h, projection2 = self.block2(h, down_sampling=True)
|
54 |
+
h, projection3 = self.block3(h, down_sampling=True)
|
55 |
+
h, projection4 = self.block4(h, down_sampling=True)
|
56 |
+
h, projection5 = self.block5(h, down_sampling=True)
|
57 |
+
return projection1, projection2, projection3, projection4, projection5
|
58 |
+
|
59 |
+
class HEDdetector:
|
60 |
+
def __init__(self, netNetwork):
|
61 |
+
self.netNetwork = netNetwork
|
62 |
+
|
63 |
+
@classmethod
|
64 |
+
def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None):
|
65 |
+
filename = filename or "ControlNetHED.pth"
|
66 |
+
|
67 |
+
if os.path.isdir(pretrained_model_or_path):
|
68 |
+
model_path = os.path.join(pretrained_model_or_path, filename)
|
69 |
+
else:
|
70 |
+
model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
|
71 |
+
|
72 |
+
netNetwork = ControlNetHED_Apache2()
|
73 |
+
netNetwork.load_state_dict(torch.load(model_path, map_location='cpu'))
|
74 |
+
netNetwork.float().eval()
|
75 |
+
|
76 |
+
return cls(netNetwork)
|
77 |
+
|
78 |
+
def to(self, device):
|
79 |
+
self.netNetwork.to(device)
|
80 |
+
return self
|
81 |
+
|
82 |
+
def __call__(self, input_image, detect_resolution=512, image_resolution=512, safe=False, output_type="pil", scribble=False, **kwargs):
|
83 |
+
if "return_pil" in kwargs:
|
84 |
+
warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
|
85 |
+
output_type = "pil" if kwargs["return_pil"] else "np"
|
86 |
+
if type(output_type) is bool:
|
87 |
+
warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
|
88 |
+
if output_type:
|
89 |
+
output_type = "pil"
|
90 |
+
|
91 |
+
device = next(iter(self.netNetwork.parameters())).device
|
92 |
+
if not isinstance(input_image, np.ndarray):
|
93 |
+
input_image = np.array(input_image, dtype=np.uint8)
|
94 |
+
|
95 |
+
input_image = HWC3(input_image)
|
96 |
+
input_image = resize_image(input_image, detect_resolution)
|
97 |
+
|
98 |
+
assert input_image.ndim == 3
|
99 |
+
H, W, C = input_image.shape
|
100 |
+
with torch.no_grad():
|
101 |
+
image_hed = torch.from_numpy(input_image.copy()).float().to(device)
|
102 |
+
image_hed = rearrange(image_hed, 'h w c -> 1 c h w')
|
103 |
+
edges = self.netNetwork(image_hed)
|
104 |
+
edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges]
|
105 |
+
edges = [cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges]
|
106 |
+
edges = np.stack(edges, axis=2)
|
107 |
+
edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
|
108 |
+
if safe:
|
109 |
+
edge = safe_step(edge)
|
110 |
+
edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
|
111 |
+
|
112 |
+
detected_map = edge
|
113 |
+
detected_map = HWC3(detected_map)
|
114 |
+
|
115 |
+
img = resize_image(input_image, image_resolution)
|
116 |
+
H, W, C = img.shape
|
117 |
+
|
118 |
+
detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
|
119 |
+
|
120 |
+
if scribble:
|
121 |
+
detected_map = nms(detected_map, 127, 3.0)
|
122 |
+
detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
|
123 |
+
detected_map[detected_map > 4] = 255
|
124 |
+
detected_map[detected_map < 255] = 0
|
125 |
+
|
126 |
+
if output_type == "pil":
|
127 |
+
detected_map = Image.fromarray(detected_map)
|
128 |
+
|
129 |
+
return detected_map
|
controlnet_aux/hed/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (5.06 kB). View file
|
|
controlnet_aux/leres/__init__.py
ADDED
@@ -0,0 +1,118 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
|
3 |
+
import cv2
|
4 |
+
import numpy as np
|
5 |
+
import torch
|
6 |
+
from huggingface_hub import hf_hub_download
|
7 |
+
from PIL import Image
|
8 |
+
|
9 |
+
from ..util import HWC3, resize_image
|
10 |
+
from .leres.depthmap import estimateboost, estimateleres
|
11 |
+
from .leres.multi_depth_model_woauxi import RelDepthModel
|
12 |
+
from .leres.net_tools import strip_prefix_if_present
|
13 |
+
from .pix2pix.models.pix2pix4depth_model import Pix2Pix4DepthModel
|
14 |
+
from .pix2pix.options.test_options import TestOptions
|
15 |
+
|
16 |
+
|
17 |
+
class LeresDetector:
|
18 |
+
def __init__(self, model, pix2pixmodel):
|
19 |
+
self.model = model
|
20 |
+
self.pix2pixmodel = pix2pixmodel
|
21 |
+
|
22 |
+
@classmethod
|
23 |
+
def from_pretrained(cls, pretrained_model_or_path, filename=None, pix2pix_filename=None, cache_dir=None):
|
24 |
+
filename = filename or "res101.pth"
|
25 |
+
pix2pix_filename = pix2pix_filename or "latest_net_G.pth"
|
26 |
+
|
27 |
+
if os.path.isdir(pretrained_model_or_path):
|
28 |
+
model_path = os.path.join(pretrained_model_or_path, filename)
|
29 |
+
else:
|
30 |
+
model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
|
31 |
+
|
32 |
+
checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
|
33 |
+
|
34 |
+
model = RelDepthModel(backbone='resnext101')
|
35 |
+
model.load_state_dict(strip_prefix_if_present(checkpoint['depth_model'], "module."), strict=True)
|
36 |
+
del checkpoint
|
37 |
+
|
38 |
+
if os.path.isdir(pretrained_model_or_path):
|
39 |
+
model_path = os.path.join(pretrained_model_or_path, pix2pix_filename)
|
40 |
+
else:
|
41 |
+
model_path = hf_hub_download(pretrained_model_or_path, pix2pix_filename, cache_dir=cache_dir)
|
42 |
+
|
43 |
+
opt = TestOptions().parse()
|
44 |
+
if not torch.cuda.is_available():
|
45 |
+
opt.gpu_ids = [] # cpu mode
|
46 |
+
pix2pixmodel = Pix2Pix4DepthModel(opt)
|
47 |
+
pix2pixmodel.save_dir = os.path.dirname(model_path)
|
48 |
+
pix2pixmodel.load_networks('latest')
|
49 |
+
pix2pixmodel.eval()
|
50 |
+
|
51 |
+
return cls(model, pix2pixmodel)
|
52 |
+
|
53 |
+
def to(self, device):
|
54 |
+
self.model.to(device)
|
55 |
+
# TODO - refactor pix2pix implementation to support device migration
|
56 |
+
# self.pix2pixmodel.to(device)
|
57 |
+
return self
|
58 |
+
|
59 |
+
def __call__(self, input_image, thr_a=0, thr_b=0, boost=False, detect_resolution=512, image_resolution=512, output_type="pil"):
|
60 |
+
device = next(iter(self.model.parameters())).device
|
61 |
+
if not isinstance(input_image, np.ndarray):
|
62 |
+
input_image = np.array(input_image, dtype=np.uint8)
|
63 |
+
|
64 |
+
input_image = HWC3(input_image)
|
65 |
+
input_image = resize_image(input_image, detect_resolution)
|
66 |
+
|
67 |
+
assert input_image.ndim == 3
|
68 |
+
height, width, dim = input_image.shape
|
69 |
+
|
70 |
+
with torch.no_grad():
|
71 |
+
|
72 |
+
if boost:
|
73 |
+
depth = estimateboost(input_image, self.model, 0, self.pix2pixmodel, max(width, height))
|
74 |
+
else:
|
75 |
+
depth = estimateleres(input_image, self.model, width, height)
|
76 |
+
|
77 |
+
numbytes=2
|
78 |
+
depth_min = depth.min()
|
79 |
+
depth_max = depth.max()
|
80 |
+
max_val = (2**(8*numbytes))-1
|
81 |
+
|
82 |
+
# check output before normalizing and mapping to 16 bit
|
83 |
+
if depth_max - depth_min > np.finfo("float").eps:
|
84 |
+
out = max_val * (depth - depth_min) / (depth_max - depth_min)
|
85 |
+
else:
|
86 |
+
out = np.zeros(depth.shape)
|
87 |
+
|
88 |
+
# single channel, 16 bit image
|
89 |
+
depth_image = out.astype("uint16")
|
90 |
+
|
91 |
+
# convert to uint8
|
92 |
+
depth_image = cv2.convertScaleAbs(depth_image, alpha=(255.0/65535.0))
|
93 |
+
|
94 |
+
# remove near
|
95 |
+
if thr_a != 0:
|
96 |
+
thr_a = ((thr_a/100)*255)
|
97 |
+
depth_image = cv2.threshold(depth_image, thr_a, 255, cv2.THRESH_TOZERO)[1]
|
98 |
+
|
99 |
+
# invert image
|
100 |
+
depth_image = cv2.bitwise_not(depth_image)
|
101 |
+
|
102 |
+
# remove bg
|
103 |
+
if thr_b != 0:
|
104 |
+
thr_b = ((thr_b/100)*255)
|
105 |
+
depth_image = cv2.threshold(depth_image, thr_b, 255, cv2.THRESH_TOZERO)[1]
|
106 |
+
|
107 |
+
detected_map = depth_image
|
108 |
+
detected_map = HWC3(detected_map)
|
109 |
+
|
110 |
+
img = resize_image(input_image, image_resolution)
|
111 |
+
H, W, C = img.shape
|
112 |
+
|
113 |
+
detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
|
114 |
+
|
115 |
+
if output_type == "pil":
|
116 |
+
detected_map = Image.fromarray(detected_map)
|
117 |
+
|
118 |
+
return detected_map
|
controlnet_aux/leres/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (3.33 kB). View file
|
|
controlnet_aux/leres/leres/Resnet.py
ADDED
@@ -0,0 +1,199 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch.nn as nn
|
2 |
+
import torch.nn as NN
|
3 |
+
|
4 |
+
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
|
5 |
+
'resnet152']
|
6 |
+
|
7 |
+
|
8 |
+
model_urls = {
|
9 |
+
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
|
10 |
+
'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
|
11 |
+
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
|
12 |
+
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
|
13 |
+
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
|
14 |
+
}
|
15 |
+
|
16 |
+
|
17 |
+
def conv3x3(in_planes, out_planes, stride=1):
|
18 |
+
"""3x3 convolution with padding"""
|
19 |
+
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
20 |
+
padding=1, bias=False)
|
21 |
+
|
22 |
+
|
23 |
+
class BasicBlock(nn.Module):
|
24 |
+
expansion = 1
|
25 |
+
|
26 |
+
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
27 |
+
super(BasicBlock, self).__init__()
|
28 |
+
self.conv1 = conv3x3(inplanes, planes, stride)
|
29 |
+
self.bn1 = NN.BatchNorm2d(planes) #NN.BatchNorm2d
|
30 |
+
self.relu = nn.ReLU(inplace=True)
|
31 |
+
self.conv2 = conv3x3(planes, planes)
|
32 |
+
self.bn2 = NN.BatchNorm2d(planes) #NN.BatchNorm2d
|
33 |
+
self.downsample = downsample
|
34 |
+
self.stride = stride
|
35 |
+
|
36 |
+
def forward(self, x):
|
37 |
+
residual = x
|
38 |
+
|
39 |
+
out = self.conv1(x)
|
40 |
+
out = self.bn1(out)
|
41 |
+
out = self.relu(out)
|
42 |
+
|
43 |
+
out = self.conv2(out)
|
44 |
+
out = self.bn2(out)
|
45 |
+
|
46 |
+
if self.downsample is not None:
|
47 |
+
residual = self.downsample(x)
|
48 |
+
|
49 |
+
out += residual
|
50 |
+
out = self.relu(out)
|
51 |
+
|
52 |
+
return out
|
53 |
+
|
54 |
+
|
55 |
+
class Bottleneck(nn.Module):
|
56 |
+
expansion = 4
|
57 |
+
|
58 |
+
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
59 |
+
super(Bottleneck, self).__init__()
|
60 |
+
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
|
61 |
+
self.bn1 = NN.BatchNorm2d(planes) #NN.BatchNorm2d
|
62 |
+
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
|
63 |
+
padding=1, bias=False)
|
64 |
+
self.bn2 = NN.BatchNorm2d(planes) #NN.BatchNorm2d
|
65 |
+
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
|
66 |
+
self.bn3 = NN.BatchNorm2d(planes * self.expansion) #NN.BatchNorm2d
|
67 |
+
self.relu = nn.ReLU(inplace=True)
|
68 |
+
self.downsample = downsample
|
69 |
+
self.stride = stride
|
70 |
+
|
71 |
+
def forward(self, x):
|
72 |
+
residual = x
|
73 |
+
|
74 |
+
out = self.conv1(x)
|
75 |
+
out = self.bn1(out)
|
76 |
+
out = self.relu(out)
|
77 |
+
|
78 |
+
out = self.conv2(out)
|
79 |
+
out = self.bn2(out)
|
80 |
+
out = self.relu(out)
|
81 |
+
|
82 |
+
out = self.conv3(out)
|
83 |
+
out = self.bn3(out)
|
84 |
+
|
85 |
+
if self.downsample is not None:
|
86 |
+
residual = self.downsample(x)
|
87 |
+
|
88 |
+
out += residual
|
89 |
+
out = self.relu(out)
|
90 |
+
|
91 |
+
return out
|
92 |
+
|
93 |
+
|
94 |
+
class ResNet(nn.Module):
|
95 |
+
|
96 |
+
def __init__(self, block, layers, num_classes=1000):
|
97 |
+
self.inplanes = 64
|
98 |
+
super(ResNet, self).__init__()
|
99 |
+
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
|
100 |
+
bias=False)
|
101 |
+
self.bn1 = NN.BatchNorm2d(64) #NN.BatchNorm2d
|
102 |
+
self.relu = nn.ReLU(inplace=True)
|
103 |
+
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
104 |
+
self.layer1 = self._make_layer(block, 64, layers[0])
|
105 |
+
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
106 |
+
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
107 |
+
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
|
108 |
+
#self.avgpool = nn.AvgPool2d(7, stride=1)
|
109 |
+
#self.fc = nn.Linear(512 * block.expansion, num_classes)
|
110 |
+
|
111 |
+
for m in self.modules():
|
112 |
+
if isinstance(m, nn.Conv2d):
|
113 |
+
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
114 |
+
elif isinstance(m, nn.BatchNorm2d):
|
115 |
+
nn.init.constant_(m.weight, 1)
|
116 |
+
nn.init.constant_(m.bias, 0)
|
117 |
+
|
118 |
+
def _make_layer(self, block, planes, blocks, stride=1):
|
119 |
+
downsample = None
|
120 |
+
if stride != 1 or self.inplanes != planes * block.expansion:
|
121 |
+
downsample = nn.Sequential(
|
122 |
+
nn.Conv2d(self.inplanes, planes * block.expansion,
|
123 |
+
kernel_size=1, stride=stride, bias=False),
|
124 |
+
NN.BatchNorm2d(planes * block.expansion), #NN.BatchNorm2d
|
125 |
+
)
|
126 |
+
|
127 |
+
layers = []
|
128 |
+
layers.append(block(self.inplanes, planes, stride, downsample))
|
129 |
+
self.inplanes = planes * block.expansion
|
130 |
+
for i in range(1, blocks):
|
131 |
+
layers.append(block(self.inplanes, planes))
|
132 |
+
|
133 |
+
return nn.Sequential(*layers)
|
134 |
+
|
135 |
+
def forward(self, x):
|
136 |
+
features = []
|
137 |
+
|
138 |
+
x = self.conv1(x)
|
139 |
+
x = self.bn1(x)
|
140 |
+
x = self.relu(x)
|
141 |
+
x = self.maxpool(x)
|
142 |
+
|
143 |
+
x = self.layer1(x)
|
144 |
+
features.append(x)
|
145 |
+
x = self.layer2(x)
|
146 |
+
features.append(x)
|
147 |
+
x = self.layer3(x)
|
148 |
+
features.append(x)
|
149 |
+
x = self.layer4(x)
|
150 |
+
features.append(x)
|
151 |
+
|
152 |
+
return features
|
153 |
+
|
154 |
+
|
155 |
+
def resnet18(pretrained=True, **kwargs):
|
156 |
+
"""Constructs a ResNet-18 model.
|
157 |
+
Args:
|
158 |
+
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
159 |
+
"""
|
160 |
+
model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
|
161 |
+
return model
|
162 |
+
|
163 |
+
|
164 |
+
def resnet34(pretrained=True, **kwargs):
|
165 |
+
"""Constructs a ResNet-34 model.
|
166 |
+
Args:
|
167 |
+
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
168 |
+
"""
|
169 |
+
model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
|
170 |
+
return model
|
171 |
+
|
172 |
+
|
173 |
+
def resnet50(pretrained=True, **kwargs):
|
174 |
+
"""Constructs a ResNet-50 model.
|
175 |
+
Args:
|
176 |
+
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
177 |
+
"""
|
178 |
+
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
|
179 |
+
|
180 |
+
return model
|
181 |
+
|
182 |
+
|
183 |
+
def resnet101(pretrained=True, **kwargs):
|
184 |
+
"""Constructs a ResNet-101 model.
|
185 |
+
Args:
|
186 |
+
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
187 |
+
"""
|
188 |
+
model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
|
189 |
+
|
190 |
+
return model
|
191 |
+
|
192 |
+
|
193 |
+
def resnet152(pretrained=True, **kwargs):
|
194 |
+
"""Constructs a ResNet-152 model.
|
195 |
+
Args:
|
196 |
+
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
197 |
+
"""
|
198 |
+
model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
|
199 |
+
return model
|
controlnet_aux/leres/leres/Resnext_torch.py
ADDED
@@ -0,0 +1,237 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#!/usr/bin/env python
|
2 |
+
# coding: utf-8
|
3 |
+
import torch.nn as nn
|
4 |
+
|
5 |
+
try:
|
6 |
+
from urllib import urlretrieve
|
7 |
+
except ImportError:
|
8 |
+
from urllib.request import urlretrieve
|
9 |
+
|
10 |
+
__all__ = ['resnext101_32x8d']
|
11 |
+
|
12 |
+
|
13 |
+
model_urls = {
|
14 |
+
'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
|
15 |
+
'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
|
16 |
+
}
|
17 |
+
|
18 |
+
|
19 |
+
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
20 |
+
"""3x3 convolution with padding"""
|
21 |
+
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
22 |
+
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
23 |
+
|
24 |
+
|
25 |
+
def conv1x1(in_planes, out_planes, stride=1):
|
26 |
+
"""1x1 convolution"""
|
27 |
+
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
28 |
+
|
29 |
+
|
30 |
+
class BasicBlock(nn.Module):
|
31 |
+
expansion = 1
|
32 |
+
|
33 |
+
def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
|
34 |
+
base_width=64, dilation=1, norm_layer=None):
|
35 |
+
super(BasicBlock, self).__init__()
|
36 |
+
if norm_layer is None:
|
37 |
+
norm_layer = nn.BatchNorm2d
|
38 |
+
if groups != 1 or base_width != 64:
|
39 |
+
raise ValueError('BasicBlock only supports groups=1 and base_width=64')
|
40 |
+
if dilation > 1:
|
41 |
+
raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
|
42 |
+
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
43 |
+
self.conv1 = conv3x3(inplanes, planes, stride)
|
44 |
+
self.bn1 = norm_layer(planes)
|
45 |
+
self.relu = nn.ReLU(inplace=True)
|
46 |
+
self.conv2 = conv3x3(planes, planes)
|
47 |
+
self.bn2 = norm_layer(planes)
|
48 |
+
self.downsample = downsample
|
49 |
+
self.stride = stride
|
50 |
+
|
51 |
+
def forward(self, x):
|
52 |
+
identity = x
|
53 |
+
|
54 |
+
out = self.conv1(x)
|
55 |
+
out = self.bn1(out)
|
56 |
+
out = self.relu(out)
|
57 |
+
|
58 |
+
out = self.conv2(out)
|
59 |
+
out = self.bn2(out)
|
60 |
+
|
61 |
+
if self.downsample is not None:
|
62 |
+
identity = self.downsample(x)
|
63 |
+
|
64 |
+
out += identity
|
65 |
+
out = self.relu(out)
|
66 |
+
|
67 |
+
return out
|
68 |
+
|
69 |
+
|
70 |
+
class Bottleneck(nn.Module):
|
71 |
+
# Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
|
72 |
+
# while original implementation places the stride at the first 1x1 convolution(self.conv1)
|
73 |
+
# according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
|
74 |
+
# This variant is also known as ResNet V1.5 and improves accuracy according to
|
75 |
+
# https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
|
76 |
+
|
77 |
+
expansion = 4
|
78 |
+
|
79 |
+
def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
|
80 |
+
base_width=64, dilation=1, norm_layer=None):
|
81 |
+
super(Bottleneck, self).__init__()
|
82 |
+
if norm_layer is None:
|
83 |
+
norm_layer = nn.BatchNorm2d
|
84 |
+
width = int(planes * (base_width / 64.)) * groups
|
85 |
+
# Both self.conv2 and self.downsample layers downsample the input when stride != 1
|
86 |
+
self.conv1 = conv1x1(inplanes, width)
|
87 |
+
self.bn1 = norm_layer(width)
|
88 |
+
self.conv2 = conv3x3(width, width, stride, groups, dilation)
|
89 |
+
self.bn2 = norm_layer(width)
|
90 |
+
self.conv3 = conv1x1(width, planes * self.expansion)
|
91 |
+
self.bn3 = norm_layer(planes * self.expansion)
|
92 |
+
self.relu = nn.ReLU(inplace=True)
|
93 |
+
self.downsample = downsample
|
94 |
+
self.stride = stride
|
95 |
+
|
96 |
+
def forward(self, x):
|
97 |
+
identity = x
|
98 |
+
|
99 |
+
out = self.conv1(x)
|
100 |
+
out = self.bn1(out)
|
101 |
+
out = self.relu(out)
|
102 |
+
|
103 |
+
out = self.conv2(out)
|
104 |
+
out = self.bn2(out)
|
105 |
+
out = self.relu(out)
|
106 |
+
|
107 |
+
out = self.conv3(out)
|
108 |
+
out = self.bn3(out)
|
109 |
+
|
110 |
+
if self.downsample is not None:
|
111 |
+
identity = self.downsample(x)
|
112 |
+
|
113 |
+
out += identity
|
114 |
+
out = self.relu(out)
|
115 |
+
|
116 |
+
return out
|
117 |
+
|
118 |
+
|
119 |
+
class ResNet(nn.Module):
|
120 |
+
|
121 |
+
def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
|
122 |
+
groups=1, width_per_group=64, replace_stride_with_dilation=None,
|
123 |
+
norm_layer=None):
|
124 |
+
super(ResNet, self).__init__()
|
125 |
+
if norm_layer is None:
|
126 |
+
norm_layer = nn.BatchNorm2d
|
127 |
+
self._norm_layer = norm_layer
|
128 |
+
|
129 |
+
self.inplanes = 64
|
130 |
+
self.dilation = 1
|
131 |
+
if replace_stride_with_dilation is None:
|
132 |
+
# each element in the tuple indicates if we should replace
|
133 |
+
# the 2x2 stride with a dilated convolution instead
|
134 |
+
replace_stride_with_dilation = [False, False, False]
|
135 |
+
if len(replace_stride_with_dilation) != 3:
|
136 |
+
raise ValueError("replace_stride_with_dilation should be None "
|
137 |
+
"or a 3-element tuple, got {}".format(replace_stride_with_dilation))
|
138 |
+
self.groups = groups
|
139 |
+
self.base_width = width_per_group
|
140 |
+
self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
|
141 |
+
bias=False)
|
142 |
+
self.bn1 = norm_layer(self.inplanes)
|
143 |
+
self.relu = nn.ReLU(inplace=True)
|
144 |
+
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
145 |
+
self.layer1 = self._make_layer(block, 64, layers[0])
|
146 |
+
self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
|
147 |
+
dilate=replace_stride_with_dilation[0])
|
148 |
+
self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
|
149 |
+
dilate=replace_stride_with_dilation[1])
|
150 |
+
self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
|
151 |
+
dilate=replace_stride_with_dilation[2])
|
152 |
+
#self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
153 |
+
#self.fc = nn.Linear(512 * block.expansion, num_classes)
|
154 |
+
|
155 |
+
for m in self.modules():
|
156 |
+
if isinstance(m, nn.Conv2d):
|
157 |
+
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
158 |
+
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
159 |
+
nn.init.constant_(m.weight, 1)
|
160 |
+
nn.init.constant_(m.bias, 0)
|
161 |
+
|
162 |
+
# Zero-initialize the last BN in each residual branch,
|
163 |
+
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
164 |
+
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
165 |
+
if zero_init_residual:
|
166 |
+
for m in self.modules():
|
167 |
+
if isinstance(m, Bottleneck):
|
168 |
+
nn.init.constant_(m.bn3.weight, 0)
|
169 |
+
elif isinstance(m, BasicBlock):
|
170 |
+
nn.init.constant_(m.bn2.weight, 0)
|
171 |
+
|
172 |
+
def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
|
173 |
+
norm_layer = self._norm_layer
|
174 |
+
downsample = None
|
175 |
+
previous_dilation = self.dilation
|
176 |
+
if dilate:
|
177 |
+
self.dilation *= stride
|
178 |
+
stride = 1
|
179 |
+
if stride != 1 or self.inplanes != planes * block.expansion:
|
180 |
+
downsample = nn.Sequential(
|
181 |
+
conv1x1(self.inplanes, planes * block.expansion, stride),
|
182 |
+
norm_layer(planes * block.expansion),
|
183 |
+
)
|
184 |
+
|
185 |
+
layers = []
|
186 |
+
layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
|
187 |
+
self.base_width, previous_dilation, norm_layer))
|
188 |
+
self.inplanes = planes * block.expansion
|
189 |
+
for _ in range(1, blocks):
|
190 |
+
layers.append(block(self.inplanes, planes, groups=self.groups,
|
191 |
+
base_width=self.base_width, dilation=self.dilation,
|
192 |
+
norm_layer=norm_layer))
|
193 |
+
|
194 |
+
return nn.Sequential(*layers)
|
195 |
+
|
196 |
+
def _forward_impl(self, x):
|
197 |
+
# See note [TorchScript super()]
|
198 |
+
features = []
|
199 |
+
x = self.conv1(x)
|
200 |
+
x = self.bn1(x)
|
201 |
+
x = self.relu(x)
|
202 |
+
x = self.maxpool(x)
|
203 |
+
|
204 |
+
x = self.layer1(x)
|
205 |
+
features.append(x)
|
206 |
+
|
207 |
+
x = self.layer2(x)
|
208 |
+
features.append(x)
|
209 |
+
|
210 |
+
x = self.layer3(x)
|
211 |
+
features.append(x)
|
212 |
+
|
213 |
+
x = self.layer4(x)
|
214 |
+
features.append(x)
|
215 |
+
|
216 |
+
#x = self.avgpool(x)
|
217 |
+
#x = torch.flatten(x, 1)
|
218 |
+
#x = self.fc(x)
|
219 |
+
|
220 |
+
return features
|
221 |
+
|
222 |
+
def forward(self, x):
|
223 |
+
return self._forward_impl(x)
|
224 |
+
|
225 |
+
|
226 |
+
|
227 |
+
def resnext101_32x8d(pretrained=True, **kwargs):
|
228 |
+
"""Constructs a ResNet-152 model.
|
229 |
+
Args:
|
230 |
+
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
231 |
+
"""
|
232 |
+
kwargs['groups'] = 32
|
233 |
+
kwargs['width_per_group'] = 8
|
234 |
+
|
235 |
+
model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
|
236 |
+
return model
|
237 |
+
|
controlnet_aux/leres/leres/__init__.py
ADDED
File without changes
|
controlnet_aux/leres/leres/__pycache__/Resnet.cpython-310.pyc
ADDED
Binary file (5.59 kB). View file
|
|
controlnet_aux/leres/leres/__pycache__/Resnext_torch.cpython-310.pyc
ADDED
Binary file (5.88 kB). View file
|
|
controlnet_aux/leres/leres/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (219 Bytes). View file
|
|
controlnet_aux/leres/leres/__pycache__/depthmap.cpython-310.pyc
ADDED
Binary file (11.7 kB). View file
|
|
controlnet_aux/leres/leres/__pycache__/multi_depth_model_woauxi.cpython-310.pyc
ADDED
Binary file (1.75 kB). View file
|
|
controlnet_aux/leres/leres/__pycache__/net_tools.cpython-310.pyc
ADDED
Binary file (1.93 kB). View file
|
|
controlnet_aux/leres/leres/__pycache__/network_auxi.cpython-310.pyc
ADDED
Binary file (9.86 kB). View file
|
|
controlnet_aux/leres/leres/depthmap.py
ADDED
@@ -0,0 +1,548 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Author: thygate
|
2 |
+
# https://github.com/thygate/stable-diffusion-webui-depthmap-script
|
3 |
+
|
4 |
+
import gc
|
5 |
+
from operator import getitem
|
6 |
+
|
7 |
+
import cv2
|
8 |
+
import numpy as np
|
9 |
+
import skimage.measure
|
10 |
+
import torch
|
11 |
+
from torchvision.transforms import transforms
|
12 |
+
|
13 |
+
from ...util import torch_gc
|
14 |
+
|
15 |
+
whole_size_threshold = 1600 # R_max from the paper
|
16 |
+
pix2pixsize = 1024
|
17 |
+
|
18 |
+
def scale_torch(img):
|
19 |
+
"""
|
20 |
+
Scale the image and output it in torch.tensor.
|
21 |
+
:param img: input rgb is in shape [H, W, C], input depth/disp is in shape [H, W]
|
22 |
+
:param scale: the scale factor. float
|
23 |
+
:return: img. [C, H, W]
|
24 |
+
"""
|
25 |
+
if len(img.shape) == 2:
|
26 |
+
img = img[np.newaxis, :, :]
|
27 |
+
if img.shape[2] == 3:
|
28 |
+
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406) , (0.229, 0.224, 0.225) )])
|
29 |
+
img = transform(img.astype(np.float32))
|
30 |
+
else:
|
31 |
+
img = img.astype(np.float32)
|
32 |
+
img = torch.from_numpy(img)
|
33 |
+
return img
|
34 |
+
|
35 |
+
def estimateleres(img, model, w, h):
|
36 |
+
device = next(iter(model.parameters())).device
|
37 |
+
# leres transform input
|
38 |
+
rgb_c = img[:, :, ::-1].copy()
|
39 |
+
A_resize = cv2.resize(rgb_c, (w, h))
|
40 |
+
img_torch = scale_torch(A_resize)[None, :, :, :]
|
41 |
+
|
42 |
+
# compute
|
43 |
+
with torch.no_grad():
|
44 |
+
img_torch = img_torch.to(device)
|
45 |
+
prediction = model.depth_model(img_torch)
|
46 |
+
|
47 |
+
prediction = prediction.squeeze().cpu().numpy()
|
48 |
+
prediction = cv2.resize(prediction, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC)
|
49 |
+
|
50 |
+
return prediction
|
51 |
+
|
52 |
+
def generatemask(size):
|
53 |
+
# Generates a Guassian mask
|
54 |
+
mask = np.zeros(size, dtype=np.float32)
|
55 |
+
sigma = int(size[0]/16)
|
56 |
+
k_size = int(2 * np.ceil(2 * int(size[0]/16)) + 1)
|
57 |
+
mask[int(0.15*size[0]):size[0] - int(0.15*size[0]), int(0.15*size[1]): size[1] - int(0.15*size[1])] = 1
|
58 |
+
mask = cv2.GaussianBlur(mask, (int(k_size), int(k_size)), sigma)
|
59 |
+
mask = (mask - mask.min()) / (mask.max() - mask.min())
|
60 |
+
mask = mask.astype(np.float32)
|
61 |
+
return mask
|
62 |
+
|
63 |
+
def resizewithpool(img, size):
|
64 |
+
i_size = img.shape[0]
|
65 |
+
n = int(np.floor(i_size/size))
|
66 |
+
|
67 |
+
out = skimage.measure.block_reduce(img, (n, n), np.max)
|
68 |
+
return out
|
69 |
+
|
70 |
+
def rgb2gray(rgb):
|
71 |
+
# Converts rgb to gray
|
72 |
+
return np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140])
|
73 |
+
|
74 |
+
def calculateprocessingres(img, basesize, confidence=0.1, scale_threshold=3, whole_size_threshold=3000):
|
75 |
+
# Returns the R_x resolution described in section 5 of the main paper.
|
76 |
+
|
77 |
+
# Parameters:
|
78 |
+
# img :input rgb image
|
79 |
+
# basesize : size the dilation kernel which is equal to receptive field of the network.
|
80 |
+
# confidence: value of x in R_x; allowed percentage of pixels that are not getting any contextual cue.
|
81 |
+
# scale_threshold: maximum allowed upscaling on the input image ; it has been set to 3.
|
82 |
+
# whole_size_threshold: maximum allowed resolution. (R_max from section 6 of the main paper)
|
83 |
+
|
84 |
+
# Returns:
|
85 |
+
# outputsize_scale*speed_scale :The computed R_x resolution
|
86 |
+
# patch_scale: K parameter from section 6 of the paper
|
87 |
+
|
88 |
+
# speed scale parameter is to process every image in a smaller size to accelerate the R_x resolution search
|
89 |
+
speed_scale = 32
|
90 |
+
image_dim = int(min(img.shape[0:2]))
|
91 |
+
|
92 |
+
gray = rgb2gray(img)
|
93 |
+
grad = np.abs(cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)) + np.abs(cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3))
|
94 |
+
grad = cv2.resize(grad, (image_dim, image_dim), cv2.INTER_AREA)
|
95 |
+
|
96 |
+
# thresholding the gradient map to generate the edge-map as a proxy of the contextual cues
|
97 |
+
m = grad.min()
|
98 |
+
M = grad.max()
|
99 |
+
middle = m + (0.4 * (M - m))
|
100 |
+
grad[grad < middle] = 0
|
101 |
+
grad[grad >= middle] = 1
|
102 |
+
|
103 |
+
# dilation kernel with size of the receptive field
|
104 |
+
kernel = np.ones((int(basesize/speed_scale), int(basesize/speed_scale)), float)
|
105 |
+
# dilation kernel with size of the a quarter of receptive field used to compute k
|
106 |
+
# as described in section 6 of main paper
|
107 |
+
kernel2 = np.ones((int(basesize / (4*speed_scale)), int(basesize / (4*speed_scale))), float)
|
108 |
+
|
109 |
+
# Output resolution limit set by the whole_size_threshold and scale_threshold.
|
110 |
+
threshold = min(whole_size_threshold, scale_threshold * max(img.shape[:2]))
|
111 |
+
|
112 |
+
outputsize_scale = basesize / speed_scale
|
113 |
+
for p_size in range(int(basesize/speed_scale), int(threshold/speed_scale), int(basesize / (2*speed_scale))):
|
114 |
+
grad_resized = resizewithpool(grad, p_size)
|
115 |
+
grad_resized = cv2.resize(grad_resized, (p_size, p_size), cv2.INTER_NEAREST)
|
116 |
+
grad_resized[grad_resized >= 0.5] = 1
|
117 |
+
grad_resized[grad_resized < 0.5] = 0
|
118 |
+
|
119 |
+
dilated = cv2.dilate(grad_resized, kernel, iterations=1)
|
120 |
+
meanvalue = (1-dilated).mean()
|
121 |
+
if meanvalue > confidence:
|
122 |
+
break
|
123 |
+
else:
|
124 |
+
outputsize_scale = p_size
|
125 |
+
|
126 |
+
grad_region = cv2.dilate(grad_resized, kernel2, iterations=1)
|
127 |
+
patch_scale = grad_region.mean()
|
128 |
+
|
129 |
+
return int(outputsize_scale*speed_scale), patch_scale
|
130 |
+
|
131 |
+
# Generate a double-input depth estimation
|
132 |
+
def doubleestimate(img, size1, size2, pix2pixsize, model, net_type, pix2pixmodel):
|
133 |
+
# Generate the low resolution estimation
|
134 |
+
estimate1 = singleestimate(img, size1, model, net_type)
|
135 |
+
# Resize to the inference size of merge network.
|
136 |
+
estimate1 = cv2.resize(estimate1, (pix2pixsize, pix2pixsize), interpolation=cv2.INTER_CUBIC)
|
137 |
+
|
138 |
+
# Generate the high resolution estimation
|
139 |
+
estimate2 = singleestimate(img, size2, model, net_type)
|
140 |
+
# Resize to the inference size of merge network.
|
141 |
+
estimate2 = cv2.resize(estimate2, (pix2pixsize, pix2pixsize), interpolation=cv2.INTER_CUBIC)
|
142 |
+
|
143 |
+
# Inference on the merge model
|
144 |
+
pix2pixmodel.set_input(estimate1, estimate2)
|
145 |
+
pix2pixmodel.test()
|
146 |
+
visuals = pix2pixmodel.get_current_visuals()
|
147 |
+
prediction_mapped = visuals['fake_B']
|
148 |
+
prediction_mapped = (prediction_mapped+1)/2
|
149 |
+
prediction_mapped = (prediction_mapped - torch.min(prediction_mapped)) / (
|
150 |
+
torch.max(prediction_mapped) - torch.min(prediction_mapped))
|
151 |
+
prediction_mapped = prediction_mapped.squeeze().cpu().numpy()
|
152 |
+
|
153 |
+
return prediction_mapped
|
154 |
+
|
155 |
+
# Generate a single-input depth estimation
|
156 |
+
def singleestimate(img, msize, model, net_type):
|
157 |
+
# if net_type == 0:
|
158 |
+
return estimateleres(img, model, msize, msize)
|
159 |
+
# else:
|
160 |
+
# return estimatemidasBoost(img, model, msize, msize)
|
161 |
+
|
162 |
+
def applyGridpatch(blsize, stride, img, box):
|
163 |
+
# Extract a simple grid patch.
|
164 |
+
counter1 = 0
|
165 |
+
patch_bound_list = {}
|
166 |
+
for k in range(blsize, img.shape[1] - blsize, stride):
|
167 |
+
for j in range(blsize, img.shape[0] - blsize, stride):
|
168 |
+
patch_bound_list[str(counter1)] = {}
|
169 |
+
patchbounds = [j - blsize, k - blsize, j - blsize + 2 * blsize, k - blsize + 2 * blsize]
|
170 |
+
patch_bound = [box[0] + patchbounds[1], box[1] + patchbounds[0], patchbounds[3] - patchbounds[1],
|
171 |
+
patchbounds[2] - patchbounds[0]]
|
172 |
+
patch_bound_list[str(counter1)]['rect'] = patch_bound
|
173 |
+
patch_bound_list[str(counter1)]['size'] = patch_bound[2]
|
174 |
+
counter1 = counter1 + 1
|
175 |
+
return patch_bound_list
|
176 |
+
|
177 |
+
# Generating local patches to perform the local refinement described in section 6 of the main paper.
|
178 |
+
def generatepatchs(img, base_size):
|
179 |
+
|
180 |
+
# Compute the gradients as a proxy of the contextual cues.
|
181 |
+
img_gray = rgb2gray(img)
|
182 |
+
whole_grad = np.abs(cv2.Sobel(img_gray, cv2.CV_64F, 0, 1, ksize=3)) +\
|
183 |
+
np.abs(cv2.Sobel(img_gray, cv2.CV_64F, 1, 0, ksize=3))
|
184 |
+
|
185 |
+
threshold = whole_grad[whole_grad > 0].mean()
|
186 |
+
whole_grad[whole_grad < threshold] = 0
|
187 |
+
|
188 |
+
# We use the integral image to speed-up the evaluation of the amount of gradients for each patch.
|
189 |
+
gf = whole_grad.sum()/len(whole_grad.reshape(-1))
|
190 |
+
grad_integral_image = cv2.integral(whole_grad)
|
191 |
+
|
192 |
+
# Variables are selected such that the initial patch size would be the receptive field size
|
193 |
+
# and the stride is set to 1/3 of the receptive field size.
|
194 |
+
blsize = int(round(base_size/2))
|
195 |
+
stride = int(round(blsize*0.75))
|
196 |
+
|
197 |
+
# Get initial Grid
|
198 |
+
patch_bound_list = applyGridpatch(blsize, stride, img, [0, 0, 0, 0])
|
199 |
+
|
200 |
+
# Refine initial Grid of patches by discarding the flat (in terms of gradients of the rgb image) ones. Refine
|
201 |
+
# each patch size to ensure that there will be enough depth cues for the network to generate a consistent depth map.
|
202 |
+
print("Selecting patches ...")
|
203 |
+
patch_bound_list = adaptiveselection(grad_integral_image, patch_bound_list, gf)
|
204 |
+
|
205 |
+
# Sort the patch list to make sure the merging operation will be done with the correct order: starting from biggest
|
206 |
+
# patch
|
207 |
+
patchset = sorted(patch_bound_list.items(), key=lambda x: getitem(x[1], 'size'), reverse=True)
|
208 |
+
return patchset
|
209 |
+
|
210 |
+
def getGF_fromintegral(integralimage, rect):
|
211 |
+
# Computes the gradient density of a given patch from the gradient integral image.
|
212 |
+
x1 = rect[1]
|
213 |
+
x2 = rect[1]+rect[3]
|
214 |
+
y1 = rect[0]
|
215 |
+
y2 = rect[0]+rect[2]
|
216 |
+
value = integralimage[x2, y2]-integralimage[x1, y2]-integralimage[x2, y1]+integralimage[x1, y1]
|
217 |
+
return value
|
218 |
+
|
219 |
+
# Adaptively select patches
|
220 |
+
def adaptiveselection(integral_grad, patch_bound_list, gf):
|
221 |
+
patchlist = {}
|
222 |
+
count = 0
|
223 |
+
height, width = integral_grad.shape
|
224 |
+
|
225 |
+
search_step = int(32/factor)
|
226 |
+
|
227 |
+
# Go through all patches
|
228 |
+
for c in range(len(patch_bound_list)):
|
229 |
+
# Get patch
|
230 |
+
bbox = patch_bound_list[str(c)]['rect']
|
231 |
+
|
232 |
+
# Compute the amount of gradients present in the patch from the integral image.
|
233 |
+
cgf = getGF_fromintegral(integral_grad, bbox)/(bbox[2]*bbox[3])
|
234 |
+
|
235 |
+
# Check if patching is beneficial by comparing the gradient density of the patch to
|
236 |
+
# the gradient density of the whole image
|
237 |
+
if cgf >= gf:
|
238 |
+
bbox_test = bbox.copy()
|
239 |
+
patchlist[str(count)] = {}
|
240 |
+
|
241 |
+
# Enlarge each patch until the gradient density of the patch is equal
|
242 |
+
# to the whole image gradient density
|
243 |
+
while True:
|
244 |
+
|
245 |
+
bbox_test[0] = bbox_test[0] - int(search_step/2)
|
246 |
+
bbox_test[1] = bbox_test[1] - int(search_step/2)
|
247 |
+
|
248 |
+
bbox_test[2] = bbox_test[2] + search_step
|
249 |
+
bbox_test[3] = bbox_test[3] + search_step
|
250 |
+
|
251 |
+
# Check if we are still within the image
|
252 |
+
if bbox_test[0] < 0 or bbox_test[1] < 0 or bbox_test[1] + bbox_test[3] >= height \
|
253 |
+
or bbox_test[0] + bbox_test[2] >= width:
|
254 |
+
break
|
255 |
+
|
256 |
+
# Compare gradient density
|
257 |
+
cgf = getGF_fromintegral(integral_grad, bbox_test)/(bbox_test[2]*bbox_test[3])
|
258 |
+
if cgf < gf:
|
259 |
+
break
|
260 |
+
bbox = bbox_test.copy()
|
261 |
+
|
262 |
+
# Add patch to selected patches
|
263 |
+
patchlist[str(count)]['rect'] = bbox
|
264 |
+
patchlist[str(count)]['size'] = bbox[2]
|
265 |
+
count = count + 1
|
266 |
+
|
267 |
+
# Return selected patches
|
268 |
+
return patchlist
|
269 |
+
|
270 |
+
def impatch(image, rect):
|
271 |
+
# Extract the given patch pixels from a given image.
|
272 |
+
w1 = rect[0]
|
273 |
+
h1 = rect[1]
|
274 |
+
w2 = w1 + rect[2]
|
275 |
+
h2 = h1 + rect[3]
|
276 |
+
image_patch = image[h1:h2, w1:w2]
|
277 |
+
return image_patch
|
278 |
+
|
279 |
+
class ImageandPatchs:
|
280 |
+
def __init__(self, root_dir, name, patchsinfo, rgb_image, scale=1):
|
281 |
+
self.root_dir = root_dir
|
282 |
+
self.patchsinfo = patchsinfo
|
283 |
+
self.name = name
|
284 |
+
self.patchs = patchsinfo
|
285 |
+
self.scale = scale
|
286 |
+
|
287 |
+
self.rgb_image = cv2.resize(rgb_image, (round(rgb_image.shape[1]*scale), round(rgb_image.shape[0]*scale)),
|
288 |
+
interpolation=cv2.INTER_CUBIC)
|
289 |
+
|
290 |
+
self.do_have_estimate = False
|
291 |
+
self.estimation_updated_image = None
|
292 |
+
self.estimation_base_image = None
|
293 |
+
|
294 |
+
def __len__(self):
|
295 |
+
return len(self.patchs)
|
296 |
+
|
297 |
+
def set_base_estimate(self, est):
|
298 |
+
self.estimation_base_image = est
|
299 |
+
if self.estimation_updated_image is not None:
|
300 |
+
self.do_have_estimate = True
|
301 |
+
|
302 |
+
def set_updated_estimate(self, est):
|
303 |
+
self.estimation_updated_image = est
|
304 |
+
if self.estimation_base_image is not None:
|
305 |
+
self.do_have_estimate = True
|
306 |
+
|
307 |
+
def __getitem__(self, index):
|
308 |
+
patch_id = int(self.patchs[index][0])
|
309 |
+
rect = np.array(self.patchs[index][1]['rect'])
|
310 |
+
msize = self.patchs[index][1]['size']
|
311 |
+
|
312 |
+
## applying scale to rect:
|
313 |
+
rect = np.round(rect * self.scale)
|
314 |
+
rect = rect.astype('int')
|
315 |
+
msize = round(msize * self.scale)
|
316 |
+
|
317 |
+
patch_rgb = impatch(self.rgb_image, rect)
|
318 |
+
if self.do_have_estimate:
|
319 |
+
patch_whole_estimate_base = impatch(self.estimation_base_image, rect)
|
320 |
+
patch_whole_estimate_updated = impatch(self.estimation_updated_image, rect)
|
321 |
+
return {'patch_rgb': patch_rgb, 'patch_whole_estimate_base': patch_whole_estimate_base,
|
322 |
+
'patch_whole_estimate_updated': patch_whole_estimate_updated, 'rect': rect,
|
323 |
+
'size': msize, 'id': patch_id}
|
324 |
+
else:
|
325 |
+
return {'patch_rgb': patch_rgb, 'rect': rect, 'size': msize, 'id': patch_id}
|
326 |
+
|
327 |
+
def print_options(self, opt):
|
328 |
+
"""Print and save options
|
329 |
+
|
330 |
+
It will print both current options and default values(if different).
|
331 |
+
It will save options into a text file / [checkpoints_dir] / opt.txt
|
332 |
+
"""
|
333 |
+
message = ''
|
334 |
+
message += '----------------- Options ---------------\n'
|
335 |
+
for k, v in sorted(vars(opt).items()):
|
336 |
+
comment = ''
|
337 |
+
default = self.parser.get_default(k)
|
338 |
+
if v != default:
|
339 |
+
comment = '\t[default: %s]' % str(default)
|
340 |
+
message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
|
341 |
+
message += '----------------- End -------------------'
|
342 |
+
print(message)
|
343 |
+
|
344 |
+
# save to the disk
|
345 |
+
"""
|
346 |
+
expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
|
347 |
+
util.mkdirs(expr_dir)
|
348 |
+
file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))
|
349 |
+
with open(file_name, 'wt') as opt_file:
|
350 |
+
opt_file.write(message)
|
351 |
+
opt_file.write('\n')
|
352 |
+
"""
|
353 |
+
|
354 |
+
def parse(self):
|
355 |
+
"""Parse our options, create checkpoints directory suffix, and set up gpu device."""
|
356 |
+
opt = self.gather_options()
|
357 |
+
opt.isTrain = self.isTrain # train or test
|
358 |
+
|
359 |
+
# process opt.suffix
|
360 |
+
if opt.suffix:
|
361 |
+
suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
|
362 |
+
opt.name = opt.name + suffix
|
363 |
+
|
364 |
+
#self.print_options(opt)
|
365 |
+
|
366 |
+
# set gpu ids
|
367 |
+
str_ids = opt.gpu_ids.split(',')
|
368 |
+
opt.gpu_ids = []
|
369 |
+
for str_id in str_ids:
|
370 |
+
id = int(str_id)
|
371 |
+
if id >= 0:
|
372 |
+
opt.gpu_ids.append(id)
|
373 |
+
#if len(opt.gpu_ids) > 0:
|
374 |
+
# torch.cuda.set_device(opt.gpu_ids[0])
|
375 |
+
|
376 |
+
self.opt = opt
|
377 |
+
return self.opt
|
378 |
+
|
379 |
+
|
380 |
+
def estimateboost(img, model, model_type, pix2pixmodel, max_res=512, depthmap_script_boost_rmax=None):
|
381 |
+
global whole_size_threshold
|
382 |
+
|
383 |
+
# get settings
|
384 |
+
if depthmap_script_boost_rmax:
|
385 |
+
whole_size_threshold = depthmap_script_boost_rmax
|
386 |
+
|
387 |
+
if model_type == 0: #leres
|
388 |
+
net_receptive_field_size = 448
|
389 |
+
patch_netsize = 2 * net_receptive_field_size
|
390 |
+
elif model_type == 1: #dpt_beit_large_512
|
391 |
+
net_receptive_field_size = 512
|
392 |
+
patch_netsize = 2 * net_receptive_field_size
|
393 |
+
else: #other midas
|
394 |
+
net_receptive_field_size = 384
|
395 |
+
patch_netsize = 2 * net_receptive_field_size
|
396 |
+
|
397 |
+
gc.collect()
|
398 |
+
torch_gc()
|
399 |
+
|
400 |
+
# Generate mask used to smoothly blend the local pathc estimations to the base estimate.
|
401 |
+
# It is arbitrarily large to avoid artifacts during rescaling for each crop.
|
402 |
+
mask_org = generatemask((3000, 3000))
|
403 |
+
mask = mask_org.copy()
|
404 |
+
|
405 |
+
# Value x of R_x defined in the section 5 of the main paper.
|
406 |
+
r_threshold_value = 0.2
|
407 |
+
#if R0:
|
408 |
+
# r_threshold_value = 0
|
409 |
+
|
410 |
+
input_resolution = img.shape
|
411 |
+
scale_threshold = 3 # Allows up-scaling with a scale up to 3
|
412 |
+
|
413 |
+
# Find the best input resolution R-x. The resolution search described in section 5-double estimation of the main paper and section B of the
|
414 |
+
# supplementary material.
|
415 |
+
whole_image_optimal_size, patch_scale = calculateprocessingres(img, net_receptive_field_size, r_threshold_value, scale_threshold, whole_size_threshold)
|
416 |
+
|
417 |
+
# print('wholeImage being processed in :', whole_image_optimal_size)
|
418 |
+
|
419 |
+
# Generate the base estimate using the double estimation.
|
420 |
+
whole_estimate = doubleestimate(img, net_receptive_field_size, whole_image_optimal_size, pix2pixsize, model, model_type, pix2pixmodel)
|
421 |
+
|
422 |
+
# Compute the multiplier described in section 6 of the main paper to make sure our initial patch can select
|
423 |
+
# small high-density regions of the image.
|
424 |
+
global factor
|
425 |
+
factor = max(min(1, 4 * patch_scale * whole_image_optimal_size / whole_size_threshold), 0.2)
|
426 |
+
# print('Adjust factor is:', 1/factor)
|
427 |
+
|
428 |
+
# Check if Local boosting is beneficial.
|
429 |
+
if max_res < whole_image_optimal_size:
|
430 |
+
# print("No Local boosting. Specified Max Res is smaller than R20, Returning doubleestimate result")
|
431 |
+
return cv2.resize(whole_estimate, (input_resolution[1], input_resolution[0]), interpolation=cv2.INTER_CUBIC)
|
432 |
+
|
433 |
+
# Compute the default target resolution.
|
434 |
+
if img.shape[0] > img.shape[1]:
|
435 |
+
a = 2 * whole_image_optimal_size
|
436 |
+
b = round(2 * whole_image_optimal_size * img.shape[1] / img.shape[0])
|
437 |
+
else:
|
438 |
+
a = round(2 * whole_image_optimal_size * img.shape[0] / img.shape[1])
|
439 |
+
b = 2 * whole_image_optimal_size
|
440 |
+
b = int(round(b / factor))
|
441 |
+
a = int(round(a / factor))
|
442 |
+
|
443 |
+
"""
|
444 |
+
# recompute a, b and saturate to max res.
|
445 |
+
if max(a,b) > max_res:
|
446 |
+
print('Default Res is higher than max-res: Reducing final resolution')
|
447 |
+
if img.shape[0] > img.shape[1]:
|
448 |
+
a = max_res
|
449 |
+
b = round(max_res * img.shape[1] / img.shape[0])
|
450 |
+
else:
|
451 |
+
a = round(max_res * img.shape[0] / img.shape[1])
|
452 |
+
b = max_res
|
453 |
+
b = int(b)
|
454 |
+
a = int(a)
|
455 |
+
"""
|
456 |
+
|
457 |
+
img = cv2.resize(img, (b, a), interpolation=cv2.INTER_CUBIC)
|
458 |
+
|
459 |
+
# Extract selected patches for local refinement
|
460 |
+
base_size = net_receptive_field_size * 2
|
461 |
+
patchset = generatepatchs(img, base_size)
|
462 |
+
|
463 |
+
# print('Target resolution: ', img.shape)
|
464 |
+
|
465 |
+
# Computing a scale in case user prompted to generate the results as the same resolution of the input.
|
466 |
+
# Notice that our method output resolution is independent of the input resolution and this parameter will only
|
467 |
+
# enable a scaling operation during the local patch merge implementation to generate results with the same resolution
|
468 |
+
# as the input.
|
469 |
+
"""
|
470 |
+
if output_resolution == 1:
|
471 |
+
mergein_scale = input_resolution[0] / img.shape[0]
|
472 |
+
print('Dynamicly change merged-in resolution; scale:', mergein_scale)
|
473 |
+
else:
|
474 |
+
mergein_scale = 1
|
475 |
+
"""
|
476 |
+
# always rescale to input res for now
|
477 |
+
mergein_scale = input_resolution[0] / img.shape[0]
|
478 |
+
|
479 |
+
imageandpatchs = ImageandPatchs('', '', patchset, img, mergein_scale)
|
480 |
+
whole_estimate_resized = cv2.resize(whole_estimate, (round(img.shape[1]*mergein_scale),
|
481 |
+
round(img.shape[0]*mergein_scale)), interpolation=cv2.INTER_CUBIC)
|
482 |
+
imageandpatchs.set_base_estimate(whole_estimate_resized.copy())
|
483 |
+
imageandpatchs.set_updated_estimate(whole_estimate_resized.copy())
|
484 |
+
|
485 |
+
print('Resulting depthmap resolution will be :', whole_estimate_resized.shape[:2])
|
486 |
+
print('Patches to process: '+str(len(imageandpatchs)))
|
487 |
+
|
488 |
+
# Enumerate through all patches, generate their estimations and refining the base estimate.
|
489 |
+
for patch_ind in range(len(imageandpatchs)):
|
490 |
+
|
491 |
+
# Get patch information
|
492 |
+
patch = imageandpatchs[patch_ind] # patch object
|
493 |
+
patch_rgb = patch['patch_rgb'] # rgb patch
|
494 |
+
patch_whole_estimate_base = patch['patch_whole_estimate_base'] # corresponding patch from base
|
495 |
+
rect = patch['rect'] # patch size and location
|
496 |
+
patch_id = patch['id'] # patch ID
|
497 |
+
org_size = patch_whole_estimate_base.shape # the original size from the unscaled input
|
498 |
+
print('\t Processing patch', patch_ind, '/', len(imageandpatchs)-1, '|', rect)
|
499 |
+
|
500 |
+
# We apply double estimation for patches. The high resolution value is fixed to twice the receptive
|
501 |
+
# field size of the network for patches to accelerate the process.
|
502 |
+
patch_estimation = doubleestimate(patch_rgb, net_receptive_field_size, patch_netsize, pix2pixsize, model, model_type, pix2pixmodel)
|
503 |
+
patch_estimation = cv2.resize(patch_estimation, (pix2pixsize, pix2pixsize), interpolation=cv2.INTER_CUBIC)
|
504 |
+
patch_whole_estimate_base = cv2.resize(patch_whole_estimate_base, (pix2pixsize, pix2pixsize), interpolation=cv2.INTER_CUBIC)
|
505 |
+
|
506 |
+
# Merging the patch estimation into the base estimate using our merge network:
|
507 |
+
# We feed the patch estimation and the same region from the updated base estimate to the merge network
|
508 |
+
# to generate the target estimate for the corresponding region.
|
509 |
+
pix2pixmodel.set_input(patch_whole_estimate_base, patch_estimation)
|
510 |
+
|
511 |
+
# Run merging network
|
512 |
+
pix2pixmodel.test()
|
513 |
+
visuals = pix2pixmodel.get_current_visuals()
|
514 |
+
|
515 |
+
prediction_mapped = visuals['fake_B']
|
516 |
+
prediction_mapped = (prediction_mapped+1)/2
|
517 |
+
prediction_mapped = prediction_mapped.squeeze().cpu().numpy()
|
518 |
+
|
519 |
+
mapped = prediction_mapped
|
520 |
+
|
521 |
+
# We use a simple linear polynomial to make sure the result of the merge network would match the values of
|
522 |
+
# base estimate
|
523 |
+
p_coef = np.polyfit(mapped.reshape(-1), patch_whole_estimate_base.reshape(-1), deg=1)
|
524 |
+
merged = np.polyval(p_coef, mapped.reshape(-1)).reshape(mapped.shape)
|
525 |
+
|
526 |
+
merged = cv2.resize(merged, (org_size[1],org_size[0]), interpolation=cv2.INTER_CUBIC)
|
527 |
+
|
528 |
+
# Get patch size and location
|
529 |
+
w1 = rect[0]
|
530 |
+
h1 = rect[1]
|
531 |
+
w2 = w1 + rect[2]
|
532 |
+
h2 = h1 + rect[3]
|
533 |
+
|
534 |
+
# To speed up the implementation, we only generate the Gaussian mask once with a sufficiently large size
|
535 |
+
# and resize it to our needed size while merging the patches.
|
536 |
+
if mask.shape != org_size:
|
537 |
+
mask = cv2.resize(mask_org, (org_size[1],org_size[0]), interpolation=cv2.INTER_LINEAR)
|
538 |
+
|
539 |
+
tobemergedto = imageandpatchs.estimation_updated_image
|
540 |
+
|
541 |
+
# Update the whole estimation:
|
542 |
+
# We use a simple Gaussian mask to blend the merged patch region with the base estimate to ensure seamless
|
543 |
+
# blending at the boundaries of the patch region.
|
544 |
+
tobemergedto[h1:h2, w1:w2] = np.multiply(tobemergedto[h1:h2, w1:w2], 1 - mask) + np.multiply(merged, mask)
|
545 |
+
imageandpatchs.set_updated_estimate(tobemergedto)
|
546 |
+
|
547 |
+
# output
|
548 |
+
return cv2.resize(imageandpatchs.estimation_updated_image, (input_resolution[1], input_resolution[0]), interpolation=cv2.INTER_CUBIC)
|
controlnet_aux/leres/leres/multi_depth_model_woauxi.py
ADDED
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
|
4 |
+
from . import network_auxi as network
|
5 |
+
from .net_tools import get_func
|
6 |
+
|
7 |
+
|
8 |
+
class RelDepthModel(nn.Module):
|
9 |
+
def __init__(self, backbone='resnet50'):
|
10 |
+
super(RelDepthModel, self).__init__()
|
11 |
+
if backbone == 'resnet50':
|
12 |
+
encoder = 'resnet50_stride32'
|
13 |
+
elif backbone == 'resnext101':
|
14 |
+
encoder = 'resnext101_stride32x8d'
|
15 |
+
self.depth_model = DepthModel(encoder)
|
16 |
+
|
17 |
+
def inference(self, rgb):
|
18 |
+
with torch.no_grad():
|
19 |
+
input = rgb.to(self.depth_model.device)
|
20 |
+
depth = self.depth_model(input)
|
21 |
+
#pred_depth_out = depth - depth.min() + 0.01
|
22 |
+
return depth #pred_depth_out
|
23 |
+
|
24 |
+
|
25 |
+
class DepthModel(nn.Module):
|
26 |
+
def __init__(self, encoder):
|
27 |
+
super(DepthModel, self).__init__()
|
28 |
+
backbone = network.__name__.split('.')[-1] + '.' + encoder
|
29 |
+
self.encoder_modules = get_func(backbone)()
|
30 |
+
self.decoder_modules = network.Decoder()
|
31 |
+
|
32 |
+
def forward(self, x):
|
33 |
+
lateral_out = self.encoder_modules(x)
|
34 |
+
out_logit = self.decoder_modules(lateral_out)
|
35 |
+
return out_logit
|
controlnet_aux/leres/leres/net_tools.py
ADDED
@@ -0,0 +1,54 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import importlib
|
2 |
+
import torch
|
3 |
+
import os
|
4 |
+
from collections import OrderedDict
|
5 |
+
|
6 |
+
|
7 |
+
def get_func(func_name):
|
8 |
+
"""Helper to return a function object by name. func_name must identify a
|
9 |
+
function in this module or the path to a function relative to the base
|
10 |
+
'modeling' module.
|
11 |
+
"""
|
12 |
+
if func_name == '':
|
13 |
+
return None
|
14 |
+
try:
|
15 |
+
parts = func_name.split('.')
|
16 |
+
# Refers to a function in this module
|
17 |
+
if len(parts) == 1:
|
18 |
+
return globals()[parts[0]]
|
19 |
+
# Otherwise, assume we're referencing a module under modeling
|
20 |
+
module_name = 'controlnet_aux.leres.leres.' + '.'.join(parts[:-1])
|
21 |
+
module = importlib.import_module(module_name)
|
22 |
+
return getattr(module, parts[-1])
|
23 |
+
except Exception:
|
24 |
+
print('Failed to f1ind function: %s', func_name)
|
25 |
+
raise
|
26 |
+
|
27 |
+
def load_ckpt(args, depth_model, shift_model, focal_model):
|
28 |
+
"""
|
29 |
+
Load checkpoint.
|
30 |
+
"""
|
31 |
+
if os.path.isfile(args.load_ckpt):
|
32 |
+
print("loading checkpoint %s" % args.load_ckpt)
|
33 |
+
checkpoint = torch.load(args.load_ckpt)
|
34 |
+
if shift_model is not None:
|
35 |
+
shift_model.load_state_dict(strip_prefix_if_present(checkpoint['shift_model'], 'module.'),
|
36 |
+
strict=True)
|
37 |
+
if focal_model is not None:
|
38 |
+
focal_model.load_state_dict(strip_prefix_if_present(checkpoint['focal_model'], 'module.'),
|
39 |
+
strict=True)
|
40 |
+
depth_model.load_state_dict(strip_prefix_if_present(checkpoint['depth_model'], "module."),
|
41 |
+
strict=True)
|
42 |
+
del checkpoint
|
43 |
+
if torch.cuda.is_available():
|
44 |
+
torch.cuda.empty_cache()
|
45 |
+
|
46 |
+
|
47 |
+
def strip_prefix_if_present(state_dict, prefix):
|
48 |
+
keys = sorted(state_dict.keys())
|
49 |
+
if not all(key.startswith(prefix) for key in keys):
|
50 |
+
return state_dict
|
51 |
+
stripped_state_dict = OrderedDict()
|
52 |
+
for key, value in state_dict.items():
|
53 |
+
stripped_state_dict[key.replace(prefix, "")] = value
|
54 |
+
return stripped_state_dict
|
controlnet_aux/leres/leres/network_auxi.py
ADDED
@@ -0,0 +1,417 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
import torch.nn.init as init
|
4 |
+
|
5 |
+
from . import Resnet, Resnext_torch
|
6 |
+
|
7 |
+
|
8 |
+
def resnet50_stride32():
|
9 |
+
return DepthNet(backbone='resnet', depth=50, upfactors=[2, 2, 2, 2])
|
10 |
+
|
11 |
+
def resnext101_stride32x8d():
|
12 |
+
return DepthNet(backbone='resnext101_32x8d', depth=101, upfactors=[2, 2, 2, 2])
|
13 |
+
|
14 |
+
|
15 |
+
class Decoder(nn.Module):
|
16 |
+
def __init__(self):
|
17 |
+
super(Decoder, self).__init__()
|
18 |
+
self.inchannels = [256, 512, 1024, 2048]
|
19 |
+
self.midchannels = [256, 256, 256, 512]
|
20 |
+
self.upfactors = [2,2,2,2]
|
21 |
+
self.outchannels = 1
|
22 |
+
|
23 |
+
self.conv = FTB(inchannels=self.inchannels[3], midchannels=self.midchannels[3])
|
24 |
+
self.conv1 = nn.Conv2d(in_channels=self.midchannels[3], out_channels=self.midchannels[2], kernel_size=3, padding=1, stride=1, bias=True)
|
25 |
+
self.upsample = nn.Upsample(scale_factor=self.upfactors[3], mode='bilinear', align_corners=True)
|
26 |
+
|
27 |
+
self.ffm2 = FFM(inchannels=self.inchannels[2], midchannels=self.midchannels[2], outchannels = self.midchannels[2], upfactor=self.upfactors[2])
|
28 |
+
self.ffm1 = FFM(inchannels=self.inchannels[1], midchannels=self.midchannels[1], outchannels = self.midchannels[1], upfactor=self.upfactors[1])
|
29 |
+
self.ffm0 = FFM(inchannels=self.inchannels[0], midchannels=self.midchannels[0], outchannels = self.midchannels[0], upfactor=self.upfactors[0])
|
30 |
+
|
31 |
+
self.outconv = AO(inchannels=self.midchannels[0], outchannels=self.outchannels, upfactor=2)
|
32 |
+
self._init_params()
|
33 |
+
|
34 |
+
def _init_params(self):
|
35 |
+
for m in self.modules():
|
36 |
+
if isinstance(m, nn.Conv2d):
|
37 |
+
init.normal_(m.weight, std=0.01)
|
38 |
+
if m.bias is not None:
|
39 |
+
init.constant_(m.bias, 0)
|
40 |
+
elif isinstance(m, nn.ConvTranspose2d):
|
41 |
+
init.normal_(m.weight, std=0.01)
|
42 |
+
if m.bias is not None:
|
43 |
+
init.constant_(m.bias, 0)
|
44 |
+
elif isinstance(m, nn.BatchNorm2d): #NN.BatchNorm2d
|
45 |
+
init.constant_(m.weight, 1)
|
46 |
+
init.constant_(m.bias, 0)
|
47 |
+
elif isinstance(m, nn.Linear):
|
48 |
+
init.normal_(m.weight, std=0.01)
|
49 |
+
if m.bias is not None:
|
50 |
+
init.constant_(m.bias, 0)
|
51 |
+
|
52 |
+
def forward(self, features):
|
53 |
+
x_32x = self.conv(features[3]) # 1/32
|
54 |
+
x_32 = self.conv1(x_32x)
|
55 |
+
x_16 = self.upsample(x_32) # 1/16
|
56 |
+
|
57 |
+
x_8 = self.ffm2(features[2], x_16) # 1/8
|
58 |
+
x_4 = self.ffm1(features[1], x_8) # 1/4
|
59 |
+
x_2 = self.ffm0(features[0], x_4) # 1/2
|
60 |
+
#-----------------------------------------
|
61 |
+
x = self.outconv(x_2) # original size
|
62 |
+
return x
|
63 |
+
|
64 |
+
class DepthNet(nn.Module):
|
65 |
+
__factory = {
|
66 |
+
18: Resnet.resnet18,
|
67 |
+
34: Resnet.resnet34,
|
68 |
+
50: Resnet.resnet50,
|
69 |
+
101: Resnet.resnet101,
|
70 |
+
152: Resnet.resnet152
|
71 |
+
}
|
72 |
+
def __init__(self,
|
73 |
+
backbone='resnet',
|
74 |
+
depth=50,
|
75 |
+
upfactors=[2, 2, 2, 2]):
|
76 |
+
super(DepthNet, self).__init__()
|
77 |
+
self.backbone = backbone
|
78 |
+
self.depth = depth
|
79 |
+
self.pretrained = False
|
80 |
+
self.inchannels = [256, 512, 1024, 2048]
|
81 |
+
self.midchannels = [256, 256, 256, 512]
|
82 |
+
self.upfactors = upfactors
|
83 |
+
self.outchannels = 1
|
84 |
+
|
85 |
+
# Build model
|
86 |
+
if self.backbone == 'resnet':
|
87 |
+
if self.depth not in DepthNet.__factory:
|
88 |
+
raise KeyError("Unsupported depth:", self.depth)
|
89 |
+
self.encoder = DepthNet.__factory[depth](pretrained=self.pretrained)
|
90 |
+
elif self.backbone == 'resnext101_32x8d':
|
91 |
+
self.encoder = Resnext_torch.resnext101_32x8d(pretrained=self.pretrained)
|
92 |
+
else:
|
93 |
+
self.encoder = Resnext_torch.resnext101(pretrained=self.pretrained)
|
94 |
+
|
95 |
+
def forward(self, x):
|
96 |
+
x = self.encoder(x) # 1/32, 1/16, 1/8, 1/4
|
97 |
+
return x
|
98 |
+
|
99 |
+
|
100 |
+
class FTB(nn.Module):
|
101 |
+
def __init__(self, inchannels, midchannels=512):
|
102 |
+
super(FTB, self).__init__()
|
103 |
+
self.in1 = inchannels
|
104 |
+
self.mid = midchannels
|
105 |
+
self.conv1 = nn.Conv2d(in_channels=self.in1, out_channels=self.mid, kernel_size=3, padding=1, stride=1,
|
106 |
+
bias=True)
|
107 |
+
# NN.BatchNorm2d
|
108 |
+
self.conv_branch = nn.Sequential(nn.ReLU(inplace=True), \
|
109 |
+
nn.Conv2d(in_channels=self.mid, out_channels=self.mid, kernel_size=3,
|
110 |
+
padding=1, stride=1, bias=True), \
|
111 |
+
nn.BatchNorm2d(num_features=self.mid), \
|
112 |
+
nn.ReLU(inplace=True), \
|
113 |
+
nn.Conv2d(in_channels=self.mid, out_channels=self.mid, kernel_size=3,
|
114 |
+
padding=1, stride=1, bias=True))
|
115 |
+
self.relu = nn.ReLU(inplace=True)
|
116 |
+
|
117 |
+
self.init_params()
|
118 |
+
|
119 |
+
def forward(self, x):
|
120 |
+
x = self.conv1(x)
|
121 |
+
x = x + self.conv_branch(x)
|
122 |
+
x = self.relu(x)
|
123 |
+
|
124 |
+
return x
|
125 |
+
|
126 |
+
def init_params(self):
|
127 |
+
for m in self.modules():
|
128 |
+
if isinstance(m, nn.Conv2d):
|
129 |
+
init.normal_(m.weight, std=0.01)
|
130 |
+
if m.bias is not None:
|
131 |
+
init.constant_(m.bias, 0)
|
132 |
+
elif isinstance(m, nn.ConvTranspose2d):
|
133 |
+
# init.kaiming_normal_(m.weight, mode='fan_out')
|
134 |
+
init.normal_(m.weight, std=0.01)
|
135 |
+
# init.xavier_normal_(m.weight)
|
136 |
+
if m.bias is not None:
|
137 |
+
init.constant_(m.bias, 0)
|
138 |
+
elif isinstance(m, nn.BatchNorm2d): # NN.BatchNorm2d
|
139 |
+
init.constant_(m.weight, 1)
|
140 |
+
init.constant_(m.bias, 0)
|
141 |
+
elif isinstance(m, nn.Linear):
|
142 |
+
init.normal_(m.weight, std=0.01)
|
143 |
+
if m.bias is not None:
|
144 |
+
init.constant_(m.bias, 0)
|
145 |
+
|
146 |
+
|
147 |
+
class ATA(nn.Module):
|
148 |
+
def __init__(self, inchannels, reduction=8):
|
149 |
+
super(ATA, self).__init__()
|
150 |
+
self.inchannels = inchannels
|
151 |
+
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
152 |
+
self.fc = nn.Sequential(nn.Linear(self.inchannels * 2, self.inchannels // reduction),
|
153 |
+
nn.ReLU(inplace=True),
|
154 |
+
nn.Linear(self.inchannels // reduction, self.inchannels),
|
155 |
+
nn.Sigmoid())
|
156 |
+
self.init_params()
|
157 |
+
|
158 |
+
def forward(self, low_x, high_x):
|
159 |
+
n, c, _, _ = low_x.size()
|
160 |
+
x = torch.cat([low_x, high_x], 1)
|
161 |
+
x = self.avg_pool(x)
|
162 |
+
x = x.view(n, -1)
|
163 |
+
x = self.fc(x).view(n, c, 1, 1)
|
164 |
+
x = low_x * x + high_x
|
165 |
+
|
166 |
+
return x
|
167 |
+
|
168 |
+
def init_params(self):
|
169 |
+
for m in self.modules():
|
170 |
+
if isinstance(m, nn.Conv2d):
|
171 |
+
# init.kaiming_normal_(m.weight, mode='fan_out')
|
172 |
+
# init.normal(m.weight, std=0.01)
|
173 |
+
init.xavier_normal_(m.weight)
|
174 |
+
if m.bias is not None:
|
175 |
+
init.constant_(m.bias, 0)
|
176 |
+
elif isinstance(m, nn.ConvTranspose2d):
|
177 |
+
# init.kaiming_normal_(m.weight, mode='fan_out')
|
178 |
+
# init.normal_(m.weight, std=0.01)
|
179 |
+
init.xavier_normal_(m.weight)
|
180 |
+
if m.bias is not None:
|
181 |
+
init.constant_(m.bias, 0)
|
182 |
+
elif isinstance(m, nn.BatchNorm2d): # NN.BatchNorm2d
|
183 |
+
init.constant_(m.weight, 1)
|
184 |
+
init.constant_(m.bias, 0)
|
185 |
+
elif isinstance(m, nn.Linear):
|
186 |
+
init.normal_(m.weight, std=0.01)
|
187 |
+
if m.bias is not None:
|
188 |
+
init.constant_(m.bias, 0)
|
189 |
+
|
190 |
+
|
191 |
+
class FFM(nn.Module):
|
192 |
+
def __init__(self, inchannels, midchannels, outchannels, upfactor=2):
|
193 |
+
super(FFM, self).__init__()
|
194 |
+
self.inchannels = inchannels
|
195 |
+
self.midchannels = midchannels
|
196 |
+
self.outchannels = outchannels
|
197 |
+
self.upfactor = upfactor
|
198 |
+
|
199 |
+
self.ftb1 = FTB(inchannels=self.inchannels, midchannels=self.midchannels)
|
200 |
+
# self.ata = ATA(inchannels = self.midchannels)
|
201 |
+
self.ftb2 = FTB(inchannels=self.midchannels, midchannels=self.outchannels)
|
202 |
+
|
203 |
+
self.upsample = nn.Upsample(scale_factor=self.upfactor, mode='bilinear', align_corners=True)
|
204 |
+
|
205 |
+
self.init_params()
|
206 |
+
|
207 |
+
def forward(self, low_x, high_x):
|
208 |
+
x = self.ftb1(low_x)
|
209 |
+
x = x + high_x
|
210 |
+
x = self.ftb2(x)
|
211 |
+
x = self.upsample(x)
|
212 |
+
|
213 |
+
return x
|
214 |
+
|
215 |
+
def init_params(self):
|
216 |
+
for m in self.modules():
|
217 |
+
if isinstance(m, nn.Conv2d):
|
218 |
+
# init.kaiming_normal_(m.weight, mode='fan_out')
|
219 |
+
init.normal_(m.weight, std=0.01)
|
220 |
+
# init.xavier_normal_(m.weight)
|
221 |
+
if m.bias is not None:
|
222 |
+
init.constant_(m.bias, 0)
|
223 |
+
elif isinstance(m, nn.ConvTranspose2d):
|
224 |
+
# init.kaiming_normal_(m.weight, mode='fan_out')
|
225 |
+
init.normal_(m.weight, std=0.01)
|
226 |
+
# init.xavier_normal_(m.weight)
|
227 |
+
if m.bias is not None:
|
228 |
+
init.constant_(m.bias, 0)
|
229 |
+
elif isinstance(m, nn.BatchNorm2d): # NN.Batchnorm2d
|
230 |
+
init.constant_(m.weight, 1)
|
231 |
+
init.constant_(m.bias, 0)
|
232 |
+
elif isinstance(m, nn.Linear):
|
233 |
+
init.normal_(m.weight, std=0.01)
|
234 |
+
if m.bias is not None:
|
235 |
+
init.constant_(m.bias, 0)
|
236 |
+
|
237 |
+
|
238 |
+
class AO(nn.Module):
|
239 |
+
# Adaptive output module
|
240 |
+
def __init__(self, inchannels, outchannels, upfactor=2):
|
241 |
+
super(AO, self).__init__()
|
242 |
+
self.inchannels = inchannels
|
243 |
+
self.outchannels = outchannels
|
244 |
+
self.upfactor = upfactor
|
245 |
+
|
246 |
+
self.adapt_conv = nn.Sequential(
|
247 |
+
nn.Conv2d(in_channels=self.inchannels, out_channels=self.inchannels // 2, kernel_size=3, padding=1,
|
248 |
+
stride=1, bias=True), \
|
249 |
+
nn.BatchNorm2d(num_features=self.inchannels // 2), \
|
250 |
+
nn.ReLU(inplace=True), \
|
251 |
+
nn.Conv2d(in_channels=self.inchannels // 2, out_channels=self.outchannels, kernel_size=3, padding=1,
|
252 |
+
stride=1, bias=True), \
|
253 |
+
nn.Upsample(scale_factor=self.upfactor, mode='bilinear', align_corners=True))
|
254 |
+
|
255 |
+
self.init_params()
|
256 |
+
|
257 |
+
def forward(self, x):
|
258 |
+
x = self.adapt_conv(x)
|
259 |
+
return x
|
260 |
+
|
261 |
+
def init_params(self):
|
262 |
+
for m in self.modules():
|
263 |
+
if isinstance(m, nn.Conv2d):
|
264 |
+
# init.kaiming_normal_(m.weight, mode='fan_out')
|
265 |
+
init.normal_(m.weight, std=0.01)
|
266 |
+
# init.xavier_normal_(m.weight)
|
267 |
+
if m.bias is not None:
|
268 |
+
init.constant_(m.bias, 0)
|
269 |
+
elif isinstance(m, nn.ConvTranspose2d):
|
270 |
+
# init.kaiming_normal_(m.weight, mode='fan_out')
|
271 |
+
init.normal_(m.weight, std=0.01)
|
272 |
+
# init.xavier_normal_(m.weight)
|
273 |
+
if m.bias is not None:
|
274 |
+
init.constant_(m.bias, 0)
|
275 |
+
elif isinstance(m, nn.BatchNorm2d): # NN.Batchnorm2d
|
276 |
+
init.constant_(m.weight, 1)
|
277 |
+
init.constant_(m.bias, 0)
|
278 |
+
elif isinstance(m, nn.Linear):
|
279 |
+
init.normal_(m.weight, std=0.01)
|
280 |
+
if m.bias is not None:
|
281 |
+
init.constant_(m.bias, 0)
|
282 |
+
|
283 |
+
|
284 |
+
|
285 |
+
# ==============================================================================================================
|
286 |
+
|
287 |
+
|
288 |
+
class ResidualConv(nn.Module):
|
289 |
+
def __init__(self, inchannels):
|
290 |
+
super(ResidualConv, self).__init__()
|
291 |
+
# NN.BatchNorm2d
|
292 |
+
self.conv = nn.Sequential(
|
293 |
+
# nn.BatchNorm2d(num_features=inchannels),
|
294 |
+
nn.ReLU(inplace=False),
|
295 |
+
# nn.Conv2d(in_channels=inchannels, out_channels=inchannels, kernel_size=3, padding=1, stride=1, groups=inchannels,bias=True),
|
296 |
+
# nn.Conv2d(in_channels=inchannels, out_channels=inchannels, kernel_size=1, padding=0, stride=1, groups=1,bias=True)
|
297 |
+
nn.Conv2d(in_channels=inchannels, out_channels=inchannels / 2, kernel_size=3, padding=1, stride=1,
|
298 |
+
bias=False),
|
299 |
+
nn.BatchNorm2d(num_features=inchannels / 2),
|
300 |
+
nn.ReLU(inplace=False),
|
301 |
+
nn.Conv2d(in_channels=inchannels / 2, out_channels=inchannels, kernel_size=3, padding=1, stride=1,
|
302 |
+
bias=False)
|
303 |
+
)
|
304 |
+
self.init_params()
|
305 |
+
|
306 |
+
def forward(self, x):
|
307 |
+
x = self.conv(x) + x
|
308 |
+
return x
|
309 |
+
|
310 |
+
def init_params(self):
|
311 |
+
for m in self.modules():
|
312 |
+
if isinstance(m, nn.Conv2d):
|
313 |
+
# init.kaiming_normal_(m.weight, mode='fan_out')
|
314 |
+
init.normal_(m.weight, std=0.01)
|
315 |
+
# init.xavier_normal_(m.weight)
|
316 |
+
if m.bias is not None:
|
317 |
+
init.constant_(m.bias, 0)
|
318 |
+
elif isinstance(m, nn.ConvTranspose2d):
|
319 |
+
# init.kaiming_normal_(m.weight, mode='fan_out')
|
320 |
+
init.normal_(m.weight, std=0.01)
|
321 |
+
# init.xavier_normal_(m.weight)
|
322 |
+
if m.bias is not None:
|
323 |
+
init.constant_(m.bias, 0)
|
324 |
+
elif isinstance(m, nn.BatchNorm2d): # NN.BatchNorm2d
|
325 |
+
init.constant_(m.weight, 1)
|
326 |
+
init.constant_(m.bias, 0)
|
327 |
+
elif isinstance(m, nn.Linear):
|
328 |
+
init.normal_(m.weight, std=0.01)
|
329 |
+
if m.bias is not None:
|
330 |
+
init.constant_(m.bias, 0)
|
331 |
+
|
332 |
+
|
333 |
+
class FeatureFusion(nn.Module):
|
334 |
+
def __init__(self, inchannels, outchannels):
|
335 |
+
super(FeatureFusion, self).__init__()
|
336 |
+
self.conv = ResidualConv(inchannels=inchannels)
|
337 |
+
# NN.BatchNorm2d
|
338 |
+
self.up = nn.Sequential(ResidualConv(inchannels=inchannels),
|
339 |
+
nn.ConvTranspose2d(in_channels=inchannels, out_channels=outchannels, kernel_size=3,
|
340 |
+
stride=2, padding=1, output_padding=1),
|
341 |
+
nn.BatchNorm2d(num_features=outchannels),
|
342 |
+
nn.ReLU(inplace=True))
|
343 |
+
|
344 |
+
def forward(self, lowfeat, highfeat):
|
345 |
+
return self.up(highfeat + self.conv(lowfeat))
|
346 |
+
|
347 |
+
def init_params(self):
|
348 |
+
for m in self.modules():
|
349 |
+
if isinstance(m, nn.Conv2d):
|
350 |
+
# init.kaiming_normal_(m.weight, mode='fan_out')
|
351 |
+
init.normal_(m.weight, std=0.01)
|
352 |
+
# init.xavier_normal_(m.weight)
|
353 |
+
if m.bias is not None:
|
354 |
+
init.constant_(m.bias, 0)
|
355 |
+
elif isinstance(m, nn.ConvTranspose2d):
|
356 |
+
# init.kaiming_normal_(m.weight, mode='fan_out')
|
357 |
+
init.normal_(m.weight, std=0.01)
|
358 |
+
# init.xavier_normal_(m.weight)
|
359 |
+
if m.bias is not None:
|
360 |
+
init.constant_(m.bias, 0)
|
361 |
+
elif isinstance(m, nn.BatchNorm2d): # NN.BatchNorm2d
|
362 |
+
init.constant_(m.weight, 1)
|
363 |
+
init.constant_(m.bias, 0)
|
364 |
+
elif isinstance(m, nn.Linear):
|
365 |
+
init.normal_(m.weight, std=0.01)
|
366 |
+
if m.bias is not None:
|
367 |
+
init.constant_(m.bias, 0)
|
368 |
+
|
369 |
+
|
370 |
+
class SenceUnderstand(nn.Module):
|
371 |
+
def __init__(self, channels):
|
372 |
+
super(SenceUnderstand, self).__init__()
|
373 |
+
self.channels = channels
|
374 |
+
self.conv1 = nn.Sequential(nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
|
375 |
+
nn.ReLU(inplace=True))
|
376 |
+
self.pool = nn.AdaptiveAvgPool2d(8)
|
377 |
+
self.fc = nn.Sequential(nn.Linear(512 * 8 * 8, self.channels),
|
378 |
+
nn.ReLU(inplace=True))
|
379 |
+
self.conv2 = nn.Sequential(
|
380 |
+
nn.Conv2d(in_channels=self.channels, out_channels=self.channels, kernel_size=1, padding=0),
|
381 |
+
nn.ReLU(inplace=True))
|
382 |
+
self.initial_params()
|
383 |
+
|
384 |
+
def forward(self, x):
|
385 |
+
n, c, h, w = x.size()
|
386 |
+
x = self.conv1(x)
|
387 |
+
x = self.pool(x)
|
388 |
+
x = x.view(n, -1)
|
389 |
+
x = self.fc(x)
|
390 |
+
x = x.view(n, self.channels, 1, 1)
|
391 |
+
x = self.conv2(x)
|
392 |
+
x = x.repeat(1, 1, h, w)
|
393 |
+
return x
|
394 |
+
|
395 |
+
def initial_params(self, dev=0.01):
|
396 |
+
for m in self.modules():
|
397 |
+
if isinstance(m, nn.Conv2d):
|
398 |
+
# print torch.sum(m.weight)
|
399 |
+
m.weight.data.normal_(0, dev)
|
400 |
+
if m.bias is not None:
|
401 |
+
m.bias.data.fill_(0)
|
402 |
+
elif isinstance(m, nn.ConvTranspose2d):
|
403 |
+
# print torch.sum(m.weight)
|
404 |
+
m.weight.data.normal_(0, dev)
|
405 |
+
if m.bias is not None:
|
406 |
+
m.bias.data.fill_(0)
|
407 |
+
elif isinstance(m, nn.Linear):
|
408 |
+
m.weight.data.normal_(0, dev)
|
409 |
+
|
410 |
+
|
411 |
+
if __name__ == '__main__':
|
412 |
+
net = DepthNet(depth=50, pretrained=True)
|
413 |
+
print(net)
|
414 |
+
inputs = torch.ones(4,3,128,128)
|
415 |
+
out = net(inputs)
|
416 |
+
print(out.size())
|
417 |
+
|
controlnet_aux/leres/pix2pix/__init__.py
ADDED
File without changes
|
controlnet_aux/leres/pix2pix/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (221 Bytes). View file
|
|
controlnet_aux/leres/pix2pix/models/__init__.py
ADDED
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""This package contains modules related to objective functions, optimizations, and network architectures.
|
2 |
+
|
3 |
+
To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.
|
4 |
+
You need to implement the following five functions:
|
5 |
+
-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
|
6 |
+
-- <set_input>: unpack data from dataset and apply preprocessing.
|
7 |
+
-- <forward>: produce intermediate results.
|
8 |
+
-- <optimize_parameters>: calculate loss, gradients, and update network weights.
|
9 |
+
-- <modify_commandline_options>: (optionally) add model-specific options and set default options.
|
10 |
+
|
11 |
+
In the function <__init__>, you need to define four lists:
|
12 |
+
-- self.loss_names (str list): specify the training losses that you want to plot and save.
|
13 |
+
-- self.model_names (str list): define networks used in our training.
|
14 |
+
-- self.visual_names (str list): specify the images that you want to display and save.
|
15 |
+
-- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.
|
16 |
+
|
17 |
+
Now you can use the model class by specifying flag '--model dummy'.
|
18 |
+
See our template model class 'template_model.py' for more details.
|
19 |
+
"""
|
20 |
+
|
21 |
+
import importlib
|
22 |
+
from .base_model import BaseModel
|
23 |
+
|
24 |
+
|
25 |
+
def find_model_using_name(model_name):
|
26 |
+
"""Import the module "models/[model_name]_model.py".
|
27 |
+
|
28 |
+
In the file, the class called DatasetNameModel() will
|
29 |
+
be instantiated. It has to be a subclass of BaseModel,
|
30 |
+
and it is case-insensitive.
|
31 |
+
"""
|
32 |
+
model_filename = "controlnet_aux.leres.pix2pix.models." + model_name + "_model"
|
33 |
+
modellib = importlib.import_module(model_filename)
|
34 |
+
model = None
|
35 |
+
target_model_name = model_name.replace('_', '') + 'model'
|
36 |
+
for name, cls in modellib.__dict__.items():
|
37 |
+
if name.lower() == target_model_name.lower() \
|
38 |
+
and issubclass(cls, BaseModel):
|
39 |
+
model = cls
|
40 |
+
|
41 |
+
if model is None:
|
42 |
+
print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))
|
43 |
+
exit(0)
|
44 |
+
|
45 |
+
return model
|
46 |
+
|
47 |
+
|
48 |
+
def get_option_setter(model_name):
|
49 |
+
"""Return the static method <modify_commandline_options> of the model class."""
|
50 |
+
model_class = find_model_using_name(model_name)
|
51 |
+
return model_class.modify_commandline_options
|
52 |
+
|
53 |
+
|
54 |
+
def create_model(opt):
|
55 |
+
"""Create a model given the option.
|
56 |
+
|
57 |
+
This function warps the class CustomDatasetDataLoader.
|
58 |
+
This is the main interface between this package and 'train.py'/'test.py'
|
59 |
+
|
60 |
+
Example:
|
61 |
+
>>> from models import create_model
|
62 |
+
>>> model = create_model(opt)
|
63 |
+
"""
|
64 |
+
model = find_model_using_name(opt.model)
|
65 |
+
instance = model(opt)
|
66 |
+
print("model [%s] was created" % type(instance).__name__)
|
67 |
+
return instance
|
controlnet_aux/leres/pix2pix/models/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (3.34 kB). View file
|
|
controlnet_aux/leres/pix2pix/models/__pycache__/base_model.cpython-310.pyc
ADDED
Binary file (10.4 kB). View file
|
|
controlnet_aux/leres/pix2pix/models/__pycache__/base_model_hg.cpython-310.pyc
ADDED
Binary file (2.7 kB). View file
|
|
controlnet_aux/leres/pix2pix/models/__pycache__/networks.cpython-310.pyc
ADDED
Binary file (23.5 kB). View file
|
|
controlnet_aux/leres/pix2pix/models/__pycache__/pix2pix4depth_model.cpython-310.pyc
ADDED
Binary file (5.6 kB). View file
|
|
controlnet_aux/leres/pix2pix/models/base_model.py
ADDED
@@ -0,0 +1,244 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gc
|
2 |
+
import os
|
3 |
+
from abc import ABC, abstractmethod
|
4 |
+
from collections import OrderedDict
|
5 |
+
|
6 |
+
import torch
|
7 |
+
|
8 |
+
from ....util import torch_gc
|
9 |
+
from . import networks
|
10 |
+
|
11 |
+
|
12 |
+
class BaseModel(ABC):
|
13 |
+
"""This class is an abstract base class (ABC) for models.
|
14 |
+
To create a subclass, you need to implement the following five functions:
|
15 |
+
-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
|
16 |
+
-- <set_input>: unpack data from dataset and apply preprocessing.
|
17 |
+
-- <forward>: produce intermediate results.
|
18 |
+
-- <optimize_parameters>: calculate losses, gradients, and update network weights.
|
19 |
+
-- <modify_commandline_options>: (optionally) add model-specific options and set default options.
|
20 |
+
"""
|
21 |
+
|
22 |
+
def __init__(self, opt):
|
23 |
+
"""Initialize the BaseModel class.
|
24 |
+
|
25 |
+
Parameters:
|
26 |
+
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
|
27 |
+
|
28 |
+
When creating your custom class, you need to implement your own initialization.
|
29 |
+
In this function, you should first call <BaseModel.__init__(self, opt)>
|
30 |
+
Then, you need to define four lists:
|
31 |
+
-- self.loss_names (str list): specify the training losses that you want to plot and save.
|
32 |
+
-- self.model_names (str list): define networks used in our training.
|
33 |
+
-- self.visual_names (str list): specify the images that you want to display and save.
|
34 |
+
-- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
|
35 |
+
"""
|
36 |
+
self.opt = opt
|
37 |
+
self.gpu_ids = opt.gpu_ids
|
38 |
+
self.isTrain = opt.isTrain
|
39 |
+
self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') # get device name: CPU or GPU
|
40 |
+
self.save_dir = os.path.join(opt.checkpoints_dir, opt.name) # save all the checkpoints to save_dir
|
41 |
+
if opt.preprocess != 'scale_width': # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.
|
42 |
+
torch.backends.cudnn.benchmark = True
|
43 |
+
self.loss_names = []
|
44 |
+
self.model_names = []
|
45 |
+
self.visual_names = []
|
46 |
+
self.optimizers = []
|
47 |
+
self.image_paths = []
|
48 |
+
self.metric = 0 # used for learning rate policy 'plateau'
|
49 |
+
|
50 |
+
@staticmethod
|
51 |
+
def modify_commandline_options(parser, is_train):
|
52 |
+
"""Add new model-specific options, and rewrite default values for existing options.
|
53 |
+
|
54 |
+
Parameters:
|
55 |
+
parser -- original option parser
|
56 |
+
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
|
57 |
+
|
58 |
+
Returns:
|
59 |
+
the modified parser.
|
60 |
+
"""
|
61 |
+
return parser
|
62 |
+
|
63 |
+
@abstractmethod
|
64 |
+
def set_input(self, input):
|
65 |
+
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
|
66 |
+
|
67 |
+
Parameters:
|
68 |
+
input (dict): includes the data itself and its metadata information.
|
69 |
+
"""
|
70 |
+
pass
|
71 |
+
|
72 |
+
@abstractmethod
|
73 |
+
def forward(self):
|
74 |
+
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
|
75 |
+
pass
|
76 |
+
|
77 |
+
@abstractmethod
|
78 |
+
def optimize_parameters(self):
|
79 |
+
"""Calculate losses, gradients, and update network weights; called in every training iteration"""
|
80 |
+
pass
|
81 |
+
|
82 |
+
def setup(self, opt):
|
83 |
+
"""Load and print networks; create schedulers
|
84 |
+
|
85 |
+
Parameters:
|
86 |
+
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
|
87 |
+
"""
|
88 |
+
if self.isTrain:
|
89 |
+
self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
|
90 |
+
if not self.isTrain or opt.continue_train:
|
91 |
+
load_suffix = 'iter_%d' % opt.load_iter if opt.load_iter > 0 else opt.epoch
|
92 |
+
self.load_networks(load_suffix)
|
93 |
+
self.print_networks(opt.verbose)
|
94 |
+
|
95 |
+
def eval(self):
|
96 |
+
"""Make models eval mode during test time"""
|
97 |
+
for name in self.model_names:
|
98 |
+
if isinstance(name, str):
|
99 |
+
net = getattr(self, 'net' + name)
|
100 |
+
net.eval()
|
101 |
+
|
102 |
+
def test(self):
|
103 |
+
"""Forward function used in test time.
|
104 |
+
|
105 |
+
This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
|
106 |
+
It also calls <compute_visuals> to produce additional visualization results
|
107 |
+
"""
|
108 |
+
with torch.no_grad():
|
109 |
+
self.forward()
|
110 |
+
self.compute_visuals()
|
111 |
+
|
112 |
+
def compute_visuals(self):
|
113 |
+
"""Calculate additional output images for visdom and HTML visualization"""
|
114 |
+
pass
|
115 |
+
|
116 |
+
def get_image_paths(self):
|
117 |
+
""" Return image paths that are used to load current data"""
|
118 |
+
return self.image_paths
|
119 |
+
|
120 |
+
def update_learning_rate(self):
|
121 |
+
"""Update learning rates for all the networks; called at the end of every epoch"""
|
122 |
+
old_lr = self.optimizers[0].param_groups[0]['lr']
|
123 |
+
for scheduler in self.schedulers:
|
124 |
+
if self.opt.lr_policy == 'plateau':
|
125 |
+
scheduler.step(self.metric)
|
126 |
+
else:
|
127 |
+
scheduler.step()
|
128 |
+
|
129 |
+
lr = self.optimizers[0].param_groups[0]['lr']
|
130 |
+
print('learning rate %.7f -> %.7f' % (old_lr, lr))
|
131 |
+
|
132 |
+
def get_current_visuals(self):
|
133 |
+
"""Return visualization images. train.py will display these images with visdom, and save the images to a HTML"""
|
134 |
+
visual_ret = OrderedDict()
|
135 |
+
for name in self.visual_names:
|
136 |
+
if isinstance(name, str):
|
137 |
+
visual_ret[name] = getattr(self, name)
|
138 |
+
return visual_ret
|
139 |
+
|
140 |
+
def get_current_losses(self):
|
141 |
+
"""Return traning losses / errors. train.py will print out these errors on console, and save them to a file"""
|
142 |
+
errors_ret = OrderedDict()
|
143 |
+
for name in self.loss_names:
|
144 |
+
if isinstance(name, str):
|
145 |
+
errors_ret[name] = float(getattr(self, 'loss_' + name)) # float(...) works for both scalar tensor and float number
|
146 |
+
return errors_ret
|
147 |
+
|
148 |
+
def save_networks(self, epoch):
|
149 |
+
"""Save all the networks to the disk.
|
150 |
+
|
151 |
+
Parameters:
|
152 |
+
epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
|
153 |
+
"""
|
154 |
+
for name in self.model_names:
|
155 |
+
if isinstance(name, str):
|
156 |
+
save_filename = '%s_net_%s.pth' % (epoch, name)
|
157 |
+
save_path = os.path.join(self.save_dir, save_filename)
|
158 |
+
net = getattr(self, 'net' + name)
|
159 |
+
|
160 |
+
if len(self.gpu_ids) > 0 and torch.cuda.is_available():
|
161 |
+
torch.save(net.module.cpu().state_dict(), save_path)
|
162 |
+
net.cuda(self.gpu_ids[0])
|
163 |
+
else:
|
164 |
+
torch.save(net.cpu().state_dict(), save_path)
|
165 |
+
|
166 |
+
def unload_network(self, name):
|
167 |
+
"""Unload network and gc.
|
168 |
+
"""
|
169 |
+
if isinstance(name, str):
|
170 |
+
net = getattr(self, 'net' + name)
|
171 |
+
del net
|
172 |
+
gc.collect()
|
173 |
+
torch_gc()
|
174 |
+
return None
|
175 |
+
|
176 |
+
def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
|
177 |
+
"""Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""
|
178 |
+
key = keys[i]
|
179 |
+
if i + 1 == len(keys): # at the end, pointing to a parameter/buffer
|
180 |
+
if module.__class__.__name__.startswith('InstanceNorm') and \
|
181 |
+
(key == 'running_mean' or key == 'running_var'):
|
182 |
+
if getattr(module, key) is None:
|
183 |
+
state_dict.pop('.'.join(keys))
|
184 |
+
if module.__class__.__name__.startswith('InstanceNorm') and \
|
185 |
+
(key == 'num_batches_tracked'):
|
186 |
+
state_dict.pop('.'.join(keys))
|
187 |
+
else:
|
188 |
+
self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)
|
189 |
+
|
190 |
+
def load_networks(self, epoch):
|
191 |
+
"""Load all the networks from the disk.
|
192 |
+
|
193 |
+
Parameters:
|
194 |
+
epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
|
195 |
+
"""
|
196 |
+
for name in self.model_names:
|
197 |
+
if isinstance(name, str):
|
198 |
+
load_filename = '%s_net_%s.pth' % (epoch, name)
|
199 |
+
load_path = os.path.join(self.save_dir, load_filename)
|
200 |
+
net = getattr(self, 'net' + name)
|
201 |
+
if isinstance(net, torch.nn.DataParallel):
|
202 |
+
net = net.module
|
203 |
+
# print('Loading depth boost model from %s' % load_path)
|
204 |
+
# if you are using PyTorch newer than 0.4 (e.g., built from
|
205 |
+
# GitHub source), you can remove str() on self.device
|
206 |
+
state_dict = torch.load(load_path, map_location=str(self.device))
|
207 |
+
if hasattr(state_dict, '_metadata'):
|
208 |
+
del state_dict._metadata
|
209 |
+
|
210 |
+
# patch InstanceNorm checkpoints prior to 0.4
|
211 |
+
for key in list(state_dict.keys()): # need to copy keys here because we mutate in loop
|
212 |
+
self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
|
213 |
+
net.load_state_dict(state_dict)
|
214 |
+
|
215 |
+
def print_networks(self, verbose):
|
216 |
+
"""Print the total number of parameters in the network and (if verbose) network architecture
|
217 |
+
|
218 |
+
Parameters:
|
219 |
+
verbose (bool) -- if verbose: print the network architecture
|
220 |
+
"""
|
221 |
+
print('---------- Networks initialized -------------')
|
222 |
+
for name in self.model_names:
|
223 |
+
if isinstance(name, str):
|
224 |
+
net = getattr(self, 'net' + name)
|
225 |
+
num_params = 0
|
226 |
+
for param in net.parameters():
|
227 |
+
num_params += param.numel()
|
228 |
+
if verbose:
|
229 |
+
print(net)
|
230 |
+
print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))
|
231 |
+
print('-----------------------------------------------')
|
232 |
+
|
233 |
+
def set_requires_grad(self, nets, requires_grad=False):
|
234 |
+
"""Set requies_grad=Fasle for all the networks to avoid unnecessary computations
|
235 |
+
Parameters:
|
236 |
+
nets (network list) -- a list of networks
|
237 |
+
requires_grad (bool) -- whether the networks require gradients or not
|
238 |
+
"""
|
239 |
+
if not isinstance(nets, list):
|
240 |
+
nets = [nets]
|
241 |
+
for net in nets:
|
242 |
+
if net is not None:
|
243 |
+
for param in net.parameters():
|
244 |
+
param.requires_grad = requires_grad
|
controlnet_aux/leres/pix2pix/models/base_model_hg.py
ADDED
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import torch
|
3 |
+
|
4 |
+
class BaseModelHG():
|
5 |
+
def name(self):
|
6 |
+
return 'BaseModel'
|
7 |
+
|
8 |
+
def initialize(self, opt):
|
9 |
+
self.opt = opt
|
10 |
+
self.gpu_ids = opt.gpu_ids
|
11 |
+
self.isTrain = opt.isTrain
|
12 |
+
self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor
|
13 |
+
self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
|
14 |
+
|
15 |
+
def set_input(self, input):
|
16 |
+
self.input = input
|
17 |
+
|
18 |
+
def forward(self):
|
19 |
+
pass
|
20 |
+
|
21 |
+
# used in test time, no backprop
|
22 |
+
def test(self):
|
23 |
+
pass
|
24 |
+
|
25 |
+
def get_image_paths(self):
|
26 |
+
pass
|
27 |
+
|
28 |
+
def optimize_parameters(self):
|
29 |
+
pass
|
30 |
+
|
31 |
+
def get_current_visuals(self):
|
32 |
+
return self.input
|
33 |
+
|
34 |
+
def get_current_errors(self):
|
35 |
+
return {}
|
36 |
+
|
37 |
+
def save(self, label):
|
38 |
+
pass
|
39 |
+
|
40 |
+
# helper saving function that can be used by subclasses
|
41 |
+
def save_network(self, network, network_label, epoch_label, gpu_ids):
|
42 |
+
save_filename = '_%s_net_%s.pth' % (epoch_label, network_label)
|
43 |
+
save_path = os.path.join(self.save_dir, save_filename)
|
44 |
+
torch.save(network.cpu().state_dict(), save_path)
|
45 |
+
if len(gpu_ids) and torch.cuda.is_available():
|
46 |
+
network.cuda(device_id=gpu_ids[0])
|
47 |
+
|
48 |
+
# helper loading function that can be used by subclasses
|
49 |
+
def load_network(self, network, network_label, epoch_label):
|
50 |
+
save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
|
51 |
+
save_path = os.path.join(self.save_dir, save_filename)
|
52 |
+
print(save_path)
|
53 |
+
model = torch.load(save_path)
|
54 |
+
return model
|
55 |
+
# network.load_state_dict(torch.load(save_path))
|
56 |
+
|
57 |
+
def update_learning_rate():
|
58 |
+
pass
|
controlnet_aux/leres/pix2pix/models/networks.py
ADDED
@@ -0,0 +1,623 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
from torch.nn import init
|
4 |
+
import functools
|
5 |
+
from torch.optim import lr_scheduler
|
6 |
+
|
7 |
+
|
8 |
+
###############################################################################
|
9 |
+
# Helper Functions
|
10 |
+
###############################################################################
|
11 |
+
|
12 |
+
|
13 |
+
class Identity(nn.Module):
|
14 |
+
def forward(self, x):
|
15 |
+
return x
|
16 |
+
|
17 |
+
|
18 |
+
def get_norm_layer(norm_type='instance'):
|
19 |
+
"""Return a normalization layer
|
20 |
+
|
21 |
+
Parameters:
|
22 |
+
norm_type (str) -- the name of the normalization layer: batch | instance | none
|
23 |
+
|
24 |
+
For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
|
25 |
+
For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
|
26 |
+
"""
|
27 |
+
if norm_type == 'batch':
|
28 |
+
norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
|
29 |
+
elif norm_type == 'instance':
|
30 |
+
norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
|
31 |
+
elif norm_type == 'none':
|
32 |
+
def norm_layer(x): return Identity()
|
33 |
+
else:
|
34 |
+
raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
|
35 |
+
return norm_layer
|
36 |
+
|
37 |
+
|
38 |
+
def get_scheduler(optimizer, opt):
|
39 |
+
"""Return a learning rate scheduler
|
40 |
+
|
41 |
+
Parameters:
|
42 |
+
optimizer -- the optimizer of the network
|
43 |
+
opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions.
|
44 |
+
opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine
|
45 |
+
|
46 |
+
For 'linear', we keep the same learning rate for the first <opt.n_epochs> epochs
|
47 |
+
and linearly decay the rate to zero over the next <opt.n_epochs_decay> epochs.
|
48 |
+
For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
|
49 |
+
See https://pytorch.org/docs/stable/optim.html for more details.
|
50 |
+
"""
|
51 |
+
if opt.lr_policy == 'linear':
|
52 |
+
def lambda_rule(epoch):
|
53 |
+
lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs_decay + 1)
|
54 |
+
return lr_l
|
55 |
+
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
|
56 |
+
elif opt.lr_policy == 'step':
|
57 |
+
scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
|
58 |
+
elif opt.lr_policy == 'plateau':
|
59 |
+
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
|
60 |
+
elif opt.lr_policy == 'cosine':
|
61 |
+
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0)
|
62 |
+
else:
|
63 |
+
return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
|
64 |
+
return scheduler
|
65 |
+
|
66 |
+
|
67 |
+
def init_weights(net, init_type='normal', init_gain=0.02):
|
68 |
+
"""Initialize network weights.
|
69 |
+
|
70 |
+
Parameters:
|
71 |
+
net (network) -- network to be initialized
|
72 |
+
init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
|
73 |
+
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
|
74 |
+
|
75 |
+
We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
|
76 |
+
work better for some applications. Feel free to try yourself.
|
77 |
+
"""
|
78 |
+
def init_func(m): # define the initialization function
|
79 |
+
classname = m.__class__.__name__
|
80 |
+
if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
|
81 |
+
if init_type == 'normal':
|
82 |
+
init.normal_(m.weight.data, 0.0, init_gain)
|
83 |
+
elif init_type == 'xavier':
|
84 |
+
init.xavier_normal_(m.weight.data, gain=init_gain)
|
85 |
+
elif init_type == 'kaiming':
|
86 |
+
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
|
87 |
+
elif init_type == 'orthogonal':
|
88 |
+
init.orthogonal_(m.weight.data, gain=init_gain)
|
89 |
+
else:
|
90 |
+
raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
|
91 |
+
if hasattr(m, 'bias') and m.bias is not None:
|
92 |
+
init.constant_(m.bias.data, 0.0)
|
93 |
+
elif classname.find('BatchNorm2d') != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
|
94 |
+
init.normal_(m.weight.data, 1.0, init_gain)
|
95 |
+
init.constant_(m.bias.data, 0.0)
|
96 |
+
|
97 |
+
# print('initialize network with %s' % init_type)
|
98 |
+
net.apply(init_func) # apply the initialization function <init_func>
|
99 |
+
|
100 |
+
|
101 |
+
def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
|
102 |
+
"""Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
|
103 |
+
Parameters:
|
104 |
+
net (network) -- the network to be initialized
|
105 |
+
init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
|
106 |
+
gain (float) -- scaling factor for normal, xavier and orthogonal.
|
107 |
+
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
|
108 |
+
|
109 |
+
Return an initialized network.
|
110 |
+
"""
|
111 |
+
if len(gpu_ids) > 0:
|
112 |
+
assert(torch.cuda.is_available())
|
113 |
+
net.to(gpu_ids[0])
|
114 |
+
net = torch.nn.DataParallel(net, gpu_ids) # multi-GPUs
|
115 |
+
init_weights(net, init_type, init_gain=init_gain)
|
116 |
+
return net
|
117 |
+
|
118 |
+
|
119 |
+
def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]):
|
120 |
+
"""Create a generator
|
121 |
+
|
122 |
+
Parameters:
|
123 |
+
input_nc (int) -- the number of channels in input images
|
124 |
+
output_nc (int) -- the number of channels in output images
|
125 |
+
ngf (int) -- the number of filters in the last conv layer
|
126 |
+
netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128
|
127 |
+
norm (str) -- the name of normalization layers used in the network: batch | instance | none
|
128 |
+
use_dropout (bool) -- if use dropout layers.
|
129 |
+
init_type (str) -- the name of our initialization method.
|
130 |
+
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
|
131 |
+
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
|
132 |
+
|
133 |
+
Returns a generator
|
134 |
+
|
135 |
+
Our current implementation provides two types of generators:
|
136 |
+
U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)
|
137 |
+
The original U-Net paper: https://arxiv.org/abs/1505.04597
|
138 |
+
|
139 |
+
Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)
|
140 |
+
Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.
|
141 |
+
We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style).
|
142 |
+
|
143 |
+
|
144 |
+
The generator has been initialized by <init_net>. It uses RELU for non-linearity.
|
145 |
+
"""
|
146 |
+
net = None
|
147 |
+
norm_layer = get_norm_layer(norm_type=norm)
|
148 |
+
|
149 |
+
if netG == 'resnet_9blocks':
|
150 |
+
net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
|
151 |
+
elif netG == 'resnet_6blocks':
|
152 |
+
net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
|
153 |
+
elif netG == 'resnet_12blocks':
|
154 |
+
net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=12)
|
155 |
+
elif netG == 'unet_128':
|
156 |
+
net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
|
157 |
+
elif netG == 'unet_256':
|
158 |
+
net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
|
159 |
+
elif netG == 'unet_672':
|
160 |
+
net = UnetGenerator(input_nc, output_nc, 5, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
|
161 |
+
elif netG == 'unet_960':
|
162 |
+
net = UnetGenerator(input_nc, output_nc, 6, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
|
163 |
+
elif netG == 'unet_1024':
|
164 |
+
net = UnetGenerator(input_nc, output_nc, 10, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
|
165 |
+
else:
|
166 |
+
raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
|
167 |
+
return init_net(net, init_type, init_gain, gpu_ids)
|
168 |
+
|
169 |
+
|
170 |
+
def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, gpu_ids=[]):
|
171 |
+
"""Create a discriminator
|
172 |
+
|
173 |
+
Parameters:
|
174 |
+
input_nc (int) -- the number of channels in input images
|
175 |
+
ndf (int) -- the number of filters in the first conv layer
|
176 |
+
netD (str) -- the architecture's name: basic | n_layers | pixel
|
177 |
+
n_layers_D (int) -- the number of conv layers in the discriminator; effective when netD=='n_layers'
|
178 |
+
norm (str) -- the type of normalization layers used in the network.
|
179 |
+
init_type (str) -- the name of the initialization method.
|
180 |
+
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
|
181 |
+
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
|
182 |
+
|
183 |
+
Returns a discriminator
|
184 |
+
|
185 |
+
Our current implementation provides three types of discriminators:
|
186 |
+
[basic]: 'PatchGAN' classifier described in the original pix2pix paper.
|
187 |
+
It can classify whether 70×70 overlapping patches are real or fake.
|
188 |
+
Such a patch-level discriminator architecture has fewer parameters
|
189 |
+
than a full-image discriminator and can work on arbitrarily-sized images
|
190 |
+
in a fully convolutional fashion.
|
191 |
+
|
192 |
+
[n_layers]: With this mode, you can specify the number of conv layers in the discriminator
|
193 |
+
with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).)
|
194 |
+
|
195 |
+
[pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not.
|
196 |
+
It encourages greater color diversity but has no effect on spatial statistics.
|
197 |
+
|
198 |
+
The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity.
|
199 |
+
"""
|
200 |
+
net = None
|
201 |
+
norm_layer = get_norm_layer(norm_type=norm)
|
202 |
+
|
203 |
+
if netD == 'basic': # default PatchGAN classifier
|
204 |
+
net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer)
|
205 |
+
elif netD == 'n_layers': # more options
|
206 |
+
net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer)
|
207 |
+
elif netD == 'pixel': # classify if each pixel is real or fake
|
208 |
+
net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)
|
209 |
+
else:
|
210 |
+
raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD)
|
211 |
+
return init_net(net, init_type, init_gain, gpu_ids)
|
212 |
+
|
213 |
+
|
214 |
+
##############################################################################
|
215 |
+
# Classes
|
216 |
+
##############################################################################
|
217 |
+
class GANLoss(nn.Module):
|
218 |
+
"""Define different GAN objectives.
|
219 |
+
|
220 |
+
The GANLoss class abstracts away the need to create the target label tensor
|
221 |
+
that has the same size as the input.
|
222 |
+
"""
|
223 |
+
|
224 |
+
def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
|
225 |
+
""" Initialize the GANLoss class.
|
226 |
+
|
227 |
+
Parameters:
|
228 |
+
gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
|
229 |
+
target_real_label (bool) - - label for a real image
|
230 |
+
target_fake_label (bool) - - label of a fake image
|
231 |
+
|
232 |
+
Note: Do not use sigmoid as the last layer of Discriminator.
|
233 |
+
LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
|
234 |
+
"""
|
235 |
+
super(GANLoss, self).__init__()
|
236 |
+
self.register_buffer('real_label', torch.tensor(target_real_label))
|
237 |
+
self.register_buffer('fake_label', torch.tensor(target_fake_label))
|
238 |
+
self.gan_mode = gan_mode
|
239 |
+
if gan_mode == 'lsgan':
|
240 |
+
self.loss = nn.MSELoss()
|
241 |
+
elif gan_mode == 'vanilla':
|
242 |
+
self.loss = nn.BCEWithLogitsLoss()
|
243 |
+
elif gan_mode in ['wgangp']:
|
244 |
+
self.loss = None
|
245 |
+
else:
|
246 |
+
raise NotImplementedError('gan mode %s not implemented' % gan_mode)
|
247 |
+
|
248 |
+
def get_target_tensor(self, prediction, target_is_real):
|
249 |
+
"""Create label tensors with the same size as the input.
|
250 |
+
|
251 |
+
Parameters:
|
252 |
+
prediction (tensor) - - tpyically the prediction from a discriminator
|
253 |
+
target_is_real (bool) - - if the ground truth label is for real images or fake images
|
254 |
+
|
255 |
+
Returns:
|
256 |
+
A label tensor filled with ground truth label, and with the size of the input
|
257 |
+
"""
|
258 |
+
|
259 |
+
if target_is_real:
|
260 |
+
target_tensor = self.real_label
|
261 |
+
else:
|
262 |
+
target_tensor = self.fake_label
|
263 |
+
return target_tensor.expand_as(prediction)
|
264 |
+
|
265 |
+
def __call__(self, prediction, target_is_real):
|
266 |
+
"""Calculate loss given Discriminator's output and grount truth labels.
|
267 |
+
|
268 |
+
Parameters:
|
269 |
+
prediction (tensor) - - tpyically the prediction output from a discriminator
|
270 |
+
target_is_real (bool) - - if the ground truth label is for real images or fake images
|
271 |
+
|
272 |
+
Returns:
|
273 |
+
the calculated loss.
|
274 |
+
"""
|
275 |
+
if self.gan_mode in ['lsgan', 'vanilla']:
|
276 |
+
target_tensor = self.get_target_tensor(prediction, target_is_real)
|
277 |
+
loss = self.loss(prediction, target_tensor)
|
278 |
+
elif self.gan_mode == 'wgangp':
|
279 |
+
if target_is_real:
|
280 |
+
loss = -prediction.mean()
|
281 |
+
else:
|
282 |
+
loss = prediction.mean()
|
283 |
+
return loss
|
284 |
+
|
285 |
+
|
286 |
+
def cal_gradient_penalty(netD, real_data, fake_data, device, type='mixed', constant=1.0, lambda_gp=10.0):
|
287 |
+
"""Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028
|
288 |
+
|
289 |
+
Arguments:
|
290 |
+
netD (network) -- discriminator network
|
291 |
+
real_data (tensor array) -- real images
|
292 |
+
fake_data (tensor array) -- generated images from the generator
|
293 |
+
device (str) -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
|
294 |
+
type (str) -- if we mix real and fake data or not [real | fake | mixed].
|
295 |
+
constant (float) -- the constant used in formula ( ||gradient||_2 - constant)^2
|
296 |
+
lambda_gp (float) -- weight for this loss
|
297 |
+
|
298 |
+
Returns the gradient penalty loss
|
299 |
+
"""
|
300 |
+
if lambda_gp > 0.0:
|
301 |
+
if type == 'real': # either use real images, fake images, or a linear interpolation of two.
|
302 |
+
interpolatesv = real_data
|
303 |
+
elif type == 'fake':
|
304 |
+
interpolatesv = fake_data
|
305 |
+
elif type == 'mixed':
|
306 |
+
alpha = torch.rand(real_data.shape[0], 1, device=device)
|
307 |
+
alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape)
|
308 |
+
interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)
|
309 |
+
else:
|
310 |
+
raise NotImplementedError('{} not implemented'.format(type))
|
311 |
+
interpolatesv.requires_grad_(True)
|
312 |
+
disc_interpolates = netD(interpolatesv)
|
313 |
+
gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv,
|
314 |
+
grad_outputs=torch.ones(disc_interpolates.size()).to(device),
|
315 |
+
create_graph=True, retain_graph=True, only_inputs=True)
|
316 |
+
gradients = gradients[0].view(real_data.size(0), -1) # flat the data
|
317 |
+
gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp # added eps
|
318 |
+
return gradient_penalty, gradients
|
319 |
+
else:
|
320 |
+
return 0.0, None
|
321 |
+
|
322 |
+
|
323 |
+
class ResnetGenerator(nn.Module):
|
324 |
+
"""Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
|
325 |
+
|
326 |
+
We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
|
327 |
+
"""
|
328 |
+
|
329 |
+
def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
|
330 |
+
"""Construct a Resnet-based generator
|
331 |
+
|
332 |
+
Parameters:
|
333 |
+
input_nc (int) -- the number of channels in input images
|
334 |
+
output_nc (int) -- the number of channels in output images
|
335 |
+
ngf (int) -- the number of filters in the last conv layer
|
336 |
+
norm_layer -- normalization layer
|
337 |
+
use_dropout (bool) -- if use dropout layers
|
338 |
+
n_blocks (int) -- the number of ResNet blocks
|
339 |
+
padding_type (str) -- the name of padding layer in conv layers: reflect | replicate | zero
|
340 |
+
"""
|
341 |
+
assert(n_blocks >= 0)
|
342 |
+
super(ResnetGenerator, self).__init__()
|
343 |
+
if type(norm_layer) == functools.partial:
|
344 |
+
use_bias = norm_layer.func == nn.InstanceNorm2d
|
345 |
+
else:
|
346 |
+
use_bias = norm_layer == nn.InstanceNorm2d
|
347 |
+
|
348 |
+
model = [nn.ReflectionPad2d(3),
|
349 |
+
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
|
350 |
+
norm_layer(ngf),
|
351 |
+
nn.ReLU(True)]
|
352 |
+
|
353 |
+
n_downsampling = 2
|
354 |
+
for i in range(n_downsampling): # add downsampling layers
|
355 |
+
mult = 2 ** i
|
356 |
+
model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
|
357 |
+
norm_layer(ngf * mult * 2),
|
358 |
+
nn.ReLU(True)]
|
359 |
+
|
360 |
+
mult = 2 ** n_downsampling
|
361 |
+
for i in range(n_blocks): # add ResNet blocks
|
362 |
+
|
363 |
+
model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
|
364 |
+
|
365 |
+
for i in range(n_downsampling): # add upsampling layers
|
366 |
+
mult = 2 ** (n_downsampling - i)
|
367 |
+
model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
|
368 |
+
kernel_size=3, stride=2,
|
369 |
+
padding=1, output_padding=1,
|
370 |
+
bias=use_bias),
|
371 |
+
norm_layer(int(ngf * mult / 2)),
|
372 |
+
nn.ReLU(True)]
|
373 |
+
model += [nn.ReflectionPad2d(3)]
|
374 |
+
model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
|
375 |
+
model += [nn.Tanh()]
|
376 |
+
|
377 |
+
self.model = nn.Sequential(*model)
|
378 |
+
|
379 |
+
def forward(self, input):
|
380 |
+
"""Standard forward"""
|
381 |
+
return self.model(input)
|
382 |
+
|
383 |
+
|
384 |
+
class ResnetBlock(nn.Module):
|
385 |
+
"""Define a Resnet block"""
|
386 |
+
|
387 |
+
def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
|
388 |
+
"""Initialize the Resnet block
|
389 |
+
|
390 |
+
A resnet block is a conv block with skip connections
|
391 |
+
We construct a conv block with build_conv_block function,
|
392 |
+
and implement skip connections in <forward> function.
|
393 |
+
Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
|
394 |
+
"""
|
395 |
+
super(ResnetBlock, self).__init__()
|
396 |
+
self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
|
397 |
+
|
398 |
+
def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
|
399 |
+
"""Construct a convolutional block.
|
400 |
+
|
401 |
+
Parameters:
|
402 |
+
dim (int) -- the number of channels in the conv layer.
|
403 |
+
padding_type (str) -- the name of padding layer: reflect | replicate | zero
|
404 |
+
norm_layer -- normalization layer
|
405 |
+
use_dropout (bool) -- if use dropout layers.
|
406 |
+
use_bias (bool) -- if the conv layer uses bias or not
|
407 |
+
|
408 |
+
Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
|
409 |
+
"""
|
410 |
+
conv_block = []
|
411 |
+
p = 0
|
412 |
+
if padding_type == 'reflect':
|
413 |
+
conv_block += [nn.ReflectionPad2d(1)]
|
414 |
+
elif padding_type == 'replicate':
|
415 |
+
conv_block += [nn.ReplicationPad2d(1)]
|
416 |
+
elif padding_type == 'zero':
|
417 |
+
p = 1
|
418 |
+
else:
|
419 |
+
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
|
420 |
+
|
421 |
+
conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
|
422 |
+
if use_dropout:
|
423 |
+
conv_block += [nn.Dropout(0.5)]
|
424 |
+
|
425 |
+
p = 0
|
426 |
+
if padding_type == 'reflect':
|
427 |
+
conv_block += [nn.ReflectionPad2d(1)]
|
428 |
+
elif padding_type == 'replicate':
|
429 |
+
conv_block += [nn.ReplicationPad2d(1)]
|
430 |
+
elif padding_type == 'zero':
|
431 |
+
p = 1
|
432 |
+
else:
|
433 |
+
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
|
434 |
+
conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
|
435 |
+
|
436 |
+
return nn.Sequential(*conv_block)
|
437 |
+
|
438 |
+
def forward(self, x):
|
439 |
+
"""Forward function (with skip connections)"""
|
440 |
+
out = x + self.conv_block(x) # add skip connections
|
441 |
+
return out
|
442 |
+
|
443 |
+
|
444 |
+
class UnetGenerator(nn.Module):
|
445 |
+
"""Create a Unet-based generator"""
|
446 |
+
|
447 |
+
def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):
|
448 |
+
"""Construct a Unet generator
|
449 |
+
Parameters:
|
450 |
+
input_nc (int) -- the number of channels in input images
|
451 |
+
output_nc (int) -- the number of channels in output images
|
452 |
+
num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
|
453 |
+
image of size 128x128 will become of size 1x1 # at the bottleneck
|
454 |
+
ngf (int) -- the number of filters in the last conv layer
|
455 |
+
norm_layer -- normalization layer
|
456 |
+
|
457 |
+
We construct the U-Net from the innermost layer to the outermost layer.
|
458 |
+
It is a recursive process.
|
459 |
+
"""
|
460 |
+
super(UnetGenerator, self).__init__()
|
461 |
+
# construct unet structure
|
462 |
+
unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True) # add the innermost layer
|
463 |
+
for i in range(num_downs - 5): # add intermediate layers with ngf * 8 filters
|
464 |
+
unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
|
465 |
+
# gradually reduce the number of filters from ngf * 8 to ngf
|
466 |
+
unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
|
467 |
+
unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
|
468 |
+
unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
|
469 |
+
self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer) # add the outermost layer
|
470 |
+
|
471 |
+
def forward(self, input):
|
472 |
+
"""Standard forward"""
|
473 |
+
return self.model(input)
|
474 |
+
|
475 |
+
|
476 |
+
class UnetSkipConnectionBlock(nn.Module):
|
477 |
+
"""Defines the Unet submodule with skip connection.
|
478 |
+
X -------------------identity----------------------
|
479 |
+
|-- downsampling -- |submodule| -- upsampling --|
|
480 |
+
"""
|
481 |
+
|
482 |
+
def __init__(self, outer_nc, inner_nc, input_nc=None,
|
483 |
+
submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
|
484 |
+
"""Construct a Unet submodule with skip connections.
|
485 |
+
|
486 |
+
Parameters:
|
487 |
+
outer_nc (int) -- the number of filters in the outer conv layer
|
488 |
+
inner_nc (int) -- the number of filters in the inner conv layer
|
489 |
+
input_nc (int) -- the number of channels in input images/features
|
490 |
+
submodule (UnetSkipConnectionBlock) -- previously defined submodules
|
491 |
+
outermost (bool) -- if this module is the outermost module
|
492 |
+
innermost (bool) -- if this module is the innermost module
|
493 |
+
norm_layer -- normalization layer
|
494 |
+
use_dropout (bool) -- if use dropout layers.
|
495 |
+
"""
|
496 |
+
super(UnetSkipConnectionBlock, self).__init__()
|
497 |
+
self.outermost = outermost
|
498 |
+
if type(norm_layer) == functools.partial:
|
499 |
+
use_bias = norm_layer.func == nn.InstanceNorm2d
|
500 |
+
else:
|
501 |
+
use_bias = norm_layer == nn.InstanceNorm2d
|
502 |
+
if input_nc is None:
|
503 |
+
input_nc = outer_nc
|
504 |
+
downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
|
505 |
+
stride=2, padding=1, bias=use_bias)
|
506 |
+
downrelu = nn.LeakyReLU(0.2, True)
|
507 |
+
downnorm = norm_layer(inner_nc)
|
508 |
+
uprelu = nn.ReLU(True)
|
509 |
+
upnorm = norm_layer(outer_nc)
|
510 |
+
|
511 |
+
if outermost:
|
512 |
+
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
|
513 |
+
kernel_size=4, stride=2,
|
514 |
+
padding=1)
|
515 |
+
down = [downconv]
|
516 |
+
up = [uprelu, upconv, nn.Tanh()]
|
517 |
+
model = down + [submodule] + up
|
518 |
+
elif innermost:
|
519 |
+
upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
|
520 |
+
kernel_size=4, stride=2,
|
521 |
+
padding=1, bias=use_bias)
|
522 |
+
down = [downrelu, downconv]
|
523 |
+
up = [uprelu, upconv, upnorm]
|
524 |
+
model = down + up
|
525 |
+
else:
|
526 |
+
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
|
527 |
+
kernel_size=4, stride=2,
|
528 |
+
padding=1, bias=use_bias)
|
529 |
+
down = [downrelu, downconv, downnorm]
|
530 |
+
up = [uprelu, upconv, upnorm]
|
531 |
+
|
532 |
+
if use_dropout:
|
533 |
+
model = down + [submodule] + up + [nn.Dropout(0.5)]
|
534 |
+
else:
|
535 |
+
model = down + [submodule] + up
|
536 |
+
|
537 |
+
self.model = nn.Sequential(*model)
|
538 |
+
|
539 |
+
def forward(self, x):
|
540 |
+
if self.outermost:
|
541 |
+
return self.model(x)
|
542 |
+
else: # add skip connections
|
543 |
+
return torch.cat([x, self.model(x)], 1)
|
544 |
+
|
545 |
+
|
546 |
+
class NLayerDiscriminator(nn.Module):
|
547 |
+
"""Defines a PatchGAN discriminator"""
|
548 |
+
|
549 |
+
def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
|
550 |
+
"""Construct a PatchGAN discriminator
|
551 |
+
|
552 |
+
Parameters:
|
553 |
+
input_nc (int) -- the number of channels in input images
|
554 |
+
ndf (int) -- the number of filters in the last conv layer
|
555 |
+
n_layers (int) -- the number of conv layers in the discriminator
|
556 |
+
norm_layer -- normalization layer
|
557 |
+
"""
|
558 |
+
super(NLayerDiscriminator, self).__init__()
|
559 |
+
if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters
|
560 |
+
use_bias = norm_layer.func == nn.InstanceNorm2d
|
561 |
+
else:
|
562 |
+
use_bias = norm_layer == nn.InstanceNorm2d
|
563 |
+
|
564 |
+
kw = 4
|
565 |
+
padw = 1
|
566 |
+
sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
|
567 |
+
nf_mult = 1
|
568 |
+
nf_mult_prev = 1
|
569 |
+
for n in range(1, n_layers): # gradually increase the number of filters
|
570 |
+
nf_mult_prev = nf_mult
|
571 |
+
nf_mult = min(2 ** n, 8)
|
572 |
+
sequence += [
|
573 |
+
nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
|
574 |
+
norm_layer(ndf * nf_mult),
|
575 |
+
nn.LeakyReLU(0.2, True)
|
576 |
+
]
|
577 |
+
|
578 |
+
nf_mult_prev = nf_mult
|
579 |
+
nf_mult = min(2 ** n_layers, 8)
|
580 |
+
sequence += [
|
581 |
+
nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
|
582 |
+
norm_layer(ndf * nf_mult),
|
583 |
+
nn.LeakyReLU(0.2, True)
|
584 |
+
]
|
585 |
+
|
586 |
+
sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)] # output 1 channel prediction map
|
587 |
+
self.model = nn.Sequential(*sequence)
|
588 |
+
|
589 |
+
def forward(self, input):
|
590 |
+
"""Standard forward."""
|
591 |
+
return self.model(input)
|
592 |
+
|
593 |
+
|
594 |
+
class PixelDiscriminator(nn.Module):
|
595 |
+
"""Defines a 1x1 PatchGAN discriminator (pixelGAN)"""
|
596 |
+
|
597 |
+
def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d):
|
598 |
+
"""Construct a 1x1 PatchGAN discriminator
|
599 |
+
|
600 |
+
Parameters:
|
601 |
+
input_nc (int) -- the number of channels in input images
|
602 |
+
ndf (int) -- the number of filters in the last conv layer
|
603 |
+
norm_layer -- normalization layer
|
604 |
+
"""
|
605 |
+
super(PixelDiscriminator, self).__init__()
|
606 |
+
if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters
|
607 |
+
use_bias = norm_layer.func == nn.InstanceNorm2d
|
608 |
+
else:
|
609 |
+
use_bias = norm_layer == nn.InstanceNorm2d
|
610 |
+
|
611 |
+
self.net = [
|
612 |
+
nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),
|
613 |
+
nn.LeakyReLU(0.2, True),
|
614 |
+
nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),
|
615 |
+
norm_layer(ndf * 2),
|
616 |
+
nn.LeakyReLU(0.2, True),
|
617 |
+
nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]
|
618 |
+
|
619 |
+
self.net = nn.Sequential(*self.net)
|
620 |
+
|
621 |
+
def forward(self, input):
|
622 |
+
"""Standard forward."""
|
623 |
+
return self.net(input)
|
controlnet_aux/leres/pix2pix/models/pix2pix4depth_model.py
ADDED
@@ -0,0 +1,155 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from .base_model import BaseModel
|
3 |
+
from . import networks
|
4 |
+
|
5 |
+
|
6 |
+
class Pix2Pix4DepthModel(BaseModel):
|
7 |
+
""" This class implements the pix2pix model, for learning a mapping from input images to output images given paired data.
|
8 |
+
|
9 |
+
The model training requires '--dataset_mode aligned' dataset.
|
10 |
+
By default, it uses a '--netG unet256' U-Net generator,
|
11 |
+
a '--netD basic' discriminator (PatchGAN),
|
12 |
+
and a '--gan_mode' vanilla GAN loss (the cross-entropy objective used in the orignal GAN paper).
|
13 |
+
|
14 |
+
pix2pix paper: https://arxiv.org/pdf/1611.07004.pdf
|
15 |
+
"""
|
16 |
+
@staticmethod
|
17 |
+
def modify_commandline_options(parser, is_train=True):
|
18 |
+
"""Add new dataset-specific options, and rewrite default values for existing options.
|
19 |
+
|
20 |
+
Parameters:
|
21 |
+
parser -- original option parser
|
22 |
+
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
|
23 |
+
|
24 |
+
Returns:
|
25 |
+
the modified parser.
|
26 |
+
|
27 |
+
For pix2pix, we do not use image buffer
|
28 |
+
The training objective is: GAN Loss + lambda_L1 * ||G(A)-B||_1
|
29 |
+
By default, we use vanilla GAN loss, UNet with batchnorm, and aligned datasets.
|
30 |
+
"""
|
31 |
+
# changing the default values to match the pix2pix paper (https://phillipi.github.io/pix2pix/)
|
32 |
+
parser.set_defaults(input_nc=2,output_nc=1,norm='none', netG='unet_1024', dataset_mode='depthmerge')
|
33 |
+
if is_train:
|
34 |
+
parser.set_defaults(pool_size=0, gan_mode='vanilla',)
|
35 |
+
parser.add_argument('--lambda_L1', type=float, default=1000, help='weight for L1 loss')
|
36 |
+
return parser
|
37 |
+
|
38 |
+
def __init__(self, opt):
|
39 |
+
"""Initialize the pix2pix class.
|
40 |
+
|
41 |
+
Parameters:
|
42 |
+
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
|
43 |
+
"""
|
44 |
+
BaseModel.__init__(self, opt)
|
45 |
+
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
|
46 |
+
|
47 |
+
self.loss_names = ['G_GAN', 'G_L1', 'D_real', 'D_fake']
|
48 |
+
# self.loss_names = ['G_L1']
|
49 |
+
|
50 |
+
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
|
51 |
+
if self.isTrain:
|
52 |
+
self.visual_names = ['outer','inner', 'fake_B', 'real_B']
|
53 |
+
else:
|
54 |
+
self.visual_names = ['fake_B']
|
55 |
+
|
56 |
+
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
|
57 |
+
if self.isTrain:
|
58 |
+
self.model_names = ['G','D']
|
59 |
+
else: # during test time, only load G
|
60 |
+
self.model_names = ['G']
|
61 |
+
|
62 |
+
# define networks (both generator and discriminator)
|
63 |
+
self.netG = networks.define_G(opt.input_nc, opt.output_nc, 64, 'unet_1024', 'none',
|
64 |
+
False, 'normal', 0.02, self.gpu_ids)
|
65 |
+
|
66 |
+
if self.isTrain: # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
|
67 |
+
self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
|
68 |
+
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
|
69 |
+
|
70 |
+
if self.isTrain:
|
71 |
+
# define loss functions
|
72 |
+
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
|
73 |
+
self.criterionL1 = torch.nn.L1Loss()
|
74 |
+
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
|
75 |
+
self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=1e-4, betas=(opt.beta1, 0.999))
|
76 |
+
self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=2e-06, betas=(opt.beta1, 0.999))
|
77 |
+
self.optimizers.append(self.optimizer_G)
|
78 |
+
self.optimizers.append(self.optimizer_D)
|
79 |
+
|
80 |
+
def set_input_train(self, input):
|
81 |
+
self.outer = input['data_outer'].to(self.device)
|
82 |
+
self.outer = torch.nn.functional.interpolate(self.outer,(1024,1024),mode='bilinear',align_corners=False)
|
83 |
+
|
84 |
+
self.inner = input['data_inner'].to(self.device)
|
85 |
+
self.inner = torch.nn.functional.interpolate(self.inner,(1024,1024),mode='bilinear',align_corners=False)
|
86 |
+
|
87 |
+
self.image_paths = input['image_path']
|
88 |
+
|
89 |
+
if self.isTrain:
|
90 |
+
self.gtfake = input['data_gtfake'].to(self.device)
|
91 |
+
self.gtfake = torch.nn.functional.interpolate(self.gtfake, (1024, 1024), mode='bilinear', align_corners=False)
|
92 |
+
self.real_B = self.gtfake
|
93 |
+
|
94 |
+
self.real_A = torch.cat((self.outer, self.inner), 1)
|
95 |
+
|
96 |
+
def set_input(self, outer, inner):
|
97 |
+
inner = torch.from_numpy(inner).unsqueeze(0).unsqueeze(0)
|
98 |
+
outer = torch.from_numpy(outer).unsqueeze(0).unsqueeze(0)
|
99 |
+
|
100 |
+
inner = (inner - torch.min(inner))/(torch.max(inner)-torch.min(inner))
|
101 |
+
outer = (outer - torch.min(outer))/(torch.max(outer)-torch.min(outer))
|
102 |
+
|
103 |
+
inner = self.normalize(inner)
|
104 |
+
outer = self.normalize(outer)
|
105 |
+
|
106 |
+
self.real_A = torch.cat((outer, inner), 1).to(self.device)
|
107 |
+
|
108 |
+
|
109 |
+
def normalize(self, input):
|
110 |
+
input = input * 2
|
111 |
+
input = input - 1
|
112 |
+
return input
|
113 |
+
|
114 |
+
def forward(self):
|
115 |
+
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
|
116 |
+
self.fake_B = self.netG(self.real_A) # G(A)
|
117 |
+
|
118 |
+
def backward_D(self):
|
119 |
+
"""Calculate GAN loss for the discriminator"""
|
120 |
+
# Fake; stop backprop to the generator by detaching fake_B
|
121 |
+
fake_AB = torch.cat((self.real_A, self.fake_B), 1) # we use conditional GANs; we need to feed both input and output to the discriminator
|
122 |
+
pred_fake = self.netD(fake_AB.detach())
|
123 |
+
self.loss_D_fake = self.criterionGAN(pred_fake, False)
|
124 |
+
# Real
|
125 |
+
real_AB = torch.cat((self.real_A, self.real_B), 1)
|
126 |
+
pred_real = self.netD(real_AB)
|
127 |
+
self.loss_D_real = self.criterionGAN(pred_real, True)
|
128 |
+
# combine loss and calculate gradients
|
129 |
+
self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
|
130 |
+
self.loss_D.backward()
|
131 |
+
|
132 |
+
def backward_G(self):
|
133 |
+
"""Calculate GAN and L1 loss for the generator"""
|
134 |
+
# First, G(A) should fake the discriminator
|
135 |
+
fake_AB = torch.cat((self.real_A, self.fake_B), 1)
|
136 |
+
pred_fake = self.netD(fake_AB)
|
137 |
+
self.loss_G_GAN = self.criterionGAN(pred_fake, True)
|
138 |
+
# Second, G(A) = B
|
139 |
+
self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_L1
|
140 |
+
# combine loss and calculate gradients
|
141 |
+
self.loss_G = self.loss_G_L1 + self.loss_G_GAN
|
142 |
+
self.loss_G.backward()
|
143 |
+
|
144 |
+
def optimize_parameters(self):
|
145 |
+
self.forward() # compute fake images: G(A)
|
146 |
+
# update D
|
147 |
+
self.set_requires_grad(self.netD, True) # enable backprop for D
|
148 |
+
self.optimizer_D.zero_grad() # set D's gradients to zero
|
149 |
+
self.backward_D() # calculate gradients for D
|
150 |
+
self.optimizer_D.step() # update D's weights
|
151 |
+
# update G
|
152 |
+
self.set_requires_grad(self.netD, False) # D requires no gradients when optimizing G
|
153 |
+
self.optimizer_G.zero_grad() # set G's gradients to zero
|
154 |
+
self.backward_G() # calculate graidents for G
|
155 |
+
self.optimizer_G.step() # udpate G's weights
|