Spaces:
Running on Zero
Running on Zero
File size: 4,413 Bytes
41ff959 1db10da 41ff959 1db10da 41ff959 1db10da 41ff959 1db10da 41ff959 1db10da 41ff959 1db10da 41ff959 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | from __future__ import annotations
from typing import Dict
import torch
from torch import nn
from .depth_pro import DepthProConfig, create_model_and_transforms
def _inference_precision(device: torch.device) -> torch.dtype:
return torch.float16 if device.type == "cuda" else torch.float32
def _as_depth_map(depth: torch.Tensor, batch_size: int) -> torch.Tensor:
if depth.ndim == 2:
depth = depth.unsqueeze(0)
if depth.ndim != 3:
raise ValueError(f"Expected depth with 2 or 3 dims, got shape {tuple(depth.shape)}")
if depth.shape[0] != batch_size:
raise ValueError(
f"Depth batch mismatch: expected {batch_size}, got {depth.shape[0]}"
)
return depth.unsqueeze(1)
def _as_normalized_intrinsics_batch(
intrinsics: torch.Tensor | None,
batch_size: int,
device: torch.device,
) -> torch.Tensor | None:
if intrinsics is None:
return None
intrinsics = intrinsics.to(device=device, dtype=torch.float32)
if intrinsics.ndim == 2:
intrinsics = intrinsics.unsqueeze(0)
if intrinsics.ndim != 3 or intrinsics.shape[-2:] != (3, 3):
raise ValueError(
f"Expected normalized intrinsics with shape [B,3,3] or [3,3], got {tuple(intrinsics.shape)}"
)
if intrinsics.shape[0] == 1 and batch_size > 1:
intrinsics = intrinsics.expand(batch_size, -1, -1)
if intrinsics.shape[0] != batch_size:
raise ValueError(
f"Intrinsics batch mismatch: expected {batch_size}, got {intrinsics.shape[0]}"
)
return intrinsics
class DepthPro(nn.Module):
"""InfiniSplat wrapper around the vendored upstream DepthPro package."""
def __init__(self):
super().__init__()
self.device_hint = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.precision = _inference_precision(self.device_hint)
depthpro_cfg = DepthProConfig(
patch_encoder_preset="dinov2l16_384",
image_encoder_preset="dinov2l16_384",
decoder_features=256,
checkpoint_uri=None,
fov_encoder_preset="dinov2l16_384",
use_fov_head=True,
)
self.model, self.transform = create_model_and_transforms(
config=depthpro_cfg,
device=self.device_hint,
precision=self.precision,
)
self.model.eval()
def _prepare_input(self, image: torch.Tensor) -> torch.Tensor:
if image.ndim != 4:
raise ValueError(f"Expected image tensor [B,3,H,W], got shape {tuple(image.shape)}")
self._ensure_runtime_dtype()
model_device = next(self.model.parameters()).device
image = image.to(device=model_device, dtype=torch.float32)
image = image.clamp(0.0, 1.0)
image = 2.0 * image - 1.0
return image.to(dtype=self.precision)
def _ensure_runtime_dtype(self) -> None:
param = next(self.model.parameters())
desired_dtype = _inference_precision(param.device)
if param.dtype != desired_dtype:
self.model.to(dtype=desired_dtype)
self.precision = desired_dtype
def forward(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
"""Run DepthPro with the standardized InfiniSplat depth-only interface.
Args:
batch: Dictionary containing:
- image: Tensor with shape [B, 3, H, W].
- intrinsics: Optional normalized intrinsics with shape [B, 3, 3]
or [3, 3]. When provided, fx is converted to pixel-space focal
length and passed to DepthPro. If missing, DepthPro falls back
to its internal FoV / focal estimation path.
Returns:
Dense metric depth with shape [B, 1, H, W].
"""
image = batch["image"]
prepared = self._prepare_input(image)
batch_size = prepared.shape[0]
focal_px: torch.Tensor | None = None
intrinsics = _as_normalized_intrinsics_batch(
batch.get("intrinsics"),
batch_size,
prepared.device,
)
if intrinsics is not None:
focal_px = intrinsics[:, 0, 0] * float(image.shape[-1])
with torch.no_grad():
prediction = self.model.infer(prepared, f_px=focal_px)
return _as_depth_map(prediction["depth"], batch_size).to(dtype=torch.float32)
|