File size: 4,891 Bytes
e605ee8 | 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 | """Camera Depth Model (ByteDance-Seed/manip-as-in-sim-suite @ 757b2ea2) as a transformers depth-estimation model.
``CameraDepthModelForDepthEstimation`` holds the submodules of the official ``RGBDDepth`` under the same names
(``pretrained``: the RGB DINOv2, ``depth_pretrained``: the depth DINOv2, ``crossAtts``, ``depth_head_rgbd``), so the
released checkpoints load by name, and computes depth with that class's own ``forward`` on the RGB + inverse-depth
input, returned as the standard ``DepthEstimatorOutput`` (``predicted_depth``: inverse depth). The official
``infer_image`` / ``image2tensor`` helpers are kept.
"""
import torch
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import DepthEstimatorOutput
from .cdm_dpt import RGBDDepth
from .configuration_camera_depth_model import CameraDepthModelConfig
class CameraDepthModelPreTrainedModel(PreTrainedModel):
config_class = CameraDepthModelConfig
config: CameraDepthModelConfig
base_model_prefix = "camera_depth_model"
main_input_name = "pixel_values"
class CameraDepthModelForDepthEstimation(CameraDepthModelPreTrainedModel):
def __init__(self, config: CameraDepthModelConfig):
super().__init__(config)
net = RGBDDepth(encoder=config.encoder, features=config.features, out_channels=list(config.out_channels),
use_bn=config.use_bn, use_clstoken=config.use_clstoken, max_depth=config.max_depth)
# the official module's attributes, so its methods run on this object unchanged
self.intermediate_layer_idx = net.intermediate_layer_idx
self.max_depth = net.max_depth
self.encoder = net.encoder
self.pretrained = net.pretrained
self.depth_pretrained = net.depth_pretrained
self.depth_head_rgbd = net.depth_head_rgbd
self.crossAtts = net.crossAtts
self.post_init()
def forward(
self,
pixel_values: torch.Tensor,
prompt_depth: torch.Tensor | None = None,
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; prompt_depth (B, 1, H, W): the
camera's inverse depth (1 / meters, 0 where invalid) at the same size; both from ``CameraDepthModelProcessor``.
``predicted_depth`` (B, H, W): inverse depth (meters = 1 / predicted_depth, as infer.py).
``output_hidden_states``: the DINOv2 features of the RGB and the depth branch (per layer, (patch tokens, class
token)), RGB layers first."""
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)")
if prompt_depth is None:
raise ValueError("prompt_depth (the camera's inverse depth) is required")
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
x = torch.cat((pixel_values, prompt_depth), dim=1) # image2tensor
predicted_depth = RGBDDepth.forward(self, x)
hidden_states = None
if output_hidden_states:
layers = self.intermediate_layer_idx[self.encoder]
hidden_states = tuple(self.pretrained.get_intermediate_layers(pixel_values, layers, return_class_token=True)) + tuple(
self.depth_pretrained.get_intermediate_layers(prompt_depth.repeat(1, 3, 1, 1), layers, 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)
# official helpers: RGB uint8 image + inverse depth (numpy) -> inverse depth at the image size (numpy)
@torch.no_grad()
def infer_image(self, raw_image, depth_low_res, input_size: int = 518):
# body of RGBDDepth.infer_image, calling the official forward (self.forward returns a DepthEstimatorOutput)
inputs, (h, w) = self.image2tensor(raw_image, depth_low_res, input_size)
pred_depth = RGBDDepth.forward(self, inputs)
pred_depth = F.interpolate(pred_depth[:, None], (h, w), mode="nearest")[0, 0]
return pred_depth.cpu().numpy()
def image2tensor(self, raw_image, depth, input_size: int = 518):
return RGBDDepth.image2tensor(self, raw_image, depth, input_size)
__all__ = ["CameraDepthModelForDepthEstimation", "CameraDepthModelPreTrainedModel"]
|