| """Depth Anything V2 (DepthAnything/Depth-Anything-V2 @ a561b849) as a transformers depth-estimation model. |
| |
| ``DepthAnythingV2ForDepthEstimation`` holds the submodules of the official ``DepthAnythingV2`` under the same names |
| (``pretrained``: DINOv2, ``depth_head``: DPT head), so the released ``depth_anything_v2_*.pth`` tensors load by name, |
| and computes depth with that class's own ``forward`` -- the relative-depth class (``dav2_dpt.py``), or with |
| ``config.max_depth`` set the metric-depth class (``dav2_metric_dpt.py``, metric_depth/ of the repo). It returns the |
| standard ``DepthEstimatorOutput`` (``predicted_depth``). The official ``infer_image`` / ``image2tensor`` helpers are |
| kept. Unlike transformers' ``DepthAnythingForDepthEstimation``, the DINOv2 position-embedding interpolation is the |
| original one, so non-square inputs give the original outputs. |
| """ |
|
|
| import torch |
| import torch.nn.functional as F |
| from transformers import PreTrainedModel |
| from transformers.modeling_outputs import DepthEstimatorOutput |
|
|
| from .configuration_depth_anything_v2 import DepthAnythingV2Config |
| from .dav2_dpt import DepthAnythingV2 as RelativeDepthAnythingV2 |
| from .dav2_metric_dpt import DepthAnythingV2 as MetricDepthAnythingV2 |
|
|
|
|
| class DepthAnythingV2PreTrainedModel(PreTrainedModel): |
| config_class = DepthAnythingV2Config |
| config: DepthAnythingV2Config |
| base_model_prefix = "depth_anything_v2" |
| main_input_name = "pixel_values" |
|
|
|
|
| class DepthAnythingV2ForDepthEstimation(DepthAnythingV2PreTrainedModel): |
| def __init__(self, config: DepthAnythingV2Config): |
| super().__init__(config) |
| kwargs = dict(encoder=config.encoder, features=config.features, out_channels=list(config.out_channels), |
| use_bn=config.use_bn, use_clstoken=config.use_clstoken) |
| if config.max_depth is None: |
| self._official = RelativeDepthAnythingV2 |
| net = RelativeDepthAnythingV2(**kwargs) |
| else: |
| self._official = MetricDepthAnythingV2 |
| net = MetricDepthAnythingV2(**kwargs, max_depth=config.max_depth) |
| self.max_depth = net.max_depth |
| |
| self.intermediate_layer_idx = net.intermediate_layer_idx |
| self.encoder = net.encoder |
| self.pretrained = net.pretrained |
| self.depth_head = net.depth_head |
| self.post_init() |
|
|
| def forward( |
| self, |
| pixel_values: torch.Tensor, |
| labels: torch.Tensor | None = None, |
| output_attentions: bool | None = None, |
| output_hidden_states: bool | None = None, |
| return_dict: bool | None = None, |
| **kwargs, |
| ) -> tuple | DepthEstimatorOutput: |
| """pixel_values (B, 3, H, W), H and W multiples of 14, ImageNet-normalised (``DepthAnythingV2Processor``). |
| ``predicted_depth`` (B, H, W): relative inverse depth, or metric depth in meters when ``config.max_depth`` |
| is set. ``output_hidden_states``: the DINOv2 features the DPT head reads (per layer, patch tokens).""" |
| if labels is not None: |
| raise NotImplementedError("Training is not implemented") |
| if output_attentions: |
| raise NotImplementedError("Attention weights are not returned (the original model does not expose them)") |
| return_dict = return_dict if return_dict is not None else self.config.return_dict |
| output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
| predicted_depth = self._official.forward(self, pixel_values) |
| hidden_states = None |
| if output_hidden_states: |
| features = self.pretrained.get_intermediate_layers(pixel_values, self.intermediate_layer_idx[self.encoder]) |
| hidden_states = tuple(features) |
| if not return_dict: |
| return (predicted_depth,) + ((hidden_states,) if hidden_states is not None else ()) |
| return DepthEstimatorOutput(predicted_depth=predicted_depth, hidden_states=hidden_states) |
|
|
| |
| @torch.no_grad() |
| def infer_image(self, raw_image, input_size: int = 518): |
| |
| image, (h, w) = self.image2tensor(raw_image, input_size) |
| depth = self._official.forward(self, image) |
| depth = F.interpolate(depth[:, None], (h, w), mode="bilinear", align_corners=True)[0, 0] |
| return depth.cpu().numpy() |
|
|
| def image2tensor(self, raw_image, input_size: int = 518): |
| return self._official.image2tensor(self, raw_image, input_size) |
|
|
|
|
| __all__ = ["DepthAnythingV2ForDepthEstimation", "DepthAnythingV2PreTrainedModel"] |
|
|
