Skip to content

Commit 77edbb9

Browse files
thias15thias15
committed
add models and update links
1 parent 4e3fbc8 commit 77edbb9

File tree

1 file changed

+143
-5
lines changed

1 file changed

+143
-5
lines changed

‎hubconf.py‎

Lines changed: 143 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,144 @@
66
from midas.midas_net import MidasNet
77
from midas.midas_net_custom import MidasNet_small
88

9+
def DPT_BEit_L_512(pretrained=True, **kwargs):
10+
""" # This docstring shows up in hub.help()
11+
MiDaS DPT_BEit_L_512 model for monocular depth estimation
12+
pretrained (bool): load pretrained weights into model
13+
"""
14+
15+
model = DPTDepthModel(
16+
path=None,
17+
backbone="beitl16_512",
18+
non_negative=True,
19+
)
20+
21+
if pretrained:
22+
checkpoint = (
23+
"https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_large_512.pt"
24+
)
25+
state_dict = torch.hub.load_state_dict_from_url(
26+
checkpoint, map_location=torch.device('cpu'), progress=True, check_hash=True
27+
)
28+
model.load_state_dict(state_dict)
29+
30+
return model
31+
32+
def DPT_BEit_L_384(pretrained=True, **kwargs):
33+
""" # This docstring shows up in hub.help()
34+
MiDaS DPT_BEit_L_384 model for monocular depth estimation
35+
pretrained (bool): load pretrained weights into model
36+
"""
37+
38+
model = DPTDepthModel(
39+
path=None,
40+
backbone="beitl16_384",
41+
non_negative=True,
42+
)
43+
44+
if pretrained:
45+
checkpoint = (
46+
"https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_large_384.pt"
47+
)
48+
state_dict = torch.hub.load_state_dict_from_url(
49+
checkpoint, map_location=torch.device('cpu'), progress=True, check_hash=True
50+
)
51+
model.load_state_dict(state_dict)
52+
53+
return model
54+
55+
def DPT_SwinV2_L_384(pretrained=True, **kwargs):
56+
""" # This docstring shows up in hub.help()
57+
MiDaS DPT_SwinV2_L_384 model for monocular depth estimation
58+
pretrained (bool): load pretrained weights into model
59+
"""
60+
61+
model = DPTDepthModel(
62+
path=None,
63+
backbone="swin2l24_384",
64+
non_negative=True,
65+
)
66+
67+
if pretrained:
68+
checkpoint = (
69+
"https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_swin2_large_384.pt"
70+
)
71+
state_dict = torch.hub.load_state_dict_from_url(
72+
checkpoint, map_location=torch.device('cpu'), progress=True, check_hash=True
73+
)
74+
model.load_state_dict(state_dict)
75+
76+
return model
77+
78+
def DPT_SwinV2_T_256(pretrained=True, **kwargs):
79+
""" # This docstring shows up in hub.help()
80+
MiDaS DPT_SwinV2_T_256 model for monocular depth estimation
81+
pretrained (bool): load pretrained weights into model
82+
"""
83+
84+
model = DPTDepthModel(
85+
path=None,
86+
backbone="swin2t16_256",
87+
non_negative=True,
88+
)
89+
90+
if pretrained:
91+
checkpoint = (
92+
"https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_swin2_tiny_256.pt"
93+
)
94+
state_dict = torch.hub.load_state_dict_from_url(
95+
checkpoint, map_location=torch.device('cpu'), progress=True, check_hash=True
96+
)
97+
model.load_state_dict(state_dict)
98+
99+
return model
100+
101+
def DPT_Next_ViT_L_384(pretrained=True, **kwargs):
102+
""" # This docstring shows up in hub.help()
103+
MiDaS DPT_Next_ViT_L_384 model for monocular depth estimation
104+
pretrained (bool): load pretrained weights into model
105+
"""
106+
107+
model = DPTDepthModel(
108+
path=None,
109+
backbone="next_vit_large_6m",
110+
non_negative=True,
111+
)
112+
113+
if pretrained:
114+
checkpoint = (
115+
"https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_next_vit_large_384.pt"
116+
)
117+
state_dict = torch.hub.load_state_dict_from_url(
118+
checkpoint, map_location=torch.device('cpu'), progress=True, check_hash=True
119+
)
120+
model.load_state_dict(state_dict)
121+
122+
return model
123+
124+
def DPT_LeViT_224(pretrained=True, **kwargs):
125+
""" # This docstring shows up in hub.help()
126+
MiDaS DPT_LeViT_224 model for monocular depth estimation
127+
pretrained (bool): load pretrained weights into model
128+
"""
129+
130+
model = DPTDepthModel(
131+
path=None,
132+
backbone="levit_384",
133+
non_negative=True,
134+
)
135+
136+
if pretrained:
137+
checkpoint = (
138+
"https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_levit_224.pt"
139+
)
140+
state_dict = torch.hub.load_state_dict_from_url(
141+
checkpoint, map_location=torch.device('cpu'), progress=True, check_hash=True
142+
)
143+
model.load_state_dict(state_dict)
144+
145+
return model
146+
9147
def DPT_Large(pretrained=True, **kwargs):
10148
""" # This docstring shows up in hub.help()
11149
MiDaS DPT-Large model for monocular depth estimation
@@ -20,7 +158,7 @@ def DPT_Large(pretrained=True, **kwargs):
20158

21159
if pretrained:
22160
checkpoint = (
23-
"https://github.com/isl-org/MiDaS/releases/download/v3/dpt_large-midas-2f21e586.pt"
161+
"https://github.com/isl-org/MiDaS/releases/download/v3/dpt_large_384.pt"
24162
)
25163
state_dict = torch.hub.load_state_dict_from_url(
26164
checkpoint, map_location=torch.device('cpu'), progress=True, check_hash=True
@@ -43,7 +181,7 @@ def DPT_Hybrid(pretrained=True, **kwargs):
43181

44182
if pretrained:
45183
checkpoint = (
46-
"https://github.com/isl-org/MiDaS/releases/download/v3/dpt_hybrid-midas-501f0c75.pt"
184+
"https://github.com/isl-org/MiDaS/releases/download/v3/dpt_hybrid_384.pt"
47185
)
48186
state_dict = torch.hub.load_state_dict_from_url(
49187
checkpoint, map_location=torch.device('cpu'), progress=True, check_hash=True
@@ -62,7 +200,7 @@ def MiDaS(pretrained=True, **kwargs):
62200

63201
if pretrained:
64202
checkpoint = (
65-
"https://github.com/isl-org/MiDaS/releases/download/v2_1/model-f6b98070.pt"
203+
"https://github.com/isl-org/MiDaS/releases/download/v2_1/midas_v21_384.pt"
66204
)
67205
state_dict = torch.hub.load_state_dict_from_url(
68206
checkpoint, map_location=torch.device('cpu'), progress=True, check_hash=True
@@ -73,15 +211,15 @@ def MiDaS(pretrained=True, **kwargs):
73211

74212
def MiDaS_small(pretrained=True, **kwargs):
75213
""" # This docstring shows up in hub.help()
76-
MiDaS small model for monocular depth estimation on resource-constrained devices
214+
MiDaS v2.1 small model for monocular depth estimation on resource-constrained devices
77215
pretrained (bool): load pretrained weights into model
78216
"""
79217

80218
model = MidasNet_small(None, features=64, backbone="efficientnet_lite3", exportable=True, non_negative=True, blocks={'expand': True})
81219

82220
if pretrained:
83221
checkpoint = (
84-
"https://github.com/isl-org/MiDaS/releases/download/v2_1/model-small-70d6b9c8.pt"
222+
"https://github.com/isl-org/MiDaS/releases/download/v2_1/midas_v21_small_256.pt"
85223
)
86224
state_dict = torch.hub.load_state_dict_from_url(
87225
checkpoint, map_location=torch.device('cpu'), progress=True, check_hash=True

0 commit comments

Comments
 (0)