| """Video Depth Anything (DepthAnything/Video-Depth-Anything @ 4f5ae231) as a transformers depth-estimation model. |
| |
| ``VideoDepthAnythingForDepthEstimation`` holds the submodules of the official ``VideoDepthAnything`` under the same |
| names (``pretrained``: DINOv2, ``head``: temporal DPT head), so the released ``*video_depth_anything_*.pth`` tensors |
| load by name, and computes depth with that class's own ``forward`` (``vda_video_depth.py``), returned as the standard |
| ``DepthEstimatorOutput`` (``predicted_depth`` (B, T, H, W)). The official whole-video inference (``infer_video_depth``: |
| 32-frame windows, key-frame overlap, scale / shift alignment) and the streaming inference (``infer_video_depth_one``, |
| ``vda_video_depth_stream.py``) run as they are on this model. |
| """ |
|
|
| import torch |
| from transformers import PreTrainedModel |
| from transformers.modeling_outputs import DepthEstimatorOutput |
|
|
| from .configuration_video_depth_anything import VideoDepthAnythingConfig |
| from .vda_video_depth import VideoDepthAnything |
| from .vda_video_depth_stream import INFER_LEN, INTERP_LEN, OVERLAP |
| from .vda_video_depth_stream import VideoDepthAnything as StreamingVideoDepthAnything |
|
|
|
|
| class _OfficialView: |
| """The model as an instance of an official class: that class's methods (``forward`` included) bound to this view, |
| other attributes read from the model; attributes the methods set (the streaming caches) stay on the view.""" |
|
|
| def __init__(self, model, official): |
| object.__setattr__(self, "_model", model) |
| object.__setattr__(self, "_official", official) |
|
|
| def __getattr__(self, name): |
| if name in vars(self._official): |
| return vars(self._official)[name].__get__(self) |
| return getattr(self._model, name) |
|
|
|
|
| class VideoDepthAnythingPreTrainedModel(PreTrainedModel): |
| config_class = VideoDepthAnythingConfig |
| config: VideoDepthAnythingConfig |
| base_model_prefix = "video_depth_anything" |
| main_input_name = "pixel_values" |
|
|
|
|
| class VideoDepthAnythingForDepthEstimation(VideoDepthAnythingPreTrainedModel): |
| def __init__(self, config: VideoDepthAnythingConfig): |
| super().__init__(config) |
| net = VideoDepthAnything(encoder=config.encoder, features=config.features, out_channels=list(config.out_channels), |
| use_bn=config.use_bn, use_clstoken=config.use_clstoken, num_frames=config.num_frames, |
| pe=config.pe, metric=config.metric) |
| |
| self.intermediate_layer_idx = net.intermediate_layer_idx |
| self.encoder = net.encoder |
| self.metric = net.metric |
| self.pretrained = net.pretrained |
| self.head = net.head |
| self._stream = None |
| 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, T, 3, H, W), H and W multiples of 14, ImageNet-normalised (``VideoDepthAnythingProcessor``); |
| T up to ``config.num_frames``. ``predicted_depth`` (B, T, H, W): relative inverse depth, or metric depth in |
| meters for the metric checkpoints. ``output_hidden_states``: the DINOv2 features the head reads (per layer, |
| (patch tokens, class token) of the B*T frames).""" |
| 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 = VideoDepthAnything.forward(self, pixel_values) |
| hidden_states = None |
| if output_hidden_states: |
| hidden_states = tuple(self.pretrained.get_intermediate_layers( |
| pixel_values.flatten(0, 1), self.intermediate_layer_idx[self.encoder], return_class_token=True)) |
| 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) |
|
|
| |
| def infer_video_depth(self, frames, target_fps, input_size: int = 518, device: str = "cuda", fp32: bool = False): |
| return VideoDepthAnything.infer_video_depth(_OfficialView(self, VideoDepthAnything), frames, target_fps, |
| input_size=input_size, device=device, fp32=fp32) |
|
|
| def infer_video_depth_one(self, frame, input_size: int = 518, device: str = "cuda", fp32: bool = False): |
| """Streaming (experimental upstream): one RGB uint8 frame (H, W, 3) -> its depth; the temporal caches persist |
| across calls until ``reset_streaming``.""" |
| if self._stream is None: |
| self._stream = _OfficialView(self, StreamingVideoDepthAnything) |
| |
| self._stream.transform = None |
| self._stream.frame_id_list = [] |
| self._stream.frame_cache_list = [] |
| self._stream.gap = (INFER_LEN - OVERLAP) * 2 - 1 - (OVERLAP - INTERP_LEN) |
| self._stream.id = -1 |
| return StreamingVideoDepthAnything.infer_video_depth_one(self._stream, frame, input_size=input_size, device=device, fp32=fp32) |
|
|
| def reset_streaming(self): |
| self._stream = None |
|
|
|
|
| __all__ = ["VideoDepthAnythingForDepthEstimation", "VideoDepthAnythingPreTrainedModel"] |
|
|
