diff --git a/README.md b/README.md index f36ae2e..ffe79dd 100644 --- a/README.md +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index 2af721b..3f6bf81 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "open3d>=0.15.1", "torch>=1.11.0", "torchvision", + "einops", ] [project.optional-dependencies] diff --git a/stereodemo/dust3r_lib/__init__.py b/stereodemo/dust3r_lib/__init__.py new file mode 100644 index 0000000..e4f33aa --- /dev/null +++ b/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) diff --git a/stereodemo/dust3r_lib/blocks.py b/stereodemo/dust3r_lib/blocks.py new file mode 100644 index 0000000..2473bfc --- /dev/null +++ b/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 diff --git a/stereodemo/dust3r_lib/decoder.py b/stereodemo/dust3r_lib/decoder.py new file mode 100644 index 0000000..be3ccd5 --- /dev/null +++ b/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) diff --git a/stereodemo/dust3r_lib/encoder.py b/stereodemo/dust3r_lib/encoder.py new file mode 100644 index 0000000..73fda19 --- /dev/null +++ b/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) diff --git a/stereodemo/dust3r_lib/head.py b/stereodemo/dust3r_lib/head.py new file mode 100644 index 0000000..3e4862f --- /dev/null +++ b/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) diff --git a/stereodemo/dust3r_lib/third_party/__init__.py b/stereodemo/dust3r_lib/third_party/__init__.py new file mode 100644 index 0000000..3dd6e54 --- /dev/null +++ b/stereodemo/dust3r_lib/third_party/__init__.py @@ -0,0 +1,2 @@ +from .dpt_head import DPTHead +from .rope2d import RoPE2D diff --git a/stereodemo/dust3r_lib/third_party/dpt_head.py b/stereodemo/dust3r_lib/third_party/dpt_head.py new file mode 100644 index 0000000..d39c08c --- /dev/null +++ b/stereodemo/dust3r_lib/third_party/dpt_head.py @@ -0,0 +1,170 @@ +# 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 +# References: https://github.com/isl-org/DPT +# https://github.com/EPFL-VILAB/MultiMAE/blob/main/multimae/output_adapters.py + +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from typing import Union, Tuple + + +def pair(t): + return t if isinstance(t, tuple) else (t, t) + + +def make_scratch(in_shape, out_shape, groups=1, expand=False): + scratch = nn.Module() + out_shape1 = out_shape + out_shape2 = out_shape + out_shape3 = out_shape + out_shape4 = out_shape + if expand: + out_shape2 = out_shape * 2 + out_shape3 = out_shape * 4 + out_shape4 = out_shape * 8 + + scratch.layer1_rn = nn.Conv2d(in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups) + scratch.layer2_rn = nn.Conv2d(in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups) + scratch.layer3_rn = nn.Conv2d(in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups) + scratch.layer4_rn = nn.Conv2d(in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups) + scratch.layer_rn = nn.ModuleList([scratch.layer1_rn, scratch.layer2_rn, scratch.layer3_rn, scratch.layer4_rn]) + return scratch + + +class ResidualConvUnit_custom(nn.Module): + def __init__(self, features, activation, bn): + super().__init__() + self.bn = bn + self.groups = 1 + self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=not self.bn, groups=self.groups) + self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=not self.bn, groups=self.groups) + if self.bn: + self.bn1 = nn.BatchNorm2d(features) + self.bn2 = nn.BatchNorm2d(features) + self.activation = activation + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, x): + out = self.activation(x) + out = self.conv1(out) + if self.bn: + out = self.bn1(out) + out = self.activation(out) + out = self.conv2(out) + if self.bn: + out = self.bn2(out) + return self.skip_add.add(out, x) + + +class FeatureFusionBlock_custom(nn.Module): + def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, width_ratio=1): + super().__init__() + self.width_ratio = width_ratio + self.deconv = deconv + self.align_corners = align_corners + self.groups = 1 + self.expand = expand + out_features = features // 2 if expand else features + self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1) + self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn) + self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn) + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, *xs): + output = xs[0] + if len(xs) == 2: + res = self.resConfUnit1(xs[1]) + if self.width_ratio != 1: + res = F.interpolate(res, size=(output.shape[2], output.shape[3]), mode='bilinear') + output = self.skip_add.add(output, res) + output = self.resConfUnit2(output) + if self.width_ratio != 1: + if (output.shape[3] / output.shape[2]) < (2 / 3) * self.width_ratio: + shape = 3 * output.shape[3] + else: + shape = int(self.width_ratio * 2 * output.shape[2]) + output = F.interpolate(output, size=(2 * output.shape[2], shape), mode='bilinear') + else: + output = nn.functional.interpolate(output, scale_factor=2, mode="bilinear", align_corners=self.align_corners) + output = self.out_conv(output) + return output + + +def make_fusion_block(features, use_bn, width_ratio=1): + return FeatureFusionBlock_custom(features, nn.ReLU(False), deconv=False, bn=use_bn, expand=False, + align_corners=True, width_ratio=width_ratio) + + +class Interpolate(nn.Module): + def __init__(self, scale_factor, mode, align_corners=False): + super().__init__() + self.interp = nn.functional.interpolate + self.scale_factor = scale_factor + self.mode = mode + self.align_corners = align_corners + + def forward(self, x): + return self.interp(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners) + + +class DPTHead(nn.Module): + def __init__(self, width=512, height=512, num_channels: int = 4, stride_level: int = 1, + patch_size: Union[int, Tuple[int, int]] = 16, layer_dims: Tuple[int] = (96, 192, 384, 768), + feature_dim: int = 256, last_dim: int = 128, use_bn: bool = False, + dim_tokens_enc: Tuple[int] = (1024, 768, 768, 768), output_width_ratio=1, **kwargs): + super().__init__() + self.num_channels = num_channels + self.stride_level = stride_level + self.patch_size = pair(patch_size) + self.layer_dims = layer_dims + self.feature_dim = feature_dim + self.dim_tokens_enc = dim_tokens_enc + self.P_H = max(1, self.patch_size[0] // stride_level) + self.P_W = max(1, self.patch_size[1] // stride_level) + self.num_w = width // (self.stride_level * self.P_W) + self.num_h = height // (self.stride_level * self.P_H) + self.scratch = make_scratch(layer_dims, feature_dim, groups=1, expand=False) + self.scratch.refinenet1 = make_fusion_block(feature_dim, use_bn, output_width_ratio) + self.scratch.refinenet2 = make_fusion_block(feature_dim, use_bn, output_width_ratio) + self.scratch.refinenet3 = make_fusion_block(feature_dim, use_bn, output_width_ratio) + self.scratch.refinenet4 = make_fusion_block(feature_dim, use_bn, output_width_ratio) + self.head = nn.Sequential( + nn.Conv2d(feature_dim, feature_dim // 2, kernel_size=3, stride=1, padding=1), + Interpolate(scale_factor=2, mode="bilinear", align_corners=True), + nn.Conv2d(feature_dim // 2, last_dim, kernel_size=3, stride=1, padding=1), + nn.ReLU(True), + nn.Conv2d(last_dim, self.num_channels, kernel_size=1, stride=1, padding=0), + ) + self.act_postprocess = self._init_act_postprocess() + + def _init_act_postprocess(self): + act_postprocess = nn.ModuleList() + act_postprocess.append(nn.Sequential( + nn.Conv2d(self.dim_tokens_enc[0], self.layer_dims[0], kernel_size=1, stride=1, padding=0), + nn.ConvTranspose2d(self.layer_dims[0], self.layer_dims[0], kernel_size=4, stride=4, padding=0, bias=True), + )) + act_postprocess.append(nn.Sequential( + nn.Conv2d(self.dim_tokens_enc[1], self.layer_dims[1], kernel_size=1, stride=1, padding=0), + nn.ConvTranspose2d(self.layer_dims[1], self.layer_dims[1], kernel_size=2, stride=2, padding=0, bias=True), + )) + act_postprocess.append(nn.Sequential( + nn.Conv2d(self.dim_tokens_enc[2], self.layer_dims[2], kernel_size=1, stride=1, padding=0), + )) + act_postprocess.append(nn.Sequential( + nn.Conv2d(self.dim_tokens_enc[3], self.layer_dims[3], kernel_size=1, stride=1, padding=0), + nn.Conv2d(self.layer_dims[3], self.layer_dims[3], kernel_size=3, stride=2, padding=1), + )) + return act_postprocess + + def forward(self, tokens_0, tokens_6, tokens_9, tokens_12): + layers = [tokens_0, tokens_6, tokens_9, tokens_12] + layers = [rearrange(l, 'b (nh nw) c -> b c nh nw', nh=self.num_h, nw=self.num_w) for l in layers] + layers = [self.act_postprocess[idx](l) for idx, l in enumerate(layers)] + layers = [self.scratch.layer_rn[idx](l) for idx, l in enumerate(layers)] + path_4 = self.scratch.refinenet4(layers[3])[:, :, :layers[2].shape[2], :layers[2].shape[3]] + path_3 = self.scratch.refinenet3(path_4, layers[2]) + path_2 = self.scratch.refinenet2(path_3, layers[1]) + path_1 = self.scratch.refinenet1(path_2, layers[0]) + return self.head(path_1) diff --git a/stereodemo/dust3r_lib/third_party/rope2d.py b/stereodemo/dust3r_lib/third_party/rope2d.py new file mode 100644 index 0000000..7a3bd59 --- /dev/null +++ b/stereodemo/dust3r_lib/third_party/rope2d.py @@ -0,0 +1,47 @@ +# 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 + + +def get_positions(b, h, w, device): + x = torch.arange(w, device=device) + y = torch.arange(h, device=device) + positions = torch.cartesian_prod(y, x) + return positions.view(1, h * w, 2).expand(b, -1, 2) + + +def get_cos_sin(base, D, seq_len, device, dtype): + inv_freq = 1.0 / (base ** (torch.arange(0, D, 2).float().to(device) / D)) + t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) + freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +class RoPE2D(torch.nn.Module): + def __init__(self, batch=2, width=512, height=288, patch_size=16, base=100.0, D=32, + device=torch.device('cpu'), dtype=torch.float32): + super().__init__() + pos = get_positions(batch, height // patch_size, width // patch_size, device) + pos_x, pos_y = pos[:, :, 1], pos[:, :, 0] + cos, sin = get_cos_sin(base, D, int(pos.max()) + 1, device, dtype) + self.cos_x = torch.nn.functional.embedding(pos_x, cos)[:, None, :, :] + self.sin_x = torch.nn.functional.embedding(pos_x, sin)[:, None, :, :] + self.cos_y = torch.nn.functional.embedding(pos_y, cos)[:, None, :, :] + self.sin_y = torch.nn.functional.embedding(pos_y, sin)[:, None, :, :] + + @staticmethod + def rotate_half(x): + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + def apply_rope1d(self, tokens, cos, sin): + return (tokens * cos) + (self.rotate_half(tokens) * sin) + + def forward(self, tokens): + y, x = tokens.chunk(2, dim=-1) + x = self.apply_rope1d(x, self.cos_x, self.sin_x) + y = self.apply_rope1d(y, self.cos_y, self.sin_y) + return torch.cat((y, x), dim=-1) diff --git a/stereodemo/main.py b/stereodemo/main.py index 6d58558..0559172 100644 --- a/stereodemo/main.py +++ b/stereodemo/main.py @@ -24,6 +24,7 @@ from .method_hitnet import HitnetStereo from .method_sttr import StereoTransformers from stereodemo.method_dist_depth import DistDepth +from stereodemo.method_dust3r import DUSt3R def parse_args(): import argparse @@ -170,7 +171,8 @@ def main(): HitnetStereo(config), StereoTransformers(config), ChangRealtimeStereo(config), - DistDepth(config) + DistDepth(config), + DUSt3R(config), ] if args.images: diff --git a/stereodemo/method_dust3r.py b/stereodemo/method_dust3r.py new file mode 100644 index 0000000..189f8d1 --- /dev/null +++ b/stereodemo/method_dust3r.py @@ -0,0 +1,270 @@ +from pathlib import Path +import time + +import cv2 +import numpy as np +import torch +import torch.nn as nn +from PIL import Image +from torchvision import transforms + +from .methods import Config, EnumParameter, IntParameter, StereoMethod, InputPair, StereoOutput, Calibration +from . import utils + + +urls = { + "DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth": "https://download.europe.naverlabs.com/ComputerVision/DUSt3R/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth", + "DUSt3R_ViTLarge_BaseDecoder_224_linear.pth": "https://download.europe.naverlabs.com/ComputerVision/DUSt3R/DUSt3R_ViTLarge_BaseDecoder_224_linear.pth", +} + + +class Dust3rModel(nn.Module): + def __init__(self, ckpt_dict, width: int, height: int, device: str): + super().__init__() + from .dust3r_lib.encoder import Dust3rEncoder + from .dust3r_lib.decoder import Dust3rDecoder + from .dust3r_lib.head import Dust3rHead + + torch_device = torch.device(device) + self.encoder = Dust3rEncoder(ckpt_dict, batch=2, width=width, height=height, device=torch_device) + self.decoder = Dust3rDecoder(ckpt_dict, batch=1, width=width, height=height, device=torch_device) + self.head = Dust3rHead(ckpt_dict, width=width, height=height, device=torch_device) + + @torch.inference_mode() + def forward(self, left, right): + encoded = self.encoder(torch.cat([left, right], dim=0)) + f1, f2 = encoded[0:1], encoded[1:2] + return self.head(*self.decoder(f1, f2)) + + +class Dust3rStereo(StereoMethod): + def __init__(self, config: Config): + super().__init__( + "DUSt3R (CVPR 2024)", + "DUSt3R: Geometric 3D Vision Made Easy. Experimental metric stereo adapter using official PyTorch inference.", + {}, + config) + self.reset_defaults() + self.net = None + self._loaded_model = None + self._loaded_device = None + self._loaded_size = None + + def reset_defaults(self): + self.parameters.update({ + "Device": EnumParameter("Inference device", 0, ["Auto", "CPU", "CUDA"]), + "Model": EnumParameter("Pre-trained model", 0, ["512-dpt", "224-linear"]), + "Min Confidence x100": IntParameter("Minimum DUSt3R confidence multiplied by 100", 25, 0, 1000), + }) + + def compute_disparity(self, input: InputPair) -> StereoOutput: + device = self._selected_device() + image_size = self._image_size() + min_conf = self.parameters["Min Confidence x100"].value / 100.0 + + left_view, _left_calib, resize_scale = self._prepare_view( + input.left_image, input.calibration, image_size, idx=0) + right_view, _right_calib, _ = self._prepare_view( + input.right_image, input.calibration, image_size, idx=1) + if left_view["img"].shape != right_view["img"].shape: + raise RuntimeError("DUSt3R preprocessing produced mismatched left/right tensor shapes.") + self._load_model(self._model_filename(), device, left_view["img"].shape[-1], left_view["img"].shape[-2]) + + start = time.time() + with torch.inference_mode(): + pts3d_left, conf_left, pts3d_right_in_left, conf_right = self.net( + left_view["img"].to(device), right_view["img"].to(device)) + elapsed_time = time.time() - start + + pts3d_left = self._as_numpy(pts3d_left)[0] + conf_left = np.squeeze(self._as_numpy(conf_left)[0]) + pts3d_right_in_left = self._as_numpy(pts3d_right_in_left)[0] + conf_right = np.squeeze(self._as_numpy(conf_right)[0]) + + # DUSt3R predicts a self-consistent (f_dust3r, pts3d) for the processed + # crop with its principal point at the geometric image center, not at + # the user's calibrated principal point. The actual pixel disparity in + # the rectified pair is f_dust3r * b_dust3r / Z_dust3r, which we can + # compute purely from DUSt3R's outputs — using the user's calibrated + # focal length here would bias the recovered baseline. + f_dust3r = self._estimate_dust3r_focal(pts3d_left, conf_left, min_conf) + if f_dust3r is None: + raise RuntimeError("DUSt3R: not enough confident predictions to estimate the implicit focal length.") + b_dust3r = self._estimate_dust3r_baseline( + pts3d_right_in_left, conf_right, f_dust3r, min_conf) + if b_dust3r is None or b_dust3r <= 1e-6: + raise RuntimeError("DUSt3R: failed to estimate the predicted baseline from the right view.") + + Z = pts3d_left[:, :, 2].astype(np.float32) + actual_pixel_disparity_resized = np.float32(f_dust3r * b_dust3r) + with np.errstate(divide="ignore", invalid="ignore"): + disparity_map = np.where(Z > 0, actual_pixel_disparity_resized / Z, np.float32(-1.0)).astype(np.float32) + invalid = (conf_left < min_conf) | ~np.isfinite(disparity_map) | (disparity_map <= 0) + disparity_map[invalid] = -1.0 + + if disparity_map.shape[:2] != input.left_image.shape[:2]: + disparity_map = cv2.resize( + disparity_map, + (input.left_image.shape[1], input.left_image.shape[0]), + interpolation=cv2.INTER_NEAREST) + valid = disparity_map > 0 + # Disparity scales with the image-resize factor only; the centred + # crop does not change disparity values. + disparity_map[valid] *= np.float32(1.0 / resize_scale) + + return StereoOutput(disparity_map, input.left_image, elapsed_time) + + def _selected_device(self) -> str: + requested = self.parameters["Device"].value + if requested == "CPU": + return "cpu" + if requested == "CUDA": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA was selected for DUSt3R, but torch.cuda.is_available() is false.") + return "cuda" + return "cuda" if torch.cuda.is_available() else "cpu" + + def _model_filename(self): + if self.parameters["Model"].value == "224-linear": + return "DUSt3R_ViTLarge_BaseDecoder_224_linear.pth" + return "DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth" + + def _image_size(self): + return 224 if self.parameters["Model"].value == "224-linear" else 512 + + def _load_model(self, model_name: str, device: str, width: int, height: int): + size = (width, height) + if self._loaded_model == model_name and self._loaded_device == device and self._loaded_size == size: + return + + model_path = self.config.models_path / model_name + if not model_path.exists(): + utils.download_model(urls[model_path.name], model_path) + try: + ckpt_dict = torch.load(model_path, map_location="cpu", weights_only=False) + except TypeError: + ckpt_dict = torch.load(model_path, map_location="cpu") + try: + self.net = Dust3rModel(ckpt_dict, width, height, device).eval() + except ImportError as e: + raise RuntimeError("DUSt3R requires the small einops package for its local DPT head implementation.") from e + self._loaded_model = model_name + self._loaded_device = device + self._loaded_size = size + + def _prepare_view(self, image_bgr: np.ndarray, calibration: Calibration, size: int, idx: int): + image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) + pil_image = Image.fromarray(image_rgb) + processed, calib, scale = self._resize_and_crop(pil_image, calibration, size) + img_norm = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), + ]) + return { + "img": img_norm(processed)[None], + "true_shape": np.int32([processed.size[::-1]]), + "idx": idx, + "instance": str(idx), + }, calib, scale + + def _resize_and_crop(self, image: Image.Image, calibration: Calibration, size: int): + width, height = image.size + if size == 224: + long_edge_size = round(size * max(width / height, height / width)) + else: + long_edge_size = size + + scale = long_edge_size / max(width, height) + resized_width = int(round(width * scale)) + resized_height = int(round(height * scale)) + resampling = getattr(Image, "Resampling", Image) + resample = resampling.LANCZOS if max(width, height) > long_edge_size else resampling.BICUBIC + image = image.resize((resized_width, resized_height), resample) + + cx, cy = resized_width // 2, resized_height // 2 + if size == 224: + half_width = min(cx, cy) + half_height = half_width + else: + patch_size = 16 + half_width = int(((2 * cx) // patch_size) * patch_size / 2) + half_height = int(((2 * cy) // patch_size) * patch_size / 2) + if resized_width == resized_height: + half_height = int(3 * half_width / 4) + + left = int(cx - half_width) + top = int(cy - half_height) + right = int(cx + half_width) + bottom = int(cy + half_height) + image = image.crop((left, top, right, bottom)) + + calib = Calibration( + width=right - left, + height=bottom - top, + fx=calibration.fx * scale, + fy=calibration.fy * scale, + cx0=calibration.cx0 * scale - left, + cx1=calibration.cx1 * scale - left, + cy=calibration.cy * scale - top, + baseline_meters=calibration.baseline_meters, + depth_range=calibration.depth_range, + left_image_rect_normalized=np.array([0., 0., 1., 1.]), + comment=calibration.comment, + ) + return image, calib, scale + + def _estimate_dust3r_focal(self, pts3d: np.ndarray, confidence: np.ndarray, min_conf: float): + """Estimate DUSt3R's implicit focal length from a per-pixel point map. + + DUSt3R places its principal point at the geometric image center of the + processed crop. For pixel (u, v) and predicted point (X, Y, Z): + (u - W/2) = f * X / Z, (v - H/2) = f * Y / Z + so f = Z * sqrt((u - W/2)^2 + (v - H/2)^2) / sqrt(X^2 + Y^2). We take + the per-pixel median for robustness. + """ + height, width = pts3d.shape[:2] + pp_x = width / 2.0 + pp_y = height / 2.0 + valid = (confidence >= min_conf) & np.isfinite(pts3d).all(axis=2) & (pts3d[:, :, 2] > 0) + ys, xs = np.nonzero(valid) + if xs.size < 64: + return None + points = pts3d[ys, xs] + pixel_radius = np.hypot(xs - pp_x, ys - pp_y) + point_radius = np.hypot(points[:, 0], points[:, 1]) + keep = point_radius > 1e-6 + if keep.sum() < 64: + return None + f_per_pixel = points[keep, 2] * pixel_radius[keep] / point_radius[keep] + return float(np.median(f_per_pixel)) + + def _estimate_dust3r_baseline(self, pts3d_right_in_left: np.ndarray, confidence: np.ndarray, + f_dust3r: float, min_conf: float): + """Estimate the predicted baseline assuming a rectified pair. + + For pixel (u, v) of the right image with predicted point (X, Y, Z) in + the left camera frame, with the right camera at (b, 0, 0) and no + rotation: + (u - W/2) = f_dust3r * (X - b) / Z + so b = X - (u - W/2) * Z / f_dust3r. Taking the per-pixel median is + robust against the outlier predictions DUSt3R produces in low-texture + or occluded regions. + """ + height, width = pts3d_right_in_left.shape[:2] + pp_x = width / 2.0 + valid = (confidence >= min_conf) & np.isfinite(pts3d_right_in_left).all(axis=2) & (pts3d_right_in_left[:, :, 2] > 0) + ys, xs = np.nonzero(valid) + if xs.size < 64: + return None + points = pts3d_right_in_left[ys, xs] + b_per_pixel = points[:, 0] - (xs - pp_x) * points[:, 2] / f_dust3r + return float(np.median(b_per_pixel)) + + @staticmethod + def _as_numpy(value): + if isinstance(value, torch.Tensor): + return value.detach().cpu().numpy() + return np.asarray(value) + + +DUSt3R = Dust3rStereo diff --git a/tests/test_methods.py b/tests/test_methods.py index 7fd2320..88acf69 100755 --- a/tests/test_methods.py +++ b/tests/test_methods.py @@ -14,6 +14,7 @@ from stereodemo import method_cre_stereo from stereodemo import method_raft_stereo from stereodemo import method_sttr +from stereodemo import method_dust3r from stereodemo import method_opencv_cuda_bp from stereodemo import method_opencv_cuda_csbp from stereodemo.methods import Config, InputPair, Calibration, StereoOutput, StereoMethod @@ -73,6 +74,19 @@ def test_sttr(self): m.parameters["Shape"].set_value ("640x480 (ds3)") self.check_method (m, 7.4636, 0.9869) + def test_dust3r_optional_registration(self): + method_dust3r.DUSt3R(config) + + def test_dust3r(self): + m = method_dust3r.DUSt3R(config) + model_path = config.models_path / m._model_filename() + if not model_path.exists(): + self.skipTest("DUSt3R checkpoint is not cached") + output = m.compute_disparity(input) + self.assertEqual(output.disparity_pixels.shape, input.left_image.shape[:2]) + valid_pixels = output.disparity_pixels[output.disparity_pixels > 0.] + self.assertGreater(valid_pixels.size, 0) + def test_cuda_bp(self): if not self._opencv_cuda_available(): self.skipTest("OpenCV CUDA is not available") diff --git a/uv.lock b/uv.lock index 0b852cc..ab9e209 100644 --- a/uv.lock +++ b/uv.lock @@ -918,7 +918,7 @@ name = "cuda-bindings" version = "13.2.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder", marker = "python_full_version >= '3.10'" }, + { name = "cuda-pathfinder", marker = "(python_full_version == '3.10.*' and sys_platform == 'emscripten') or (python_full_version == '3.10.*' and sys_platform == 'win32') or (python_full_version >= '3.10' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/1a/fe/7351d7e586a8b4c9f89731bfe4cf0148223e8f9903ff09571f78b3fb0682/cuda_bindings-13.2.0-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:08b395f79cb89ce0cd8effff07c4a1e20101b873c256a1aeb286e8fd7bd0f556", size = 5744254, upload-time = "2026-03-11T00:12:29.798Z" }, @@ -953,37 +953,37 @@ wheels = [ [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version >= '3.10' and sys_platform == 'win32')" }, + { name = "nvidia-cublas", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'win32')" }, ] cudart = [ - { name = "nvidia-cuda-runtime", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version >= '3.10' and sys_platform == 'win32')" }, + { name = "nvidia-cuda-runtime", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'win32')" }, ] cufft = [ - { name = "nvidia-cufft", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version >= '3.10' and sys_platform == 'win32')" }, + { name = "nvidia-cufft", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'win32')" }, ] cufile = [ { name = "nvidia-cufile", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version >= '3.10' and sys_platform == 'win32')" }, + { name = "nvidia-cuda-cupti", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'win32')" }, ] curand = [ - { name = "nvidia-curand", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version >= '3.10' and sys_platform == 'win32')" }, + { name = "nvidia-curand", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'win32')" }, ] cusolver = [ - { name = "nvidia-cusolver", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version >= '3.10' and sys_platform == 'win32')" }, + { name = "nvidia-cusolver", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'win32')" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version >= '3.10' and sys_platform == 'win32')" }, + { name = "nvidia-cusparse", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'win32')" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version >= '3.10' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'win32')" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version >= '3.10' and sys_platform == 'win32')" }, + { name = "nvidia-cuda-nvrtc", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'win32')" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version >= '3.10' and sys_platform == 'win32')" }, + { name = "nvidia-nvtx", marker = "(python_full_version >= '3.10' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'win32')" }, ] [[package]] @@ -1101,6 +1101,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/10/5da547df7a391dcde17f59520a231527b8571e6f46fc8efb02ccb370ab12/docutils-0.22.4-py3-none-any.whl", hash = "sha256:d0013f540772d1420576855455d050a2180186c91c15779301ac2ccb3eeb68de", size = 633196, upload-time = "2025-12-18T19:00:18.077Z" }, ] +[[package]] +name = "einops" +version = "0.8.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.9'", +] +sdist = { url = "https://files.pythonhosted.org/packages/e5/81/df4fbe24dff8ba3934af99044188e20a98ed441ad17a274539b74e82e126/einops-0.8.1.tar.gz", hash = "sha256:de5d960a7a761225532e0f1959e5315ebeafc0cd43394732f103ca44b9837e84", size = 54805, upload-time = "2025-02-09T03:17:00.434Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/62/9773de14fe6c45c23649e98b83231fffd7b9892b6cf863251dc2afa73643/einops-0.8.1-py3-none-any.whl", hash = "sha256:919387eb55330f5757c6bea9165c5ff5cfe63a642682ea788a6d472576d81737", size = 64359, upload-time = "2025-02-09T03:17:01.998Z" }, +] + +[[package]] +name = "einops" +version = "0.8.2" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.11' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.11' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.11' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.10.*'", + "python_full_version > '3.9' and python_full_version < '3.10'", + "python_full_version == '3.9'", +] +sdist = { url = "https://files.pythonhosted.org/packages/2c/77/850bef8d72ffb9219f0b1aac23fbc1bf7d038ee6ea666f331fa273031aa2/einops-0.8.2.tar.gz", hash = "sha256:609da665570e5e265e27283aab09e7f279ade90c4f01bcfca111f3d3e13f2827", size = 56261, upload-time = "2026-01-26T04:13:17.638Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl", hash = "sha256:54058201ac7087911181bfec4af6091bb59380360f069276601256a76af08193", size = 65638, upload-time = "2026-01-26T04:13:18.546Z" }, +] + [[package]] name = "fastjsonschema" version = "2.21.2" @@ -3259,7 +3291,7 @@ name = "nvidia-cudnn-cu13" version = "9.19.0.56" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "python_full_version >= '3.10'" }, + { name = "nvidia-cublas", marker = "(python_full_version == '3.10.*' and sys_platform == 'emscripten') or (python_full_version == '3.10.*' and sys_platform == 'win32') or (python_full_version >= '3.10' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f1/84/26025437c1e6b61a707442184fa0c03d083b661adf3a3eecfd6d21677740/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:6ed29ffaee1176c612daf442e4dd6cfeb6a0caa43ddcbeb59da94953030b1be4", size = 433781201, upload-time = "2026-02-03T20:40:53.805Z" }, @@ -3271,7 +3303,7 @@ name = "nvidia-cufft" version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "python_full_version >= '3.10'" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version == '3.10.*' and sys_platform == 'emscripten') or (python_full_version == '3.10.*' and sys_platform == 'win32') or (python_full_version >= '3.10' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, @@ -3361,9 +3393,9 @@ name = "nvidia-cusolver" version = "12.0.4.66" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "python_full_version >= '3.10'" }, - { name = "nvidia-cusparse", marker = "python_full_version >= '3.10'" }, - { name = "nvidia-nvjitlink", marker = "python_full_version >= '3.10'" }, + { name = "nvidia-cublas", marker = "(python_full_version == '3.10.*' and sys_platform == 'emscripten') or (python_full_version == '3.10.*' and sys_platform == 'win32') or (python_full_version >= '3.10' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cusparse", marker = "(python_full_version == '3.10.*' and sys_platform == 'emscripten') or (python_full_version == '3.10.*' and sys_platform == 'win32') or (python_full_version >= '3.10' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version == '3.10.*' and sys_platform == 'emscripten') or (python_full_version == '3.10.*' and sys_platform == 'win32') or (python_full_version >= '3.10' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, @@ -3408,7 +3440,7 @@ name = "nvidia-cusparse" version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "python_full_version >= '3.10'" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version == '3.10.*' and sys_platform == 'emscripten') or (python_full_version == '3.10.*' and sys_platform == 'win32') or (python_full_version >= '3.10' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, @@ -3623,10 +3655,10 @@ resolution-markers = [ "python_full_version >= '3.11' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "flatbuffers", marker = "python_full_version >= '3.11'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "packaging", marker = "python_full_version >= '3.11'" }, - { name = "protobuf", version = "7.34.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "flatbuffers", marker = "python_full_version >= '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "packaging", marker = "python_full_version >= '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "protobuf", version = "7.34.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/d4/81/29a9eb470994a75eb7b3ccf32be314d7c66675a00ac7b50294816cc2db27/onnxruntime-1.26.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:ee1109ef4ef27cad90e823399e61e03b3c6c7bfe0fb820b4baf3678c15be8b3c", size = 18005108, upload-time = "2026-05-08T19:08:11.728Z" }, @@ -5828,6 +5860,8 @@ name = "stereodemo" version = "0.6.2" source = { editable = "." } dependencies = [ + { name = "einops", version = "0.8.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.9'" }, + { name = "einops", version = "0.8.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.9'" }, { name = "numpy", version = "1.24.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.9'" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.9.*'" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, @@ -5866,6 +5900,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "depthai", marker = "extra == 'oak'" }, + { name = "einops" }, { name = "numpy" }, { name = "onnxruntime", marker = "sys_platform == 'darwin'", specifier = ">=1.10.0" }, { name = "onnxruntime-gpu", marker = "sys_platform != 'darwin'", specifier = ">=1.10.0" },