modeling_dpt.py

"""TIPSv2 DPT dense prediction model for HuggingFace."""

from dataclasses import dataclass
from typing import Optional

import torch
from transformers import AutoConfig, AutoModel, PreTrainedModel

from .configuration_dpt import TIPSv2DPTConfig
from .dpt_head import DPTDepthHead, DPTNormalsHead, DPTSegmentationHead


@dataclass
class TIPSv2DPTOutput:
    depth: Optional[torch.Tensor] = None
    normals: Optional[torch.Tensor] = None
    segmentation: Optional[torch.Tensor] = None


class TIPSv2DPTModel(PreTrainedModel):
    """TIPSv2 DPT dense prediction model (depth, normals, segmentation).

    The backbone is loaded automatically from the base TIPSv2 model repo.

    Usage::

        model = AutoModel.from_pretrained("google/tipsv2-l14-dpt", trust_remote_code=True)
        model.eval().cuda()

        outputs = model(pixel_values)
        outputs.depth          # (B, 1, H, W)
        outputs.normals        # (B, 3, H, W)
        outputs.segmentation   # (B, 150, H, W)

        # Individual tasks
        depth = model.predict_depth(pixel_values)
        normals = model.predict_normals(pixel_values)
        seg = model.predict_segmentation(pixel_values)
    """

    config_class = TIPSv2DPTConfig
    _no_split_modules = []
    _supports_cache_class = False
    _tied_weights_keys = []

    @property
    def all_tied_weights_keys(self):
        return {}

    def __init__(self, config: TIPSv2DPTConfig):
        super().__init__(config)

        ppc = tuple(config.post_process_channels)

        backbone_config = AutoConfig.from_pretrained(config.backbone_repo, trust_remote_code=True)
        backbone = AutoModel.from_config(backbone_config, trust_remote_code=True)
        self.vision_encoder = backbone.vision_encoder

        self.depth_head = DPTDepthHead(
            input_embed_dim=config.embed_dim, channels=config.channels,
            post_process_channels=ppc, readout_type=config.readout_type,
            num_depth_bins=config.num_depth_bins,
            min_depth=config.min_depth, max_depth=config.max_depth,
        )
        self.normals_head = DPTNormalsHead(
            input_embed_dim=config.embed_dim, channels=config.channels,
            post_process_channels=ppc, readout_type=config.readout_type,
        )
        self.segmentation_head = DPTSegmentationHead(
            input_embed_dim=config.embed_dim, channels=config.channels,
            post_process_channels=ppc, readout_type=config.readout_type,
            num_classes=config.num_seg_classes,
        )

    def _extract_intermediate(self, pixel_values):
        intermediate = self.vision_encoder.get_intermediate_layers(
            pixel_values, n=self.config.block_indices,
            reshape=True, return_class_token=True, norm=True,
        )
        return [(cls_tok, patch_feat) for patch_feat, cls_tok in intermediate]

    @torch.no_grad()
    def predict_depth(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Predict depth map. Returns (B, 1, H, W)."""
        pixel_values = pixel_values.to(self.device)
        h, w = pixel_values.shape[2:]
        dpt_inputs = self._extract_intermediate(pixel_values)
        return self.depth_head(dpt_inputs, image_size=(h, w))

    @torch.no_grad()
    def predict_normals(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Predict surface normals. Returns (B, 3, H, W)."""
        pixel_values = pixel_values.to(self.device)
        h, w = pixel_values.shape[2:]
        dpt_inputs = self._extract_intermediate(pixel_values)
        return self.normals_head(dpt_inputs, image_size=(h, w))

    @torch.no_grad()
    def predict_segmentation(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Predict semantic segmentation (ADE20K). Returns (B, 150, H, W)."""
        pixel_values = pixel_values.to(self.device)
        h, w = pixel_values.shape[2:]
        dpt_inputs = self._extract_intermediate(pixel_values)
        return self.segmentation_head(dpt_inputs, image_size=(h, w))

    def forward(self, pixel_values: torch.Tensor) -> TIPSv2DPTOutput:
        """Run all three tasks. Returns TIPSv2DPTOutput."""
        pixel_values = pixel_values.to(self.device)
        h, w = pixel_values.shape[2:]
        dpt_inputs = self._extract_intermediate(pixel_values)
        return TIPSv2DPTOutput(
            depth=self.depth_head(dpt_inputs, image_size=(h, w)),
            normals=self.normals_head(dpt_inputs, image_size=(h, w)),
            segmentation=self.segmentation_head(dpt_inputs, image_size=(h, w)),
        )