Spaces:
Running on Zero
Running on Zero
File size: 1,656 Bytes
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 | from __future__ import annotations
import torch
from einops import rearrange
from torch import nn
from src.model.encoder.depth.infinidepth.implicit_pda import (
InfiniDepth as InfiniDepthModel,
)
from src.model.encoder.depth.infinidepth.sampling_utils import make_2d_uniform_coord
class InfiniDepth(nn.Module):
"""Prompt-conditioned InfiniDepth dense-depth predictor."""
def __init__(self) -> None:
super().__init__()
self.model = InfiniDepthModel(
model_path=None,
geometry_type="disparity",
use_prompt=True,
)
self.model.eval()
@torch.inference_mode()
def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
"""Predict a dense depth map from RGB and sparse disparity prompts."""
required_keys = ("image", "prompt_disparity", "prompt_mask")
missing_keys = [key for key in required_keys if key not in batch]
if missing_keys:
raise AssertionError(
"InfiniDepth requires prompt-conditioned inputs. "
f"Missing keys: {missing_keys}"
)
image = batch["image"]
batch_size, _, height, width = image.shape
query_coords = make_2d_uniform_coord((height, width)).to(image.device)
query_coords = query_coords.unsqueeze(0).expand(batch_size, -1, -1)
depth, _, _, _ = self.model.inference(
image=image,
query_coord=query_coords,
prompt_depth=batch["prompt_disparity"],
prompt_mask=batch["prompt_mask"],
)
return rearrange(depth, "b (h w) c -> b c h w", h=height, w=width)
|