AdaBins: Depth Estimation using Adaptive Bins

In plain words

Given one ordinary photo, the job is to guess how far away each pixel is. A common approach is to split the possible distances into fixed slices and ask the network which slice each pixel falls into. That works, but the slices are the same for every photo, so a close-up of furniture and a long corridor get the same slices, and the output looks stepped.

This paper lets the network choose the slices for each photo. A small transformer looks at the whole image and decides how wide each slice should be, so narrow slices go where the scene's distances are likely to be. Each pixel then gets a set of weights over the slices, and its distance is the weighted average of the slice centres. Because the answer is an average rather than a single slice, the depth map comes out smooth. The authors report that this design beats earlier methods on two standard indoor and outdoor datasets.

TL;DR

• Monocular depth network. Its head predicts per-image bin widths over the depth range with a small transformer (mini-ViT), and per-pixel softmax weights over bins from range-attention maps §3.3.
• Final depth is the probability-weighted sum of bin centres, not the centre of the most likely bin §3.3 eq. 3.
• Trained with a scaled scale-invariant loss plus a Chamfer term that pulls bin centres toward ground-truth depths §3.4 eq. 4 eq. 5.
• Author-reported state of the art on NYU-Depth-v2 and KITTI; the code is public under GPL-3.0, and it differs from the paper in several training values (see Paper vs code).

Problem

Depth from one image is ill-posed, and direct regression from a convolutional encoder-decoder is hard §1. Methods that use fixed bins over the depth range give coarse, stepped outputs and cannot adapt to the depth distribution of each image §3.1 fig. 3. Methods that restrict the scene with planarity assumptions may fail outdoors §1. The paper asks whether global processing at high resolution, applied to the output depth distribution, improves accuracy §1 §3.1.

Method

Overview §3.3 fig. 2. An EfficientNet-B5 encoder with a standard upsampling decoder produces decoded features of size h x w x C_d at half the input resolution (h = H/2, w = W/2) §3.3 §3.2. The AdaBins block takes those features and outputs a depth map, which is bilinearly upsampled to full size §3.3.

Design choices §3.2. The paper compares four ways to set bin widths: fixed uniform bins, fixed log-scale bins, trained bin widths shared by all images, and per-image adaptive bins. It recommends per-image adaptive bins §3.2 tab. 6 fig. 3. It also predicts depth as a linear combination of bin centres rather than a single bin centre §3.2. The attention block is placed after the decoder, at high resolution §3.2.

Mini-ViT §3.3 fig. 4 tab. 1. A small vision transformer produces two things: the bin widths vector b, and the range-attention maps R. The decoded features are cut into non-overlapping patches with a convolution (kernel and stride equal to patch size p), flattened into S = hw/p^2 tokens, and given learned positional encodings. A transformer encoder runs over the tokens. The first output embedding goes to an MLP that outputs N raw bin scores, which are normalised by eq. 1. Output embeddings 2 through C+1 are used as 1x1 kernels over the decoded features, after a 3x3 convolution, giving C range-attention maps. The paper describes this as dot-product attention between pixel features (keys) and output embeddings (queries) §3.3.

Bin widths and centres §3.3 eq. 1 eq. 2. Raw scores get a small positive offset and are normalised to sum to 1. The bin edges are placed across [d_min, d_max] by the cumulative widths, and each centre is the midpoint of its bin eq. 2. The paper says d_min and d_max are fixed per dataset §3.2.

Hybrid regression §3.3 eq. 2 eq. 3. The range-attention maps pass through a 1x1 convolution to N channels and a softmax, giving per-pixel probabilities over the N bin centres §3.3. The depth at each pixel is the sum of centre times probability eq. 3. The paper says this avoids the discretisation artifacts of picking the most likely bin §3.3 fig. 5.

Loss §3.4 eq. 4 eq. 5 eq. 6. The pixel term is a scaled scale-invariant loss on log depth errors over valid ground-truth pixels eq. 4. The bin term is a two-way Chamfer distance between bin centres and all ground-truth depth values eq. 5. The total is the pixel term plus beta times the bin term eq. 6. Other losses tried (RMSE; SSIM plus L1) did worse in the authors' experiments §3.4.

# AdaBins head: bin widths -> bin centres -> depth  [§3.3]
# Input:  decoded features (B, C_d, h, w), h = H/2, w = W/2  [§3.3]
# Output: depth (B, 1, h, w) in metres; bin edges (B, N+1)

def adabins_head(decoded, d_min, d_max, num_bins):
    # Mini-ViT predicts raw bin scores and range-attention maps  [§3.3]
    raw_scores, range_maps = mini_vit(decoded)   # (B, N), (B, C, h, w)
    # Offset each score by eps, then normalise so the widths sum to 1
    # Eps is 1e-3 in the paper [eq. 1]
    # The code uses 0.1 [code: models/miniViT.py#L38]
    floored = raw_scores + eps
    widths = floored / sum(floored, over bins)   # (B, N)
    # Stretch the widths over [d_min, d_max] to get the bin edges  [eq. 2]
    edges = d_min + (d_max - d_min) * padded_cumsum(widths)   # (B, N+1)
    # Each bin centre is the midpoint of its two edges  [eq. 2]
    centres = (edges[:, :-1] + edges[:, 1:]) / 2   # (B, N)
    # Per-pixel softmax over bins, from a 1x1 conv on the range maps  [§3.3]
    probs = softmax(conv1x1(range_maps), over bins)   # (B, N, h, w)
    # Depth is the probability-weighted sum of centres  [eq. 3]
    depth = sum(probs * centres[:, :, None, None], over bins)
    return depth, edges

# Mini-ViT: raw bin scores and range-attention maps  [§3.3] [fig. 4]
# Input:  decoded features (B, C_d, h, w)
# Output: raw scores (B, N); range maps (B, C, h, w)

def mini_vit(decoded):
    # Patch embedding: conv with kernel = stride = p = 16  [§3.3] [tab. 1]
    patches = conv_embed(decoded)                 # (B, E, h/p, w/p), E = 128
    # Flatten to S = h*w/p^2 tokens, then add learned positional encodings
    # [§3.3]
    tokens = flatten(patches) + positional_encodings[:S]   # (B, S, E)
    # Transformer encoder: 4 layers, 4 heads, MLP size 1024  [tab. 1]
    outputs = transformer_encoder(tokens)         # (B, S, E)
    # First output embedding goes to an MLP with a ReLU output  [§3.3] [tab. 7]
    raw_scores = mlp_head(outputs[:, 0])          # (B, N)
    # Embeddings 2..C+1 are queries; keys come from a 3x3 conv  [§3.3]
    queries = outputs[:, 1 : C + 1]               # (B, C, E), C = 128 [tab. 1]
    keys = conv3x3(decoded)                       # (B, E, h, w)
    # Dot product of each pixel's key with each query  [§3.3]
    range_maps = dot_product(keys, queries)       # (B, C, h, w)
    return raw_scores, range_maps

# Training loss: scaled SI term plus Chamfer bin regulariser  [§3.4]
# Input:  pred (B, 1, H, W) in metres; target (B, 1, H, W); edges (B, N+1)
# Output: scalar loss

def adabins_loss(pred, target, edges):
    # Valid pixels only; the threshold is not stated in the paper [§3.4]
    # The code uses depth > 1e-3 [code: train.py#L189]
    valid = target > valid_threshold
    log_error = log(pred[valid]) - log(target[valid])   # g_i
    # Scaled SI term with lambda = 0.85 and alpha = 10  [eq. 4] [§3.4]
    spread = mean(log_error ** 2) - 0.85 * mean(log_error) ** 2
    pixel_loss = 10 * sqrt(spread)
    # Bin centres, then two-way Chamfer distance to ground-truth depths
    # [eq. 5] [§3.4]
    centres = (edges[:, 1:] + edges[:, :-1]) / 2
    bins_loss = (chamfer(centres, target_values)
                 + chamfer(target_values, centres))
    # Weighted sum with beta = 0.1  [eq. 6] [§3.4]
    return pixel_loss + 0.1 * bins_loss

Key equations, as written in the paper:

b_i = \frac{b'_i + \epsilon}{\sum_{j=1}^{N} (b'_j + \epsilon)} eq. 1

c(b_i) = d_{min} + (d_{max} - d_{min})\Big(\frac{b_i}{2} + \sum_{j=1}^{i-1} b_j\Big) eq. 2

\tilde{d} = \sum_{k=1}^{N} c(b_k)\, p_k eq. 3

\mathcal{L}_{pixel} = \alpha \sqrt{\frac{1}{T}\sum_i g_{i}^{2} - \frac{\lambda}{T^2}\Big(\sum_i g_i\Big)^2} eq. 4

\mathcal{L}_{bins} = chamfer(X, c(\mathbf{b})) + chamfer(c(\mathbf{b}), X) eq. 5

\mathcal{L}_{total} = \mathcal{L}_{pixel} + \beta\, \mathcal{L}_{bins} eq. 6

Hyperparameters

"Paper" values are from the paper. "Code" values come from the official training and test argument files, as noted. "not stated" means the paper does not say.

Results

Author-reported. Setup: NYU-Depth-v2 trained on a 50K-image subset and tested on the official test set with Eigen cropping, up to 10 m §4.1 §4.2. KITTI trained on about 26K left-view images outside the 697 Eigen test images, tested with the Garg crop, up to 80 m §4.1. Baseline numbers are copied from each method's paper; BTS was retrained by the authors to check its numbers §4.3. DAV's numbers come from its paper, but the authors note that DAV evaluates at 1/4 resolution §4.3, footnote.

NYU-Depth-v2 §4.3 tab. 2. Higher is better for delta metrics; lower for REL and RMS. "--" means the paper gives no value. The source marks best entries in bold and second-best in underline; those marks are not reproduced here.

The paper states a significant improvement over the state of the art on all metrics §4.3.

KITTI (Eigen split) §4.3 tab. 3. Columns: delta1, delta2, delta3, REL, Sq Rel, RMS, RMS log.

The text gives RMS and Sq Rel gains of about 13.5% and 22.4% over the previous state of the art §4.3; see Pitfalls for which rows these match.

Cross-dataset: trained on NYU, tested on SUN RGB-D, no fine-tuning §4.3 tab. 5. Test set of 5050 images §4.1. Metrics: delta1, delta2, delta3, REL, RMS, log10.

Ablation: loss function (NYU) §4.4 tab. 4. Columns: delta1, delta2, delta3, REL, RMS, log10.

Ablation: bin widths and regression (NYU) §4.4 tab. 6. "R" is standard regression, "HR" is hybrid regression, "Fix" means the widths are fixed.

Number of bins N §4.4 fig. 6. Error rises first and then falls as N grows, and gains flatten above N = 256, which the paper uses §4.4. Values are plotted only; no table is given.

Paper vs code

Official code: shariqfarooq123/AdaBins at commit 0952d91e9e762be310bb4cd055cbfe2448c0ce20, GPL-3.0 LICENSE. It is linked from a footnote on the paper's abstract, and the owner name matches the first author. The code is copyleft, so this card cites it by path and line and does not quote it. Everything below was checked by reading the code and its argument files.

• Bin floor epsilon. Paper: 1e-3 eq. 1. Code: 0.1, added after a ReLU miniViT.py. The effect is larger than the value suggests: with N = 256, the floors add 25.6 to the normaliser in the code and 0.256 with 1e-3.
• Weight decay. Paper: 1e-2 §4.2. Code: 0.1 args_train_nyu.txt args_train_kitti_eigen.txt.
• Max learning rate. Paper: 3.5e-4 §4.2. Code: 0.000357 args_train_nyu.txt args_train_kitti_eigen.txt.
• Final learning rate. Paper: cosine annealing to max_lr/75 §4.2. Code: OneCycleLR with div factor 25 and final div factor 100 args_train_nyu.txt train.py. Under PyTorch's OneCycleLR definition (library behaviour, not in this repo), the minimum is max_lr/(25 x 100) = max_lr/2500. This is my reading of the library, not a value the repo prints.
• Warm-up fraction. Paper: first 30% of iterations §4.2. Code: no pct_start passed train.py, so the library default applies (not in this repo).
• Momentum. Not in the paper. Code: OneCycleLR cycles momentum between 0.85 and 0.95 train.py.
• Encoder learning rate. Not in the paper. Code: encoder at lr/10, decoder, mini-ViT and output conv at lr train.py unet_adaptive_bins.py.
• Gradient clipping. Not in the paper. Code: clip to norm 0.1, with an "optional" comment train.py.
• Bin count default. Paper: N = 256 §4.4. Code args files: 256 args_train_nyu.txt. The training script's default is 80 if the flag is omitted train.py; the evaluation script's default is 256 evaluate.py.
• Depth range. Paper: d_min and d_max fixed per dataset §3.2. Code: d_min = 1e-3 m for both datasets, d_max = 10 m (NYU) or 80 m (KITTI) args_train_nyu.txt args_train_kitti_eigen.txt.
• SI loss form. Paper eq. 4 eq. 4. Code: 10 x sqrt(var(g) + 0.15 mean(g)^2) loss.py. The two are algebraically the same if the variance is 1/T, but torch.var is called without a correction argument loss.py, and PyTorch's default is the unbiased estimator (1/(T-1)). This is a library default, not stated in the repo. Check it before comparing numbers.
• Valid pixels. Paper: pixels with valid ground truth §3.4. Code: depth above 1e-3 for the loss train.py and for the Chamfer term loss.py.
• Chamfer term. Paper: all ground-truth depth values §3.4. Code: valid ground-truth pixels of each image, flattened loss.py, through pytorch3d's chamfer_distance loss.py loss.py. The paper does not name the library or the distance convention.
• Bin centres and edges. Paper eq. 2 eq. 2. Code: widths scaled by (d_max - d_min), a leading edge at min_val, cumulative sum, centres at midpoints unet_adaptive_bins.py. The formula matches with d_min = min_val.
• Depth prediction. Paper eq. 3 eq. 3. Code: sum of softmax probabilities times centres unet_adaptive_bins.py. Matches.
• Mini-ViT configuration. Paper: p = 16, E = 128, 4 layers, 4 heads, MLP 1024, C = 128 tab. 1. Code: same values unet_adaptive_bins.py layers.py. Matches.
• Positional encodings. Paper: learned, no table size or init given §3.3. Code: learned table of 500 rows with uniform random init layers.py. Sequences longer than 500 tokens will fail.
• Decoder skip features. Not in the paper. Code takes encoder outputs 4, 5, 6, 8 and 11 as skips unet_adaptive_bins.py.
• Decoder first conv. Not in the paper. The 1x1 projection has padding=1 unet_adaptive_bins.py. My reading: this adds 2 pixels per spatial side, and the bilinear resize in the first upsampling block absorbs it unet_adaptive_bins.py.
• Output head and ReLU. Paper: MLP with ReLU output §3.3 tab. 7. Code: LeakyReLU hidden layers, ReLU applied in the normalisation step miniViT.py miniViT.py. Matches.
• Test-time clipping. Not in the paper. Code clips predictions to [min_depth, max_depth] before averaging evaluate.py evaluate.py and to the eval range at validation train.py.
• Test-time flip. Paper: mirror averaging §4.1. Code: same evaluate.py. Matches.
• Augmentation. Not in the paper. Code: random rotation, flips, gamma 0.9 to 1.1, brightness (NYU 0.75 to 1.25, KITTI 0.9 to 1.1), per-channel colour scaling dataloader.py dataloader.py.
• Depth units. Not stated in the paper. Code: divide stored depth by 1000 (NYU) or 256 (KITTI) dataloader.py.
• KITTI crop and eval naming. Paper: 704 x 352 random crop; Garg crop for evaluation §4.1 §4.2. Code: 352 x 704 crop args_train_kitti_eigen.txt; Garg crop and KB crop at test args_test_kitti_eigen.txt. Matches, but the test argument file is named kitti_eigen.
• Training resolution. Paper: not stated for NYU training crops §4. Code: 416 x 544 args_train_nyu.txt.

Pitfalls

• Two configurations, one paper. The paper's optimiser values (weight decay 1e-2, LR 3.5e-4, final LR max_lr/75) differ from the released configs (0.1, 0.000357, max_lr/2500 under OneCycleLR). Pick one and say which.
• Epsilon changes the widths. The floor enters the normaliser of every bin width eq. 1. With N = 256 the code's floors sum to 25.6 and the paper's to 0.256 (see Paper vs code). Pick one value, state it, and check the bin-width distribution when you change it.
• Encoder learning rate. The paper does not mention a lower LR for the encoder. An implementer following only the paper will train the encoder at the full rate, which the code does not do.
• Depth range. The paper says d_min is set per dataset. The code uses 1e-3 m for every dataset. A different d_min changes every bin centre.
• Sequence length limit. Mini-ViT tokens are the decoder output after 16 x 16 patching, so S = HW/1024 for a 2x-downsampled decoder output with p = 16. The code's 500-row positional table caps S at 500 layers.py. At 480 x 640 input, S = 300 args_test_nyu.txt. Check the limit when using larger inputs.
• Cost of the head. The range-attention maps hold C = 128 channels at half resolution tab. 1 §3.3. The paper reports the total parameter count but no runtime or memory figure for the AdaBins head tab. 1 §4.2.
• Training-only details. Gradient clipping, momentum cycling, encoder LR and augmentation are all absent from the paper. Dropping any of them changes training.
• Ablation text vs table. The text says the Trained-but-Fixed variant performs worst among all choices §4.4. In table 6, standard regression (Base + R) is worst on delta1, delta2, REL and RMS, and Train-Fix is worst among the fixed-bin variants only on REL. Its delta1 is higher than Uniform-Fix's. Use the table.
• Which baselines the text gains are against. The text gives about 13.5% RMS and 22.4% Sq Rel gains §4.3. These match the second-best rows in table 3 for those metrics (RMS: Fu et al. 2.727; Sq Rel: BTS 0.245), by my arithmetic. The paper does not name the baseline.
• Comparisons across papers. Most baseline numbers are copied from original papers, not re-run under one protocol. DAV's evaluation uses 1/4 resolution §4.3, footnote; BTS was retrained by the authors and the other rows were not §4.3.
• Figure and equation numbers. Counted from the LaTeX source in order; the compiled PDF was not checked. Equations 1 to 6 are the only numbered equations in the source.

Lineage

• Builds on: the encoder-decoder baseline of Alhashim and Wonka [Alhashim2018, §2, §3.3]; ordinal regression with fixed bins from Fu et al. [Fu2018DeepOR, §3.1]; the Vision Transformer, for the patch tokens [dosovitskiy2020, §3.3]; EfficientNet-B5 as the encoder [efficentnet_TanL19, §3.2]; the scale-invariant loss of Eigen et al. [Eigen2014, §3.4]; the Chamfer loss [Fan2017_chamfer_loss, §3.4].
• Compared against: Eigen et al., Laina et al., Hao et al., Lee et al., Fu et al. (DORN), SharpNet, Hu et al., Chen et al., Yin et al., BTS, DAV on NYU tab. 2; Eigen, Liu, Godard, Kuznietsov, Gan, Fu (DORN), Yin, BTS, Saxena on KITTI tab. 3; Chen, Yin, BTS on SUN RGB-D tab. 5.