| |
| |
| """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 |
|
|
|
|
| |
| |
| |
| |
| |
| |
|
|
| 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 |
|
|
|
|
| |
| |
| |
| |
| |
| |
|
|
| 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): |
| |
| scale_height = self.__height / height |
| scale_width = self.__width / width |
|
|
| if self.__keep_aspect_ratio: |
| if self.__resize_method == "lower_bound": |
| |
| if scale_width > scale_height: |
| |
| scale_height = scale_width |
| else: |
| |
| scale_width = scale_height |
| elif self.__resize_method == "upper_bound": |
| |
| if scale_width < scale_height: |
| |
| scale_height = scale_width |
| else: |
| |
| scale_width = scale_height |
| elif self.__resize_method == "minimal": |
| |
| if abs(1 - scale_width) < abs(1 - scale_height): |
| |
| scale_height = scale_width |
| else: |
| |
| 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] |
| ) |
|
|
| |
| 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 |
|
|
|
|
| |
| |
| |
| |
| |
| |
|
|
| 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_rgbd = DPTHead( |
| self.pretrained.embed_dim * 2, |
| features, |
| use_bn, |
| out_channels=out_channels, |
| use_clstoken=use_clstoken, |
| sigact_out=False, |
| ) |
|
|
| |
| 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) |
|
|
