Camera-Depth-Model-Base / modeling_camera_depth_model.py

Upload Camera-Depth-Model-Base (transformers port, original code)

169a39f verified 2 days ago

4.89 kB

"""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"]

	"""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"]