broyang commited on
Commit
43d4ed2
·
verified ·
1 Parent(s): 634e9e3

Upload 299 files

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. controlnet_aux-0.0.7.dist-info/INSTALLER +1 -0
  2. controlnet_aux-0.0.7.dist-info/LICENSE.txt +201 -0
  3. controlnet_aux-0.0.7.dist-info/METADATA +148 -0
  4. controlnet_aux-0.0.7.dist-info/RECORD +298 -0
  5. controlnet_aux-0.0.7.dist-info/REQUESTED +0 -0
  6. controlnet_aux-0.0.7.dist-info/WHEEL +5 -0
  7. controlnet_aux-0.0.7.dist-info/top_level.txt +1 -0
  8. controlnet_aux/.DS_Store +0 -0
  9. controlnet_aux/__init__.py +18 -0
  10. controlnet_aux/__pycache__/__init__.cpython-310.pyc +0 -0
  11. controlnet_aux/__pycache__/processor.cpython-310.pyc +0 -0
  12. controlnet_aux/__pycache__/util.cpython-310.pyc +0 -0
  13. controlnet_aux/canny/__init__.py +36 -0
  14. controlnet_aux/canny/__pycache__/__init__.cpython-310.pyc +0 -0
  15. controlnet_aux/dwpose/__init__.py +91 -0
  16. controlnet_aux/dwpose/__pycache__/__init__.cpython-310.pyc +0 -0
  17. controlnet_aux/dwpose/__pycache__/util.cpython-310.pyc +0 -0
  18. controlnet_aux/dwpose/__pycache__/wholebody.cpython-310.pyc +0 -0
  19. controlnet_aux/dwpose/util.py +303 -0
  20. controlnet_aux/dwpose/wholebody.py +121 -0
  21. controlnet_aux/hed/__init__.py +129 -0
  22. controlnet_aux/hed/__pycache__/__init__.cpython-310.pyc +0 -0
  23. controlnet_aux/leres/__init__.py +118 -0
  24. controlnet_aux/leres/__pycache__/__init__.cpython-310.pyc +0 -0
  25. controlnet_aux/leres/leres/Resnet.py +199 -0
  26. controlnet_aux/leres/leres/Resnext_torch.py +237 -0
  27. controlnet_aux/leres/leres/__init__.py +0 -0
  28. controlnet_aux/leres/leres/__pycache__/Resnet.cpython-310.pyc +0 -0
  29. controlnet_aux/leres/leres/__pycache__/Resnext_torch.cpython-310.pyc +0 -0
  30. controlnet_aux/leres/leres/__pycache__/__init__.cpython-310.pyc +0 -0
  31. controlnet_aux/leres/leres/__pycache__/depthmap.cpython-310.pyc +0 -0
  32. controlnet_aux/leres/leres/__pycache__/multi_depth_model_woauxi.cpython-310.pyc +0 -0
  33. controlnet_aux/leres/leres/__pycache__/net_tools.cpython-310.pyc +0 -0
  34. controlnet_aux/leres/leres/__pycache__/network_auxi.cpython-310.pyc +0 -0
  35. controlnet_aux/leres/leres/depthmap.py +548 -0
  36. controlnet_aux/leres/leres/multi_depth_model_woauxi.py +35 -0
  37. controlnet_aux/leres/leres/net_tools.py +54 -0
  38. controlnet_aux/leres/leres/network_auxi.py +417 -0
  39. controlnet_aux/leres/pix2pix/__init__.py +0 -0
  40. controlnet_aux/leres/pix2pix/__pycache__/__init__.cpython-310.pyc +0 -0
  41. controlnet_aux/leres/pix2pix/models/__init__.py +67 -0
  42. controlnet_aux/leres/pix2pix/models/__pycache__/__init__.cpython-310.pyc +0 -0
  43. controlnet_aux/leres/pix2pix/models/__pycache__/base_model.cpython-310.pyc +0 -0
  44. controlnet_aux/leres/pix2pix/models/__pycache__/base_model_hg.cpython-310.pyc +0 -0
  45. controlnet_aux/leres/pix2pix/models/__pycache__/networks.cpython-310.pyc +0 -0
  46. controlnet_aux/leres/pix2pix/models/__pycache__/pix2pix4depth_model.cpython-310.pyc +0 -0
  47. controlnet_aux/leres/pix2pix/models/base_model.py +244 -0
  48. controlnet_aux/leres/pix2pix/models/base_model_hg.py +58 -0
  49. controlnet_aux/leres/pix2pix/models/networks.py +623 -0
  50. 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