src/model/encoder/depth/infinidepth/implicit_pda.py

import os
import torch
import torch.nn as nn
import sys
import numpy as np
from pathlib import Path
from sklearn.linear_model import RANSACRegressor
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import make_pipeline
sys.path.append(str(Path(__file__).parent.parent.parent))
sys.path.append(str(Path(__file__).parent))

from src.model.encoder.depth.infinidepth.implicit_promptda_simple.config import (
    dinov3_model_configs,
    model_configs,
)
from src.model.encoder.depth.infinidepth.implicit_promptda_simple.low_level_implicit import (
    LowLevelImplicitHead,
)
from src.model.encoder.depth.infinidepth.implicit_promptda_simple.prompt_models import (
    GeneralPromptModel,
    SelfAttnPromptModel,
)
from src.model.encoder.depth.infinidepth.implicit_promptda_simple.convolution import BasicEncoder
from src.model.encoder.depth.infinidepth.warp_utils import WarpMedian
from typing import Optional

EPS = 1e-6
acc_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16


def _clear_torchhub_package(package_name: str) -> None:
    stale_modules = [
        module_name
        for module_name in list(sys.modules)
        if module_name == package_name or module_name.startswith(f"{package_name}.")
    ]
    for module_name in stale_modules:
        sys.modules.pop(module_name, None)


class InfiniDepth(nn.Module):
    use_bn = False
    use_clstoken = False

    def __init__(self,
                 model_path: Optional[str] = None,
                 geometry_type: str = "disparity",
                 encoder: str = "vitl16",  # dinov2 -> vitl  dinov3 -> vitl16
                 backbone: str = "dinov3",
                 use_prompt=True,
                 ):
        super().__init__()
        if backbone != "dinov3":
            model_config = model_configs[encoder]
        else:
            model_config = dinov3_model_configs[encoder]
        self.model_config = model_config
        self.use_prompt = use_prompt

        # Learnable module definitions
        # backbone
        if backbone == "dinov2":
            raise NotImplementedError("Dinov2 backbone is not implemented in this version.")
        elif backbone == "dinov3":
            _clear_torchhub_package("dinov3")
            self.pretrained = torch.hub.load(
                "src/model/encoder/depth/infinidepth/implicit_promptda_simple/torchhub/dinov3",
                f"dinov3_{encoder}",  # vitl16, vith16plus, vit7b16
                source="local",
                pretrained=False,
            )
            _clear_torchhub_package("dinov3")
            self.patch_size = 16

        dim = self.pretrained.blocks[0].attn.qkv.in_features

        self.basic_encoder = BasicEncoder(
            input_dim=3,
            output_dim=128,
            stride=4
        )
        
        # prompt
        if use_prompt:
            self.prompt_model = GeneralPromptModel(prompt_stage=[3], block=SelfAttnPromptModel(num_blocks=4, pe="qk"))
            self.warp_func = WarpMedian()
        else:
            poly_features = PolynomialFeatures(degree=1, include_bias=False)
            ransac = RANSACRegressor(max_trials=1000)
            self.align_model = make_pipeline(poly_features, ransac)
            self.scale = 1.0
            self.shift = 0.0

        # implicit depth decoder head
        self.depth_implicit_head = LowLevelImplicitHead(                
            hidden_dim=dim,
            basic_dim=128,  # BasicEncoder output dim
            fusion_type="concat",  # concat, gated
            out_dim=1,
            hidden_list=[1024, 256, 32],
        )

        self.register_buffer("_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer("_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
        self.geometry_type = geometry_type

        if model_path is not None:
            if os.path.exists(model_path):
                checkpoint = torch.load(model_path, map_location="cpu")
                self.load_state_dict({k[9:]: v for k, v in checkpoint["state_dict"].items()})
            else:
                raise FileNotFoundError(f"Model file {model_path} not found")

        # The wrapper owns device placement and train/eval mode.

    @torch.no_grad() 
    def inference(self,
                image: torch.Tensor,
                query_coord: torch.Tensor,
                prompt_depth: torch.Tensor,
                prompt_mask=None,
                use_batch_infer=True,
                align_mode: str = "fit_and_apply"):
        if prompt_mask is None:
            prompt_mask = prompt_depth > 0
        
        if self.use_prompt:
            prompt_depth, prompt_mask, reference_meta = self.warp_func.warp(
                prompt_depth,
                prompt_depth=prompt_depth,
                prompt_mask=prompt_mask,
            )

        if use_batch_infer:
            pred, dino_feat, basic_feat = self.batch_forward(image, query_coord, prompt_depth, prompt_mask, bsize=100000)
        else:
            pred, dino_feat, basic_feat = self.forward(image, query_coord, prompt_depth, prompt_mask)

        if self.use_prompt:
            pred = self.warp_func.unwarp(
                pred,
                reference_meta=reference_meta[...,0],
            )
        else:
            if align_mode == "fit_and_apply":
                pred = self._ransac_align_depth(pred, prompt_depth, prompt_mask, update_params=True).to(image.device)
            elif align_mode == "apply_cached":
                pred = self._apply_scale_shift(pred, self.scale, self.shift)
            elif align_mode == "none":
                pass
            else:
                raise ValueError(f"Unknown align_mode: {align_mode}")

        if self.geometry_type == "depth":
            pred_depth = pred
            pred_disparity = 1.0 / torch.clamp(pred, min=5e-3)
        elif self.geometry_type == "disparity":
            pred_disparity = pred
            pred_depth = 1.0 / torch.clamp(pred, min=5e-3)

        return pred_depth, pred_disparity, dino_feat, basic_feat

    def batch_forward(self, x, coord, prompt_depth=None, prompt_mask=None, bsize=3000):
        """
        Forward pass with batching to avoid OOM.
        """
        h, w = x.shape[-2:]

        # DINOv3 branch (semantic features) - uses ImageNet normalization
        x_dino = (x - self._mean) / self._std
        with torch.autocast("cuda", enabled=True, dtype=acc_dtype):
            features = self.pretrained.get_intermediate_layers(
                x_dino,
                n=self.model_config["layer_idxs"],
                return_class_token=True,
            )
        dino_feat = features[-1][0].clone()
        features = [list(feature) for feature in features]
        patch_h, patch_w = h // self.patch_size, w // self.patch_size
        
        if self.use_prompt:
            features = self.prompt_model(features, prompt_depth, prompt_mask, patch_h, patch_w)

        # BasicEncoder branch (low-level features) - uses [-1, 1] normalization
        x_basic = 2.0 * x - 1.0
        basic_feat = self.basic_encoder(x_basic)  # [B, 128, H/4, W/4]
  
        # generate feature map for learning implicit function
        feat = self.depth_implicit_head.encode_feat(features, patch_h, patch_w)
        n = coord.shape[1]
        ql = 0
        preds = []
        while ql < n:
            qr = min(ql + bsize, n)
            # batch querying
            pred = self.depth_implicit_head.decode_dpt(feat, basic_feat, coord[:, ql: qr, :], cell=None)
            preds.append(pred)
            ql = qr
        pred = torch.cat(preds, dim=1)
        return pred, dino_feat, basic_feat

    def forward(self, x, coords, prompt_depth, prompt_mask):
        h, w = x.shape[-2:]

        # DINOv3 branch (semantic features) - uses ImageNet normalization
        x_dino = (x - self._mean) / self._std
        with torch.autocast("cuda", enabled=True, dtype=acc_dtype):
            features = self.pretrained.get_intermediate_layers(
                x_dino,
                n=self.model_config["layer_idxs"],
                return_class_token=True,
            )
        dino_feat = features[-1][0].clone()
        features = [list(feature) for feature in features]
        patch_h, patch_w = h // self.patch_size, w // self.patch_size

        if self.use_prompt:
            features = self.prompt_model(features, prompt_depth, prompt_mask, patch_h, patch_w)

        # BasicEncoder branch (low-level features) - uses [-1, 1] normalization
        x_basic = 2.0 * x - 1.0
        basic_feat = self.basic_encoder(x_basic)  # [B, 128, H/4, W/4]

        with torch.autocast("cuda", enabled=True, dtype=torch.float32):
            depth = self.depth_implicit_head(features, basic_feat, patch_h, patch_w, coords, cell=None)

        return depth, dino_feat, basic_feat

    def _apply_scale_shift(self, pred, scale, shift):
        if type(pred).__module__ == torch.__name__:
            return pred * float(scale) + float(shift)

        pred_np = np.asarray(pred, dtype=np.float32)
        return pred_np * float(scale) + float(shift)

    def _ransac_align_depth(self, pred, gt, mask0=None, geometry_type=None, update_params=True):
        if geometry_type is None:
            geometry_type = self.geometry_type

        if type(pred).__module__ == torch.__name__:
            pred = pred.detach().cpu().numpy()
        if type(gt).__module__ == torch.__name__:
            gt = gt.detach().cpu().numpy()

        pred = np.asarray(pred, dtype=np.float32)
        gt = np.asarray(gt, dtype=np.float32)

        # Keep the expected tensor layout as (B, N, C) and avoid squeeze-induced shape changes.
        if pred.ndim == 1:
            pred = pred[None, :, None]
        elif pred.ndim == 2:
            pred = pred[:, :, None]
        elif pred.ndim != 3:
            raise ValueError(f"Expected pred with 1/2/3 dims, but got shape {pred.shape}")

        bsz = pred.shape[0]

        if gt.ndim == 0:
            gt_flat = gt.reshape(1, 1)
        elif gt.ndim == 1:
            gt_flat = gt.reshape(1, -1)
        else:
            gt_flat = gt.reshape(gt.shape[0], -1)

        mask_flat = None
        if mask0 is not None:
            if type(mask0).__module__ == torch.__name__:
                mask0 = mask0.detach().cpu().numpy()
            mask0 = np.asarray(mask0)
            if mask0.ndim == 0:
                mask_flat = (mask0.reshape(1, 1) > 0)
            elif mask0.ndim == 1:
                mask_flat = (mask0.reshape(1, -1) > 0)
            else:
                mask_flat = (mask0.reshape(mask0.shape[0], -1) > 0)

        if gt_flat.shape[0] == 1 and bsz > 1:
            gt_flat = np.repeat(gt_flat, bsz, axis=0)
        if mask_flat is not None and mask_flat.shape[0] == 1 and bsz > 1:
            mask_flat = np.repeat(mask_flat, bsz, axis=0)

        if gt_flat.shape[0] != bsz:
            bsz = min(bsz, gt_flat.shape[0])
            pred = pred[:bsz]
            gt_flat = gt_flat[:bsz]
            if mask_flat is not None:
                mask_flat = mask_flat[:bsz]
        elif mask_flat is not None and mask_flat.shape[0] != bsz:
            bsz = min(bsz, mask_flat.shape[0])
            pred = pred[:bsz]
            gt_flat = gt_flat[:bsz]
            mask_flat = mask_flat[:bsz]

        pred_metric = pred.copy()
        scales = []
        shifts = []

        for bi in range(bsz):
            pred_vec = pred[bi, :, 0].reshape(-1)
            gt_vec = gt_flat[bi].reshape(-1)
            n = min(pred_vec.shape[0], gt_vec.shape[0])

            curr_mask = None
            if mask_flat is not None:
                n = min(n, mask_flat[bi].shape[0])
                curr_mask = mask_flat[bi][:n]

            scale = 1.0
            shift = 0.0
            if n >= 2:
                pred_n = pred_vec[:n]
                gt_n = gt_vec[:n]
                valid = gt_n > 1e-8
                if curr_mask is not None:
                    valid = valid & curr_mask

                if valid.sum() >= 2:
                    gt_mask = gt_n[valid].astype(np.float32)
                    pred_mask = pred_n[valid].astype(np.float32)

                    if geometry_type == "disparity":
                        gt_mask = np.clip(gt_mask, 1e-8, None)
                    if geometry_type == "depth":
                        gt_mask = np.log(gt_mask + 1.0)

                    try:
                        self.align_model.fit(pred_mask[:, None], gt_mask[:, None])
                        a = self.align_model.named_steps['ransacregressor'].estimator_.coef_.item()
                        b = self.align_model.named_steps['ransacregressor'].estimator_.intercept_.item()
                    except Exception:
                        a, b = 1.0, 0.0

                    if a > 0:
                        scale, shift = a, b
                    else:
                        pred_mean = np.mean(pred_mask)
                        gt_mean = np.mean(gt_mask)
                        scale = gt_mean / (pred_mean + EPS)
                        shift = 0.0

            pred_metric[bi, :n, :] = self._apply_scale_shift(pred[bi, :n, :], scale, shift)
            scales.append(float(scale))
            shifts.append(float(shift))

        if update_params and not self.use_prompt and len(scales) > 0:
            self.scale = float(np.mean(scales))
            self.shift = float(np.mean(shifts))

        return torch.from_numpy(pred_metric)