File size: 4,523 Bytes
e8b3753 | 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 | """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)),
)
|