selmee/depth-pro
4
1# Copyright (C) 2024 Apple Inc. All Rights Reserved.
2# Depth Pro: Sharp Monocular Metric Depth in Less Than a Second
3
4
5from __future__ import annotations
6
7from dataclasses import dataclass
8from typing import Mapping, Optional, Tuple, Union
9
10import torch
11from torch import nn
12from torchvision.transforms import (
13 Compose,
14 ConvertImageDtype,
15 Lambda,
16 Normalize,
17 ToTensor,
18)
19
20from .network.decoder import MultiresConvDecoder
21from .network.encoder import DepthProEncoder
22from .network.fov import FOVNetwork
23from .network.vit_factory import VIT_CONFIG_DICT, ViTPreset, create_vit
24
25
26@dataclass
27class DepthProConfig:
28 """Configuration for DepthPro."""
29
30 patch_encoder_preset: ViTPreset
31 image_encoder_preset: ViTPreset
32 decoder_features: int
33
34 checkpoint_uri: Optional[str] = None
35 fov_encoder_preset: Optional[ViTPreset] = None
36 use_fov_head: bool = True
37
38
39DEFAULT_MONODEPTH_CONFIG_DICT = DepthProConfig(
40 patch_encoder_preset="dinov2l16_384",
41 image_encoder_preset="dinov2l16_384",
42 checkpoint_uri="./checkpoints/depth_pro.pt",
43 decoder_features=256,
44 use_fov_head=True,
45 fov_encoder_preset="dinov2l16_384",
46)
47
48
49def create_backbone_model(
50 preset: ViTPreset
51) -> Tuple[nn.Module, ViTPreset]:
52 """Create and load a backbone model given a config.
53
54 Args:
55 ----
56 preset: A backbone preset to load pre-defind configs.
57
58 Returns:
59 -------
60 A Torch module and the associated config.
61
62 """
63 if preset in VIT_CONFIG_DICT:
64 config = VIT_CONFIG_DICT[preset]
65 model = create_vit(preset=preset, use_pretrained=False)
66 else:
67 raise KeyError(f"Preset {preset} not found.")
68
69 return model, config
70
71
72def create_model_and_transforms(
73 config: DepthProConfig = DEFAULT_MONODEPTH_CONFIG_DICT,
74 device: torch.device = torch.device("cpu"),
75 precision: torch.dtype = torch.float32,
76) -> Tuple[DepthPro, Compose]:
77 """Create a DepthPro model and load weights from `config.checkpoint_uri`.
78
79 Args:
80 ----
81 config: The configuration for the DPT model architecture.
82 device: The optional Torch device to load the model onto, default runs on "cpu".
83 precision: The optional precision used for the model, default is FP32.
84
85 Returns:
86 -------
87 The Torch DepthPro model and associated Transform.
88
89 """
90 patch_encoder, patch_encoder_config = create_backbone_model(
91 preset=config.patch_encoder_preset
92 )
93 image_encoder, _ = create_backbone_model(
94 preset=config.image_encoder_preset
95 )
96
97 fov_encoder = None
98 if config.use_fov_head and config.fov_encoder_preset is not None:
99 fov_encoder, _ = create_backbone_model(preset=config.fov_encoder_preset)
100
101 dims_encoder = patch_encoder_config.encoder_feature_dims
102 hook_block_ids = patch_encoder_config.encoder_feature_layer_ids
103 encoder = DepthProEncoder(
104 dims_encoder=dims_encoder,
105 patch_encoder=patch_encoder,
106 image_encoder=image_encoder,
107 hook_block_ids=hook_block_ids,
108 decoder_features=config.decoder_features,
109 )
110 decoder = MultiresConvDecoder(
111 dims_encoder=[config.decoder_features] + list(encoder.dims_encoder),
112 dim_decoder=config.decoder_features,
113 )
114 model = DepthPro(
115 encoder=encoder,
116 decoder=decoder,
117 last_dims=(32, 1),
118 use_fov_head=config.use_fov_head,
119 fov_encoder=fov_encoder,
120 ).to(device)
121
122 if precision == torch.half:
123 model.half()
124
125 transform = Compose(
126 [
127 ToTensor(),
128 Lambda(lambda x: x.to(device)),
129 Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
130 ConvertImageDtype(precision),
131 ]
132 )
133
134 if config.checkpoint_uri is not None:
135 state_dict = torch.load(config.checkpoint_uri, map_location="cpu")
136 missing_keys, unexpected_keys = model.load_state_dict(
137 state_dict=state_dict, strict=True
138 )
139
140 if len(unexpected_keys) != 0:
141 raise KeyError(
142 f"Found unexpected keys when loading monodepth: {unexpected_keys}"
143 )
144
145 # fc_norm is only for the classification head,
146 # which we would not use. We only use the encoding.
147 missing_keys = [key for key in missing_keys if "fc_norm" not in key]
148 if len(missing_keys) != 0:
149 raise KeyError(f"Keys are missing when loading monodepth: {missing_keys}")
150
151 return model, transform
152
153
154class DepthPro(nn.Module):
155 """DepthPro network."""
156
157 def __init__(
158 self,
159 encoder: DepthProEncoder,
160 decoder: MultiresConvDecoder,
161 last_dims: tuple[int, int],
162 use_fov_head: bool = True,
163 fov_encoder: Optional[nn.Module] = None,
164 ):
165 """Initialize DepthPro.
166
167 Args:
168 ----
169 encoder: The DepthProEncoder backbone.
170 decoder: The MultiresConvDecoder decoder.
171 last_dims: The dimension for the last convolution layers.
172 use_fov_head: Whether to use the field-of-view head.
173 fov_encoder: A separate encoder for the field of view.
174
175 """
176 super().__init__()
177
178 self.encoder = encoder
179 self.decoder = decoder
180
181 dim_decoder = decoder.dim_decoder
182 self.head = nn.Sequential(
183 nn.Conv2d(
184 dim_decoder, dim_decoder // 2, kernel_size=3, stride=1, padding=1
185 ),
186 nn.ConvTranspose2d(
187 in_channels=dim_decoder // 2,
188 out_channels=dim_decoder // 2,
189 kernel_size=2,
190 stride=2,
191 padding=0,
192 bias=True,
193 ),
194 nn.Conv2d(
195 dim_decoder // 2,
196 last_dims[0],
197 kernel_size=3,
198 stride=1,
199 padding=1,
200 ),
201 nn.ReLU(True),
202 nn.Conv2d(last_dims[0], last_dims[1], kernel_size=1, stride=1, padding=0),
203 nn.ReLU(),
204 )
205
206 # Set the final convoultion layer's bias to be 0.
207 self.head[4].bias.data.fill_(0)
208
209 # Set the FOV estimation head.
210 if use_fov_head:
211 self.fov = FOVNetwork(num_features=dim_decoder, fov_encoder=fov_encoder)
212
213 @property
214 def img_size(self) -> int:
215 """Return the internal image size of the network."""
216 return self.encoder.img_size
217
218 def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
219 """Decode by projection and fusion of multi-resolution encodings.
220
221 Args:
222 ----
223 x (torch.Tensor): Input image.
224
225 Returns:
226 -------
227 The canonical inverse depth map [m] and the optional estimated field of view [deg].
228
229 """
230 _, _, H, W = x.shape
231 print("Width:", W)
232 print("Height:", H)
233 assert H == self.img_size and W == self.img_size
234
235 encodings = self.encoder(x)
236 features, features_0 = self.decoder(encodings)
237 canonical_inverse_depth = self.head(features)
238
239 fov_deg = None
240 if hasattr(self, "fov"):
241 fov_deg = self.fov.forward(x, features_0.detach())
242
243 return canonical_inverse_depth, fov_deg
244
245 @torch.no_grad()
246 def infer(
247 self,
248 x: torch.Tensor,
249 f_px: Optional[Union[float, torch.Tensor]] = None,
250 interpolation_mode="bilinear",
251 ) -> Mapping[str, torch.Tensor]:
252 """Infer depth and fov for a given image.
253
254 If the image is not at network resolution, it is resized to 1536x1536 and
255 the estimated depth is resized to the original image resolution.
256 Note: if the focal length is given, the estimated value is ignored and the provided
257 focal length is use to generate the metric depth values.
258
259 Args:
260 ----
261 x (torch.Tensor): Input image
262 f_px (torch.Tensor): Optional focal length in pixels corresponding to `x`.
263 interpolation_mode (str): Interpolation function for downsampling/upsampling.
264
265 Returns:
266 -------
267 Tensor dictionary (torch.Tensor): depth [m], focallength [pixels].
268
269 """
270 if len(x.shape) == 3:
271 x = x.unsqueeze(0)
272 _, _, H, W = x.shape
273 resize = H != self.img_size or W != self.img_size
274
275 if resize:
276 x = nn.functional.interpolate(
277 x,
278 size=(self.img_size, self.img_size),
279 mode=interpolation_mode,
280 align_corners=False,
281 )
282
283 canonical_inverse_depth, fov_deg = self.forward(x)
284 if f_px is None:
285 f_px = 0.5 * W / torch.tan(0.5 * torch.deg2rad(fov_deg.to(torch.float)))
286
287 inverse_depth = canonical_inverse_depth * (W / f_px)
288 f_px = f_px.squeeze()
289
290 if resize:
291 inverse_depth = nn.functional.interpolate(
292 inverse_depth, size=(H, W), mode=interpolation_mode, align_corners=False
293 )
294
295 depth = 1.0 / torch.clamp(inverse_depth, min=1e-4, max=1e4)
296
297 return {
298 "depth": depth.squeeze(),
299 "focallength_px": f_px,
300 }
301