1use std::sync::Arc;
8
9use candle::D::Minus1;
10use candle::{Module, Result, Tensor};
11use candle_nn::ops::Identity;
12use candle_nn::{
13 batch_norm, conv2d, conv2d_no_bias, conv_transpose2d, linear, seq, Activation, BatchNorm,
14 BatchNormConfig, Conv2d, Conv2dConfig, ConvTranspose2dConfig, Sequential, VarBuilder,
15};
16
17use crate::models::dinov2::DinoVisionTransformer;
18
19pub struct DepthAnythingV2Config {
20 out_channel_sizes: [usize; 4],
21 in_channel_size: usize, num_features: usize,
23 use_batch_norm: bool,
24 use_class_token: bool,
25 layer_ids_vits: Vec<usize>,
26 input_image_size: usize,
27 target_patch_size: usize,
28}
29
30impl DepthAnythingV2Config {
31 #[allow(clippy::too_many_arguments)]
32 pub fn new(
33 out_channel_sizes: [usize; 4],
34 in_channel_size: usize,
35 num_features: usize,
36 use_batch_norm: bool,
37 use_class_token: bool,
38 layer_ids_vits: Vec<usize>,
39 input_image_size: usize,
40 target_patch_size: usize,
41 ) -> Self {
42 Self {
43 out_channel_sizes,
44 in_channel_size,
45 num_features,
46 use_batch_norm,
47 use_class_token,
48 layer_ids_vits,
49 input_image_size,
50 target_patch_size,
51 }
52 }
53
54 pub fn vit_small() -> Self {
55 Self {
56 out_channel_sizes: [48, 96, 192, 384],
57 in_channel_size: 384,
58 num_features: 64,
59 use_batch_norm: false,
60 use_class_token: false,
61 layer_ids_vits: vec![2, 5, 8, 11],
62 input_image_size: 518,
63 target_patch_size: 518 / 14,
64 }
65 }
66
67 pub fn vit_base() -> Self {
68 Self {
69 out_channel_sizes: [96, 192, 384, 768],
70 in_channel_size: 768,
71 num_features: 128,
72 use_batch_norm: false,
73 use_class_token: false,
74 layer_ids_vits: vec![2, 5, 8, 11],
75 input_image_size: 518,
76 target_patch_size: 518 / 14,
77 }
78 }
79
80 pub fn vit_large() -> Self {
81 Self {
82 out_channel_sizes: [256, 512, 1024, 1024],
83 in_channel_size: 1024,
84 num_features: 256,
85 use_batch_norm: false,
86 use_class_token: false,
87 layer_ids_vits: vec![4, 11, 17, 23],
88 input_image_size: 518,
89 target_patch_size: 518 / 14,
90 }
91 }
92
93 pub fn vit_giant() -> Self {
94 Self {
95 out_channel_sizes: [1536, 1536, 1536, 1536],
96 in_channel_size: 1536,
97 num_features: 384,
98 use_batch_norm: false,
99 use_class_token: false,
100 layer_ids_vits: vec![9, 19, 29, 39],
101 input_image_size: 518,
102 target_patch_size: 518 / 14,
103 }
104 }
105}
106
107pub struct ResidualConvUnit {
108 activation: Activation,
109 conv1: Conv2d,
110 conv2: Conv2d,
111 batch_norm1: Option<BatchNorm>,
112 batch_norm2: Option<BatchNorm>,
113}
114
115impl ResidualConvUnit {
116 pub fn new(
117 conf: &DepthAnythingV2Config,
118 activation: Activation,
119 vb: VarBuilder,
120 ) -> Result<Self> {
121 const KERNEL_SIZE: usize = 3;
122 let conv_cfg = Conv2dConfig {
123 padding: 1,
124 stride: 1,
125 dilation: 1,
126 groups: 1,
127 cudnn_fwd_algo: None,
128 };
129 let conv1 = conv2d(
130 conf.num_features,
131 conf.num_features,
132 KERNEL_SIZE,
133 conv_cfg,
134 vb.pp("conv1"),
135 )?;
136 let conv2 = conv2d(
137 conf.num_features,
138 conf.num_features,
139 KERNEL_SIZE,
140 conv_cfg,
141 vb.pp("conv2"),
142 )?;
143
144 let (batch_norm1, batch_norm2) = match conf.use_batch_norm {
145 true => {
146 let batch_norm_cfg = BatchNormConfig {
147 eps: 1e-05,
148 remove_mean: false,
149 affine: true,
150 momentum: 0.1,
151 };
152 (
153 Some(batch_norm(conf.num_features, batch_norm_cfg, vb.pp("bn1"))?),
154 Some(batch_norm(conf.num_features, batch_norm_cfg, vb.pp("bn2"))?),
155 )
156 }
157 false => (None, None),
158 };
159
160 Ok(Self {
161 activation,
162 conv1,
163 conv2,
164 batch_norm1,
165 batch_norm2,
166 })
167 }
168}
169
170impl Module for ResidualConvUnit {
171 fn forward(&self, xs: &Tensor) -> Result<Tensor> {
172 let out = self.activation.forward(xs)?;
173 let out = self.conv1.forward(&out)?;
174 let out = if let Some(batch_norm1) = &self.batch_norm1 {
175 batch_norm1.forward_train(&out)?
176 } else {
177 out
178 };
179
180 let out = self.activation.forward(&out)?;
181 let out = self.conv2.forward(&out)?;
182 let out = if let Some(batch_norm2) = &self.batch_norm2 {
183 batch_norm2.forward_train(&out)?
184 } else {
185 out
186 };
187
188 out + xs
189 }
190}
191
192pub struct FeatureFusionBlock {
193 res_conv_unit1: ResidualConvUnit,
194 res_conv_unit2: ResidualConvUnit,
195 output_conv: Conv2d,
196 target_patch_size: usize,
197}
198
199impl FeatureFusionBlock {
200 pub fn new(
201 conf: &DepthAnythingV2Config,
202 target_patch_size: usize,
203 activation: Activation,
204 vb: VarBuilder,
205 ) -> Result<Self> {
206 const KERNEL_SIZE: usize = 1;
207 let conv_cfg = Conv2dConfig {
208 padding: 0,
209 stride: 1,
210 dilation: 1,
211 groups: 1,
212 cudnn_fwd_algo: None,
213 };
214 let output_conv = conv2d(
215 conf.num_features,
216 conf.num_features,
217 KERNEL_SIZE,
218 conv_cfg,
219 vb.pp("out_conv"),
220 )?;
221 let res_conv_unit1 = ResidualConvUnit::new(conf, activation, vb.pp("resConfUnit1"))?;
222 let res_conv_unit2 = ResidualConvUnit::new(conf, activation, vb.pp("resConfUnit2"))?;
223
224 Ok(Self {
225 res_conv_unit1,
226 res_conv_unit2,
227 output_conv,
228 target_patch_size,
229 })
230 }
231}
232
233impl Module for FeatureFusionBlock {
234 fn forward(&self, xs: &Tensor) -> Result<Tensor> {
235 let out = self.res_conv_unit2.forward(xs)?;
236 let out = out.interpolate2d(self.target_patch_size, self.target_patch_size)?;
237
238 self.output_conv.forward(&out)
239 }
240}
241
242pub struct Scratch {
243 layer1_rn: Conv2d,
244 layer2_rn: Conv2d,
245 layer3_rn: Conv2d,
246 layer4_rn: Conv2d,
247 refine_net1: FeatureFusionBlock,
248 refine_net2: FeatureFusionBlock,
249 refine_net3: FeatureFusionBlock,
250 refine_net4: FeatureFusionBlock,
251 output_conv1: Conv2d,
252 output_conv2: Sequential,
253}
254
255impl Scratch {
256 pub fn new(conf: &DepthAnythingV2Config, vb: VarBuilder) -> Result<Self> {
257 const KERNEL_SIZE: usize = 3;
258 let conv_cfg = Conv2dConfig {
259 padding: 1,
260 stride: 1,
261 dilation: 1,
262 groups: 1,
263 cudnn_fwd_algo: None,
264 };
265
266 let layer1_rn = conv2d_no_bias(
267 conf.out_channel_sizes[0],
268 conf.num_features,
269 KERNEL_SIZE,
270 conv_cfg,
271 vb.pp("layer1_rn"),
272 )?;
273 let layer2_rn = conv2d_no_bias(
274 conf.out_channel_sizes[1],
275 conf.num_features,
276 KERNEL_SIZE,
277 conv_cfg,
278 vb.pp("layer2_rn"),
279 )?;
280 let layer3_rn = conv2d_no_bias(
281 conf.out_channel_sizes[2],
282 conf.num_features,
283 KERNEL_SIZE,
284 conv_cfg,
285 vb.pp("layer3_rn"),
286 )?;
287 let layer4_rn = conv2d_no_bias(
288 conf.out_channel_sizes[3],
289 conf.num_features,
290 KERNEL_SIZE,
291 conv_cfg,
292 vb.pp("layer4_rn"),
293 )?;
294
295 let refine_net1 = FeatureFusionBlock::new(
296 conf,
297 conf.target_patch_size * 8,
298 Activation::Relu,
299 vb.pp("refinenet1"),
300 )?;
301 let refine_net2 = FeatureFusionBlock::new(
302 conf,
303 conf.target_patch_size * 4,
304 Activation::Relu,
305 vb.pp("refinenet2"),
306 )?;
307 let refine_net3 = FeatureFusionBlock::new(
308 conf,
309 conf.target_patch_size * 2,
310 Activation::Relu,
311 vb.pp("refinenet3"),
312 )?;
313 let refine_net4 = FeatureFusionBlock::new(
314 conf,
315 conf.target_patch_size,
316 Activation::Relu,
317 vb.pp("refinenet4"),
318 )?;
319
320 let conv_cfg = Conv2dConfig {
321 padding: 1,
322 stride: 1,
323 dilation: 1,
324 groups: 1,
325 cudnn_fwd_algo: None,
326 };
327 let output_conv1 = conv2d(
328 conf.num_features,
329 conf.num_features / 2,
330 KERNEL_SIZE,
331 conv_cfg,
332 vb.pp("output_conv1"),
333 )?;
334
335 let output_conv2 = seq();
336 const HEAD_FEATURES_2: usize = 32;
337 const OUT_CHANNELS_2: usize = 1;
338 const KERNEL_SIZE_2: usize = 1;
339 let output_conv2 = output_conv2.add(conv2d(
340 conf.num_features / 2,
341 HEAD_FEATURES_2,
342 KERNEL_SIZE,
343 conv_cfg,
344 vb.pp("output_conv2").pp("0"),
345 )?);
346 let output_conv2 = output_conv2
347 .add(Activation::Relu)
348 .add(conv2d(
349 HEAD_FEATURES_2,
350 OUT_CHANNELS_2,
351 KERNEL_SIZE_2,
352 conv_cfg,
353 vb.pp("output_conv2").pp("2"),
354 )?)
355 .add(Activation::Relu);
356
357 Ok(Self {
358 layer1_rn,
359 layer2_rn,
360 layer3_rn,
361 layer4_rn,
362 refine_net1,
363 refine_net2,
364 refine_net3,
365 refine_net4,
366 output_conv1,
367 output_conv2,
368 })
369 }
370}
371
372const NUM_CHANNELS: usize = 4;
373
374pub struct DPTHead {
375 projections: Vec<Conv2d>,
376 resize_layers: Vec<Box<dyn Module>>,
377 readout_projections: Vec<Sequential>,
378 scratch: Scratch,
379 use_class_token: bool,
380 input_image_size: usize,
381 target_patch_size: usize,
382}
383
384impl DPTHead {
385 pub fn new(conf: &DepthAnythingV2Config, vb: VarBuilder) -> Result<Self> {
386 let mut projections: Vec<Conv2d> = Vec::with_capacity(conf.out_channel_sizes.len());
387 for (conv_index, out_channel_size) in conf.out_channel_sizes.iter().enumerate() {
388 projections.push(conv2d(
389 conf.in_channel_size,
390 *out_channel_size,
391 1,
392 Default::default(),
393 vb.pp("projects").pp(conv_index.to_string()),
394 )?);
395 }
396
397 let resize_layers: Vec<Box<dyn Module>> = vec![
398 Box::new(conv_transpose2d(
399 conf.out_channel_sizes[0],
400 conf.out_channel_sizes[0],
401 4,
402 ConvTranspose2dConfig {
403 padding: 0,
404 stride: 4,
405 dilation: 1,
406 output_padding: 0,
407 },
408 vb.pp("resize_layers").pp("0"),
409 )?),
410 Box::new(conv_transpose2d(
411 conf.out_channel_sizes[1],
412 conf.out_channel_sizes[1],
413 2,
414 ConvTranspose2dConfig {
415 padding: 0,
416 stride: 2,
417 dilation: 1,
418 output_padding: 0,
419 },
420 vb.pp("resize_layers").pp("1"),
421 )?),
422 Box::new(Identity::new()),
423 Box::new(conv2d(
424 conf.out_channel_sizes[3],
425 conf.out_channel_sizes[3],
426 3,
427 Conv2dConfig {
428 padding: 1,
429 stride: 2,
430 dilation: 1,
431 groups: 1,
432 cudnn_fwd_algo: None,
433 },
434 vb.pp("resize_layers").pp("3"),
435 )?),
436 ];
437
438 let readout_projections = if conf.use_class_token {
439 let rop = Vec::with_capacity(NUM_CHANNELS);
440 for rop_index in 0..NUM_CHANNELS {
441 seq()
442 .add(linear(
443 2 * conf.in_channel_size,
444 conf.in_channel_size,
445 vb.pp("readout_projects").pp(rop_index.to_string()),
446 )?)
447 .add(Activation::Gelu);
448 }
449 rop
450 } else {
451 vec![]
452 };
453
454 let scratch = Scratch::new(conf, vb.pp("scratch"))?;
455
456 Ok(Self {
457 projections,
458 resize_layers,
459 readout_projections,
460 scratch,
461 use_class_token: conf.use_class_token,
462 input_image_size: conf.input_image_size,
463 target_patch_size: conf.target_patch_size,
464 })
465 }
466}
467
468impl Module for DPTHead {
469 fn forward(&self, xs: &Tensor) -> Result<Tensor> {
470 let mut out: Vec<Tensor> = Vec::with_capacity(NUM_CHANNELS);
471 for i in 0..NUM_CHANNELS {
472 let x = if self.use_class_token {
473 let x = xs.get(i)?.get(0)?;
474 let class_token = xs.get(i)?.get(1)?;
475 let readout = class_token.unsqueeze(1)?.expand(x.shape())?;
476 let to_cat = [x, readout];
477 let cat = Tensor::cat(&to_cat, Minus1)?;
478 self.readout_projections[i].forward(&cat)?
479 } else {
480 xs.get(i)?
481 };
482 let x_dims = x.dims();
483
484 let x = x.permute((0, 2, 1))?.reshape((
485 x_dims[0],
486 x_dims[x_dims.len() - 1],
487 self.target_patch_size,
488 self.target_patch_size,
489 ))?;
490 let x = self.projections[i].forward(&x)?;
491
492 let x = self.resize_layers[i].forward(&x)?;
493 out.push(x);
494 }
495
496 let layer_1_rn = self.scratch.layer1_rn.forward(&out[0])?;
497 let layer_2_rn = self.scratch.layer2_rn.forward(&out[1])?;
498 let layer_3_rn = self.scratch.layer3_rn.forward(&out[2])?;
499 let layer_4_rn = self.scratch.layer4_rn.forward(&out[3])?;
500
501 let path4 = self.scratch.refine_net4.forward(&layer_4_rn)?;
502
503 let res3_out = self
504 .scratch
505 .refine_net3
506 .res_conv_unit1
507 .forward(&layer_3_rn)?;
508 let res3_out = path4.add(&res3_out)?;
509 let path3 = self.scratch.refine_net3.forward(&res3_out)?;
510
511 let res2_out = self
512 .scratch
513 .refine_net2
514 .res_conv_unit1
515 .forward(&layer_2_rn)?;
516 let res2_out = path3.add(&res2_out)?;
517 let path2 = self.scratch.refine_net2.forward(&res2_out)?;
518
519 let res1_out = self
520 .scratch
521 .refine_net1
522 .res_conv_unit1
523 .forward(&layer_1_rn)?;
524 let res1_out = path2.add(&res1_out)?;
525 let path1 = self.scratch.refine_net1.forward(&res1_out)?;
526
527 let out = self.scratch.output_conv1.forward(&path1)?;
528
529 let out = out.interpolate2d(self.input_image_size, self.input_image_size)?;
530
531 self.scratch.output_conv2.forward(&out)
532 }
533}
534
535pub struct DepthAnythingV2 {
536 pretrained: Arc<DinoVisionTransformer>,
537 depth_head: DPTHead,
538 conf: DepthAnythingV2Config,
539}
540
541impl DepthAnythingV2 {
542 pub fn new(
543 pretrained: Arc<DinoVisionTransformer>,
544 conf: DepthAnythingV2Config,
545 vb: VarBuilder,
546 ) -> Result<Self> {
547 let depth_head = DPTHead::new(&conf, vb.pp("depth_head"))?;
548
549 Ok(Self {
550 pretrained,
551 depth_head,
552 conf,
553 })
554 }
555}
556
557impl Module for DepthAnythingV2 {
558 fn forward(&self, xs: &Tensor) -> Result<Tensor> {
559 let features = self.pretrained.get_intermediate_layers(
560 xs,
561 &self.conf.layer_ids_vits,
562 false,
563 false,
564 true,
565 )?;
566 let depth = self.depth_head.forward(&features)?;
567
568 depth.relu()
569 }
570}