# ruff: noqa # fmt: off """Camera Depth Model (CDM), vendored verbatim from ByteDance-Seed/manip-as-in-sim-suite @ 757b2ea26bcf7992c92ac80485d82b905e961850 (cdm/rgbddepth/): the files below concatenated, package-internal imports replaced by relative imports of the other vendored files (generated by scripts/model_ports/common/vendor_flat.py; do not edit by hand). Code license: Apache-2.0.""" from .cdm_dinov2 import DINOv2 # ============================================================================================================== # rgbddepth/util/blocks.py # ============================================================================================================== #!/usr/bin/env python3 # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: Apache-2.0 import torch.nn as nn def _make_scratch(in_shape, out_shape, groups=1, expand=False): scratch = nn.Module() out_shape1 = out_shape out_shape2 = out_shape out_shape3 = out_shape if len(in_shape) >= 4: out_shape4 = out_shape if expand: out_shape1 = out_shape out_shape2 = out_shape * 2 out_shape3 = out_shape * 4 if len(in_shape) >= 4: out_shape4 = out_shape * 8 scratch.layer1_rn = nn.Conv2d( in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups, ) scratch.layer2_rn = nn.Conv2d( in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups, ) scratch.layer3_rn = nn.Conv2d( in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups, ) if len(in_shape) >= 4: scratch.layer4_rn = nn.Conv2d( in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups, ) return scratch class ResidualConvUnit(nn.Module): """Residual convolution module.""" def __init__(self, features, activation, bn): """Init. Args: features (int): number of features """ super().__init__() self.bn = bn self.groups = 1 self.conv1 = nn.Conv2d( features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups, ) self.conv2 = nn.Conv2d( features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups, ) if self.bn == True: self.bn1 = nn.BatchNorm2d(features) self.bn2 = nn.BatchNorm2d(features) self.activation = activation self.skip_add = nn.quantized.FloatFunctional() def forward(self, x): """Forward pass. Args: x (tensor): input Returns: tensor: output """ out = self.activation(x) out = self.conv1(out) if self.bn == True: out = self.bn1(out) out = self.activation(out) out = self.conv2(out) if self.bn == True: out = self.bn2(out) if self.groups > 1: out = self.conv_merge(out) return self.skip_add.add(out, x) class FeatureFusionBlock(nn.Module): """Feature fusion block.""" def __init__( self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, size=None, ): """Init. Args: features (int): number of features """ super(FeatureFusionBlock, self).__init__() self.deconv = deconv self.align_corners = align_corners self.groups = 1 self.expand = expand out_features = features if self.expand == True: out_features = features // 2 self.out_conv = nn.Conv2d( features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1, ) self.resConfUnit1 = ResidualConvUnit(features, activation, bn) self.resConfUnit2 = ResidualConvUnit(features, activation, bn) self.skip_add = nn.quantized.FloatFunctional() self.size = size def forward(self, *xs, size=None): """Forward pass. Returns: tensor: output """ output = xs[0] if len(xs) == 2: res = self.resConfUnit1(xs[1]) output = self.skip_add.add(output, res) output = self.resConfUnit2(output) if (size is None) and (self.size is None): modifier = {"scale_factor": 2} elif size is None: modifier = {"size": self.size} else: modifier = {"size": size} output = nn.functional.interpolate( output, **modifier, mode="bilinear", align_corners=self.align_corners ) output = self.out_conv(output) return output # ============================================================================================================== # rgbddepth/util/transform.py # ============================================================================================================== #!/usr/bin/env python3 # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: Apache-2.0 import cv2 import numpy as np class Resize(object): """Resize sample to given size (width, height).""" def __init__( self, width, height, resize_target=True, keep_aspect_ratio=False, ensure_multiple_of=1, resize_method="lower_bound", image_interpolation_method=cv2.INTER_AREA, ): """Init. Args: width (int): desired output width height (int): desired output height resize_target (bool, optional): True: Resize the full sample (image, mask, target). False: Resize image only. Defaults to True. keep_aspect_ratio (bool, optional): True: Keep the aspect ratio of the input sample. Output sample might not have the given width and height, and resize behaviour depends on the parameter 'resize_method'. Defaults to False. ensure_multiple_of (int, optional): Output width and height is constrained to be multiple of this parameter. Defaults to 1. resize_method (str, optional): "lower_bound": Output will be at least as large as the given size. "upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.) "minimal": Scale as least as possible. (Output size might be smaller than given size.) Defaults to "lower_bound". """ self.__width = width self.__height = height self.__resize_target = resize_target self.__keep_aspect_ratio = keep_aspect_ratio self.__multiple_of = ensure_multiple_of self.__resize_method = resize_method self.__image_interpolation_method = image_interpolation_method def constrain_to_multiple_of(self, x, min_val=0, max_val=None): y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int) if max_val is not None and y > max_val: y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int) if y < min_val: y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int) return y def get_size(self, width, height): # determine new height and width scale_height = self.__height / height scale_width = self.__width / width if self.__keep_aspect_ratio: if self.__resize_method == "lower_bound": # scale such that output size is lower bound if scale_width > scale_height: # fit width scale_height = scale_width else: # fit height scale_width = scale_height elif self.__resize_method == "upper_bound": # scale such that output size is upper bound if scale_width < scale_height: # fit width scale_height = scale_width else: # fit height scale_width = scale_height elif self.__resize_method == "minimal": # scale as least as possbile if abs(1 - scale_width) < abs(1 - scale_height): # fit width scale_height = scale_width else: # fit height scale_width = scale_height else: raise ValueError( f"resize_method {self.__resize_method} not implemented" ) if self.__resize_method == "lower_bound": new_height = self.constrain_to_multiple_of( scale_height * height, min_val=self.__height ) new_width = self.constrain_to_multiple_of( scale_width * width, min_val=self.__width ) elif self.__resize_method == "upper_bound": new_height = self.constrain_to_multiple_of( scale_height * height, max_val=self.__height ) new_width = self.constrain_to_multiple_of( scale_width * width, max_val=self.__width ) elif self.__resize_method == "minimal": new_height = self.constrain_to_multiple_of(scale_height * height) new_width = self.constrain_to_multiple_of(scale_width * width) else: raise ValueError(f"resize_method {self.__resize_method} not implemented") return (new_width, new_height) def __call__(self, sample): width, height = self.get_size( sample["image"].shape[1], sample["image"].shape[0] ) # resize sample sample["image"] = cv2.resize( sample["image"], (width, height), interpolation=self.__image_interpolation_method, ) if self.__resize_target: if "depth" in sample: sample["depth"] = cv2.resize( sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST ) if "mask" in sample: sample["mask"] = cv2.resize( sample["mask"].astype(np.float32), (width, height), interpolation=cv2.INTER_NEAREST, ) return sample class NormalizeImage(object): """Normlize image by given mean and std.""" def __init__(self, mean, std): self.__mean = mean self.__std = std def __call__(self, sample): sample["image"] = (sample["image"] - self.__mean) / self.__std return sample class PrepareForNet(object): """Prepare sample for usage as network input.""" def __init__(self): pass def __call__(self, sample): image = np.transpose(sample["image"], (2, 0, 1)) sample["image"] = np.ascontiguousarray(image).astype(np.float32) if "depth" in sample: depth = sample["depth"].astype(np.float32) sample["depth"] = np.ascontiguousarray(depth) if "mask" in sample: sample["mask"] = sample["mask"].astype(np.float32) sample["mask"] = np.ascontiguousarray(sample["mask"]) return sample # ============================================================================================================== # rgbddepth/dpt.py # ============================================================================================================== #!/usr/bin/env python3 # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: Apache-2.0 import cv2 import torch import torch.nn as nn import torch.nn.functional as F from torchvision.transforms import Compose def _make_fusion_block(features, use_bn, size=None): return FeatureFusionBlock( features, nn.ReLU(False), deconv=False, bn=use_bn, expand=False, align_corners=True, size=size, ) class ConvBlock(nn.Module): def __init__(self, in_feature, out_feature): super().__init__() self.conv_block = nn.Sequential( nn.Conv2d(in_feature, out_feature, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(out_feature), nn.ReLU(True), ) def forward(self, x): return self.conv_block(x) class DPTHead(nn.Module): def __init__( self, in_channels, features=256, use_bn=False, out_channels=[256, 512, 1024, 1024], use_clstoken=False, sigact_out=False, ): super(DPTHead, self).__init__() self.use_clstoken = use_clstoken self.projects = nn.ModuleList( [ nn.Conv2d( in_channels=in_channels, out_channels=out_channel, kernel_size=1, stride=1, padding=0, ) for out_channel in out_channels ] ) self.resize_layers = nn.ModuleList( [ nn.ConvTranspose2d( in_channels=out_channels[0], out_channels=out_channels[0], kernel_size=4, stride=4, padding=0, ), nn.ConvTranspose2d( in_channels=out_channels[1], out_channels=out_channels[1], kernel_size=2, stride=2, padding=0, ), nn.Identity(), nn.Conv2d( in_channels=out_channels[3], out_channels=out_channels[3], kernel_size=3, stride=2, padding=1, ), ] ) if use_clstoken: self.readout_projects = nn.ModuleList() for _ in range(len(self.projects)): self.readout_projects.append( nn.Sequential(nn.Linear(2 * in_channels, in_channels), nn.GELU()) ) self.scratch = _make_scratch( out_channels, features, groups=1, expand=False, ) self.scratch.stem_transpose = None self.scratch.refinenet1 = _make_fusion_block(features, use_bn) self.scratch.refinenet2 = _make_fusion_block(features, use_bn) self.scratch.refinenet3 = _make_fusion_block(features, use_bn) self.scratch.refinenet4 = _make_fusion_block(features, use_bn) head_features_1 = features head_features_2 = 32 self.scratch.output_conv1 = nn.Conv2d( head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1 ) if not sigact_out: self.scratch.output_conv2 = nn.Sequential( nn.Conv2d( head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1, ), nn.ReLU(True), nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0), nn.ReLU(True), nn.Identity(), ) else: self.scratch.output_conv2 = nn.Sequential( nn.Conv2d( head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1, ), nn.ReLU(True), nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0), nn.Sigmoid(), ) def forward(self, out_features, patch_h, patch_w): out = [] for i, x in enumerate(out_features): if self.use_clstoken: x, cls_token = x[0], x[1] readout = cls_token.unsqueeze(1).expand_as(x) x = self.readout_projects[i](torch.cat((x, readout), -1)) else: x = x[0] x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w)) x = self.projects[i](x) x = self.resize_layers[i](x) out.append(x) layer_1, layer_2, layer_3, layer_4 = out layer_1_rn = self.scratch.layer1_rn(layer_1) layer_2_rn = self.scratch.layer2_rn(layer_2) layer_3_rn = self.scratch.layer3_rn(layer_3) layer_4_rn = self.scratch.layer4_rn(layer_4) path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:]) path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:]) path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:]) path_1 = self.scratch.refinenet1(path_2, layer_1_rn) out = self.scratch.output_conv1(path_1) out = F.interpolate( out, (int(patch_h * 14), int(patch_w * 14)), mode="bilinear", align_corners=True, ) out = self.scratch.output_conv2(out) return out class RGBDDepth(nn.Module): def __init__( self, encoder="vitl", features=256, out_channels=[256, 512, 1024, 1024], use_bn=False, use_clstoken=False, max_depth=20.0, ): super(RGBDDepth, self).__init__() self.intermediate_layer_idx = { "vits": [2, 5, 8, 11], "vitb": [2, 5, 8, 11], "vitl": [4, 11, 17, 23], "vitg": [9, 19, 29, 39], } self.max_depth = max_depth self.encoder = encoder self.pretrained = DINOv2(model_name=encoder) self.depth_pretrained = DINOv2(model_name=encoder) # self.depth_head = DPTHead(self.pretrained.embed_dim, features, use_bn, out_channels=out_channels, use_clstoken=use_clstoken, sigact_out=False) self.depth_head_rgbd = DPTHead( self.pretrained.embed_dim * 2, features, use_bn, out_channels=out_channels, use_clstoken=use_clstoken, sigact_out=False, ) # cross att num_heads = 4 self.crossAtts = nn.ModuleList( [ nn.MultiheadAttention( self.pretrained.embed_dim, num_heads, batch_first=True ), nn.MultiheadAttention( self.pretrained.embed_dim, num_heads, batch_first=True ), nn.MultiheadAttention( self.pretrained.embed_dim, num_heads, batch_first=True ), nn.MultiheadAttention( self.pretrained.embed_dim, num_heads, batch_first=True ), ] ) def forward(self, x): rgb, depth = x[:, :3], x[:, 3:] patch_h, patch_w = x.shape[-2] // 14, x.shape[-1] // 14 with torch.no_grad(): features_rgb = self.pretrained.get_intermediate_layers( rgb, self.intermediate_layer_idx[self.encoder], return_class_token=True ) features_depth = self.depth_pretrained.get_intermediate_layers( depth.repeat(1, 3, 1, 1), self.intermediate_layer_idx[self.encoder], return_class_token=True, ) features = [] for f_rgb, f_depth, crossAtt in zip( features_rgb, features_depth, self.crossAtts ): B, N, C = f_rgb[0].shape tf_rgb = f_rgb[0].reshape(B * N, 1, C) tf_depth = f_depth[0].reshape(B * N, 1, C) token_feat = torch.concat((tf_rgb, tf_depth), axis=1) att_feat, _ = crossAtt(token_feat, token_feat, token_feat) att_feat = att_feat.reshape(B * N, 2, C).sum(axis=1).reshape(B, N, C) feat = torch.concat((f_rgb[0], att_feat), axis=2) cls_t = torch.concat((f_rgb[1], f_depth[1]), axis=1) tuples = (feat, cls_t) features.append(tuples) depth = self.depth_head_rgbd(features, patch_h, patch_w) depth = F.relu(depth) return depth.squeeze(1) @torch.no_grad() def infer_image(self, raw_image, depth_low_res, input_size=518): inputs, (h, w) = self.image2tensor(raw_image, depth_low_res, input_size) pred_depth = self.forward(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=518): transform = Compose( [ Resize( width=input_size, height=input_size, resize_target=True, keep_aspect_ratio=True, ensure_multiple_of=14, resize_method="lower_bound", image_interpolation_method=cv2.INTER_CUBIC, ), NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), PrepareForNet(), ] ) h, w = raw_image.shape[:2] image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) / 255.0 prepared = transform({"image": image, "depth": depth}) image = prepared["image"] image = torch.from_numpy(image).unsqueeze(0) depth = prepared["depth"] depth = torch.from_numpy(depth).unsqueeze(0).unsqueeze(0) inputs = torch.cat((image, depth), dim=1) DEVICE = ( "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" ) inputs = inputs.to(DEVICE) return inputs, (h, w)