Add DUSt3R stereo method

README.md

@@ -23,6 +23,8 @@ Small Python utility to **compare and visualize** the output of various **stereo
 
 - [DistDepth](https://github.com/facebookresearch/DistDepth): "Toward Practical Monocular Indoor Depth Estimation" (CVPR 2022). This one is actually a **monocular** method, only using the left image.
 
+- [DUSt3R](https://github.com/naver/dust3r): "Geometric 3D Vision Made Easy" (CVPR 2024). Dense Unconstrained Stereo 3D Reconstruction using a ViT-Large encoder/cross-attention decoder pair. Weights are ~2.3 GB and downloaded on first use. **Non-commercial only (CC BY-NC-SA 4.0).**
+
 See below for more details / credits to get each of these working, and check this [blog post for more results, including performance numbers](https://nicolas.burrus.name/stereo-comparison/).
 
 https://user-images.githubusercontent.com/541507/169557430-48e62510-60c2-4a2b-8747-f9606e405f74.mp4
@@ -97,6 +99,11 @@ Sample images included in this repository:
 - [pytorch](https://pytorch.org/). To run pretrained models exported as torch script.
 - [depthai](https://docs.luxonis.com/en/latest/). Optional, to grab images from a Luxonis OAK camera.
 
+DUSt3R is supported as an experimental method through a minimal local PyTorch
+inference port, so it does not require installing the official DUSt3R package.
+Its checkpoints are large, and the official code and models are licensed for
+non-commercial use only.
+
 # Credits for each method
 
 I did not implement any of these myself, but just collected pre-trained models or converted them to torch script / ONNX.
@@ -129,6 +136,12 @@ I did not implement any of these myself, but just collected pre-trained models o
   - Official implementation and pre-trained models https://github.com/facebookresearch/DistDepth
   - I exported the pytorch implementaton to torch script via tracing, see [the changes](https://github.com/facebookresearch/DistDepth/commit/fde3b427ef2ff31c34f08e99c51c8e6a2427b720).
 
+- DUSt3R
+  - Official implementation and pre-trained models: https://github.com/naver/dust3r
+  - Minimal inference code adapted from: https://github.com/ibaiGorordo/dust3r-pytorch-inference-minimal
+  - DUSt3R predicts 3D point maps rather than stereo disparity. This adapter estimates DUSt3R's implicit focal length from the left point map and the predicted baseline from the right point map (assuming a rectified pair), then uses `f_dust3r * b_dust3r / Z_dust3r` directly as the pixel disparity. No metric-scale recovery via the user's calibration is needed for the disparity output.
+  - The official DUSt3R code and checkpoints are licensed under CC BY-NC-SA 4.0.
+
 # License
 
 The code of stereodemo is MIT licensed, but the pre-trained models are subject to the license of their respective implementation.

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "open3d>=0.15.1",
     "torch>=1.11.0",
     "torchvision",
+    "einops",
 ]
 
 [project.optional-dependencies]

stereodemo/dust3r_lib/__init__.py

@@ -0,0 +1,3 @@
+# Minimal DUSt3R inference implementation.
+# Sourced from https://github.com/ibaiGorordo/dust3r-pytorch-inference-minimal
+# Original DUSt3R: https://github.com/naver/dust3r (CC BY-NC-SA 4.0)

stereodemo/dust3r_lib/blocks.py

@@ -0,0 +1,105 @@
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+# Sourced from https://github.com/ibaiGorordo/dust3r-pytorch-inference-minimal
+
+import torch.nn as nn
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
+        super().__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.drop1 = nn.Dropout(drop)
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop2 = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop1(x)
+        x = self.fc2(x)
+        x = self.drop2(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, rope, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.rope = rope
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).transpose(1, 3)
+        q, k, v = [qkv[:, :, i] for i in range(3)]
+        q = self.rope(q)
+        k = self.rope(k)
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim, rope, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+        self.projq = nn.Linear(dim, dim, bias=qkv_bias)
+        self.projk = nn.Linear(dim, dim, bias=qkv_bias)
+        self.projv = nn.Linear(dim, dim, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.rope = rope
+
+    def forward(self, query, key, value):
+        B, Nq, C = query.shape
+        Nk = key.shape[1]
+        Nv = value.shape[1]
+        q = self.projq(query).reshape(B, Nq, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        k = self.projk(key).reshape(B, Nk, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        v = self.projv(value).reshape(B, Nv, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        q = self.rope(q)
+        k = self.rope(k)
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = (attn @ v).transpose(1, 2).reshape(B, Nq, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

stereodemo/dust3r_lib/decoder.py

@@ -0,0 +1,75 @@
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+# Sourced from https://github.com/ibaiGorordo/dust3r-pytorch-inference-minimal
+
+from functools import partial
+from copy import deepcopy
+
+import torch
+import torch.nn as nn
+
+from .blocks import DropPath, Mlp, Attention, CrossAttention
+from .third_party import RoPE2D
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, norm_mem=True, rope=None):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(dim, rope=rope, num_heads=num_heads, qkv_bias=qkv_bias,
+                              attn_drop=attn_drop, proj_drop=drop)
+        self.cross_attn = CrossAttention(dim, rope=rope, num_heads=num_heads, qkv_bias=qkv_bias,
+                                         attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        self.norm3 = norm_layer(dim)
+        self.mlp = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)
+        self.norm_y = norm_layer(dim) if norm_mem else nn.Identity()
+
+    def forward(self, x, y):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        y_ = self.norm_y(y)
+        x = x + self.drop_path(self.cross_attn(self.norm2(x), y_, y_))
+        x = x + self.drop_path(self.mlp(self.norm3(x)))
+        return x, y
+
+
+class Dust3rDecoder(nn.Module):
+    def __init__(self, ckpt_dict, batch=1, width=512, height=512, patch_size=16,
+                 enc_embed_dim=1024, dec_embed_dim=768, dec_num_heads=12, dec_depth=12,
+                 mlp_ratio=4., norm_im2_in_dec=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
+                 device=torch.device('cpu')):
+        super().__init__()
+        self.rope = RoPE2D(batch, width, height, patch_size, base=100.0, device=device)
+        self.decoder_embed = nn.Linear(enc_embed_dim, dec_embed_dim, bias=True)
+        self.dec_blocks = nn.ModuleList([
+            DecoderBlock(dec_embed_dim, dec_num_heads, mlp_ratio=mlp_ratio, qkv_bias=True,
+                         norm_layer=norm_layer, norm_mem=norm_im2_in_dec, rope=self.rope)
+            for _ in range(dec_depth)
+        ])
+        self.dec_blocks2 = deepcopy(self.dec_blocks)
+        self.dec_norm = norm_layer(dec_embed_dim)
+        self._load_checkpoint(ckpt_dict)
+        self.to(device)
+
+    @torch.inference_mode()
+    def forward(self, f1, f2):
+        f1_0 = f1_6 = f1_9 = f1
+        f2_0 = f2_6 = f2_9 = f2
+        f1_prev, f2_prev = self.decoder_embed(f1), self.decoder_embed(f2)
+        for i, (blk1, blk2) in enumerate(zip(self.dec_blocks, self.dec_blocks2), start=1):
+            f1, _ = blk1(f1_prev, f2_prev)
+            f2, _ = blk2(f2_prev, f1_prev)
+            f1_prev, f2_prev = f1, f2
+            if i == 6:
+                f1_6, f2_6 = f1, f2
+            elif i == 9:
+                f1_9, f2_9 = f1, f2
+        f1_12, f2_12 = self.dec_norm(f1), self.dec_norm(f2)
+        return f1_0, f1_6, f1_9, f1_12, f2_0, f2_6, f2_9, f2_12
+
+    def _load_checkpoint(self, ckpt_dict):
+        state = {k: v for k, v in ckpt_dict['model'].items()
+                 if k.startswith(('decoder_embed', 'dec_blocks', 'dec_norm'))}
+        self.load_state_dict(state, strict=True)

stereodemo/dust3r_lib/encoder.py

@@ -0,0 +1,70 @@
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+# Sourced from https://github.com/ibaiGorordo/dust3r-pytorch-inference-minimal
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+from .third_party import RoPE2D
+from .blocks import DropPath, Mlp, Attention
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, rope, mlp_ratio=4., qkv_bias=False, drop=0.,
+                 attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(dim, rope=rope, num_heads=num_heads, qkv_bias=qkv_bias,
+                              attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        self.mlp = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, img_size=(512, 512), patch_size=(16, 16), in_chans=3, embed_dim=768, norm_layer=None):
+        super().__init__()
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        x = self.proj(x)
+        x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        return self.norm(x)
+
+
+class Dust3rEncoder(nn.Module):
+    def __init__(self, ckpt_dict, batch=2, width=512, height=512, patch_size=16,
+                 enc_embed_dim=1024, enc_num_heads=16, enc_depth=24, mlp_ratio=4.,
+                 norm_layer=partial(nn.LayerNorm, eps=1e-6), device=torch.device('cpu')):
+        super().__init__()
+        self.patch_embed = PatchEmbed((height, width), (patch_size, patch_size), 3, enc_embed_dim)
+        self.rope = RoPE2D(batch, width, height, patch_size, base=100.0, device=device)
+        self.enc_blocks = nn.ModuleList([
+            Block(enc_embed_dim, enc_num_heads, self.rope, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
+            for _ in range(enc_depth)
+        ])
+        self.enc_norm = norm_layer(enc_embed_dim)
+        self._load_checkpoint(ckpt_dict)
+        self.to(device)
+
+    @torch.inference_mode()
+    def forward(self, x):
+        x = self.patch_embed(x)
+        for blk in self.enc_blocks:
+            x = blk(x)
+        return self.enc_norm(x)
+
+    def _load_checkpoint(self, ckpt_dict):
+        state = {k: v for k, v in ckpt_dict['model'].items()
+                 if k.startswith(('patch_embed', 'enc_blocks', 'enc_norm'))}
+        self.load_state_dict(state, strict=True)

stereodemo/dust3r_lib/head.py

@@ -0,0 +1,64 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+# Sourced from https://github.com/ibaiGorordo/dust3r-pytorch-inference-minimal
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .third_party import DPTHead
+
+
+def _reg_dense_depth(xyz):
+    d = xyz.norm(dim=-1, keepdim=True)
+    xyz = xyz / d.clip(min=1e-8)
+    return xyz * (torch.exp(d) - 1)
+
+
+def _reg_dense_conf(x, vmin=1):
+    return vmin + x.exp()
+
+
+def _postprocess(out):
+    fmap = out.permute(0, 2, 3, 1)  # B,H,W,4
+    depth = _reg_dense_depth(fmap[:, :, :, 0:3])
+    conf = _reg_dense_conf(fmap[:, :, :, 3])
+    return depth, conf
+
+
+class LinearPts3d(nn.Module):
+    def __init__(self, width=512, height=512, patch_size=16, dec_embed_dim=768, has_conf=True):
+        super().__init__()
+        self.patch_size = patch_size
+        self.has_conf = has_conf
+        self.num_h = height // patch_size
+        self.num_w = width // patch_size
+        self.proj = nn.Linear(dec_embed_dim, (3 + has_conf) * patch_size ** 2)
+
+    def forward(self, tokens_0, tokens_6, tokens_9, tokens_12):
+        B, S, D = tokens_12.shape
+        feat = self.proj(tokens_12)
+        feat = feat.transpose(-1, -2).view(B, -1, self.num_h, self.num_w)
+        return F.pixel_shuffle(feat, self.patch_size)
+
+
+class Dust3rHead(nn.Module):
+    def __init__(self, ckpt_dict, width=512, height=512, device=torch.device('cpu')):
+        super().__init__()
+        is_dpt = any('dpt' in k for k in ckpt_dict['model'].keys())
+        self.downstream_head1 = DPTHead(width, height) if is_dpt else LinearPts3d(width, height)
+        self.downstream_head2 = DPTHead(width, height) if is_dpt else LinearPts3d(width, height)
+        self._load_checkpoint(ckpt_dict)
+        self.to(device)
+
+    @torch.inference_mode()
+    def forward(self, d1_0, d1_6, d1_9, d1_12, d2_0, d2_6, d2_9, d2_12):
+        out1 = self.downstream_head1(d1_0, d1_6, d1_9, d1_12)
+        out2 = self.downstream_head2(d2_0, d2_6, d2_9, d2_12)
+        pts3d1, conf1 = _postprocess(out1)
+        pts3d2, conf2 = _postprocess(out2)
+        return pts3d1, conf1, pts3d2, conf2
+
+    def _load_checkpoint(self, ckpt_dict):
+        state = {k.replace('.dpt', ''): v for k, v in ckpt_dict['model'].items() if 'head' in k}
+        self.load_state_dict(state, strict=True)

stereodemo/dust3r_lib/third_party/__init__.py

@@ -0,0 +1,2 @@
+from .dpt_head import DPTHead
+from .rope2d import RoPE2D