diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/__init__.cpython-311.pyc b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e99725dface47d66e6596dea4a565ddc914e3127 Binary files /dev/null and b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/classifiers.cpython-311.pyc b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/classifiers.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c9f72e6ee3398bacf2f51dc383ee9159af96594 Binary files /dev/null and b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/classifiers.cpython-311.pyc differ diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/depthers.cpython-311.pyc b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/depthers.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c8eb73a3300e6662b98a01ad5d2e799f5aa36a77 Binary files /dev/null and b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/depthers.cpython-311.pyc differ diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/dinotxt.cpython-311.pyc b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/dinotxt.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0693a0f783e3718e633bdc21ea58f0ee11da2824 Binary files /dev/null and b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/dinotxt.cpython-311.pyc differ diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/utils.cpython-311.pyc b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a43bfc6b71e676ccd107aacf6dddf08bc5f16032 Binary files /dev/null and b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/__pycache__/utils.cpython-311.pyc differ diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/cell_dino/__pycache__/backbones.cpython-311.pyc b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/cell_dino/__pycache__/backbones.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad85902842a32433e359c00847368a13c5516890 Binary files /dev/null and b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/cell_dino/__pycache__/backbones.cpython-311.pyc differ diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/cell_dino/backbones.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/cell_dino/backbones.py new file mode 100644 index 0000000000000000000000000000000000000000..0e6b90b0204cadcd20de1363d332d71d2a68e544 --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/cell_dino/backbones.py @@ -0,0 +1,182 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the CC-by-NC licence, +# found in the LICENSE_CELL_DINO_CODE file in the root directory of this source tree. + +from enum import Enum +from typing import Optional, Union + +import torch + + +class Weights(Enum): + CELL_DINO = "CELL-DINO" + + +def _make_cell_dino_model( + *, + arch_name: str = "vit_large", + img_size: int = 518, + patch_size: int = 14, + init_values: float = 1.0, + ffn_layer: str = "mlp", + block_chunks: int = 0, + num_register_tokens: int = 0, + interpolate_antialias: bool = False, + interpolate_offset: float = 0.1, + pretrained: bool = True, + channel_adaptive: bool = False, + weights: Union[Weights, str] = Weights.CELL_DINO, + pretrained_url: Optional[str] = None, + pretrained_path: Optional[str] = None, + **kwargs, +): + from ...models import vision_transformer as vits + + if isinstance(weights, str): + try: + weights = Weights[weights] + except KeyError: + raise AssertionError(f"Unsupported weights: {weights}") + + vit_kwargs = dict( + img_size=img_size, + patch_size=patch_size, + init_values=init_values, + ffn_layer=ffn_layer, + block_chunks=block_chunks, + num_register_tokens=num_register_tokens, + interpolate_antialias=interpolate_antialias, + interpolate_offset=interpolate_offset, + channel_adaptive=channel_adaptive, + ) + vit_kwargs.update(**kwargs) + model = vits.__dict__[arch_name](**vit_kwargs) + + if pretrained: + if pretrained_path is not None: + state_dict = torch.load(pretrained_path, map_location="cpu") + else: + pretrained_url is not None + state_dict = torch.hub.load_state_dict_from_url(pretrained_url, map_location="cpu") + model.load_state_dict(state_dict, strict=True) + + return model + + +def cell_dino_hpa_vitl16( + *, + pretrained_url: Optional[str] = None, + pretrained_path: Optional[str] = None, + pretrained: bool = True, + weights: Union[Weights, str] = Weights.CELL_DINO, + in_channels: int = 4, + **kwargs, +): + """ + Cell-DINO ViT-L/16 model dataset pretrained on HPA dataset. + """ + return _make_cell_dino_model( + arch_name="vit_large", + patch_size=16, + img_size=224, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + block_chunks=4, + pretrained_url=pretrained_url, + pretrained_path=pretrained_path, + pretrained=pretrained, + weights=weights, + in_chans=in_channels, + **kwargs, + ) + + +def cell_dino_hpa_vitl14( + *, + pretrained_url: Optional[str] = None, + pretrained_path: Optional[str] = None, + pretrained: bool = True, + weights: Union[Weights, str] = Weights.CELL_DINO, + in_channels: int = 4, + **kwargs, +): + """ + Cell-DINO ViT-L/14 model dataset pretrained on LVD, then on HPA dataset. + """ + return _make_cell_dino_model( + arch_name="vit_large", + patch_size=14, + img_size=518, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + block_chunks=4, + pretrained_url=pretrained_url, + pretrained_path=pretrained_path, + pretrained=pretrained, + weights=weights, + in_chans=in_channels, + **kwargs, + ) + + +def cell_dino_cp_vits8( + *, + pretrained_url: Optional[str] = None, + pretrained_path: Optional[str] = None, + pretrained: bool = True, + weights: Union[Weights, str] = Weights.CELL_DINO, + in_channels: int = 5, + **kwargs, +): + """ + Cell-DINO ViT-S/8 model dataset pretrained on the combined cell painting dataset. + """ + return _make_cell_dino_model( + arch_name="vit_small", + patch_size=8, + img_size=128, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + block_chunks=4, + pretrained_url=pretrained_url, + pretrained_path=pretrained_path, + pretrained=pretrained, + weights=weights, + in_chans=in_channels, + **kwargs, + ) + + +def channel_adaptive_dino_vitl16( + *, + pretrained_url: Optional[str] = None, + pretrained_path: Optional[str] = None, + pretrained: bool = True, + weights: Union[Weights, str] = Weights.CELL_DINO, + in_channels: int = 1, + channel_adaptive: bool = True, + **kwargs, +): + """ + Cell-DINO ViT-L/16 model dataset pretrained on HPA dataset. + """ + return _make_cell_dino_model( + arch_name="vit_large", + patch_size=16, + img_size=224, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + block_chunks=4, + pretrained_url=pretrained_url, + pretrained_path=pretrained_path, + pretrained=pretrained, + weights=weights, + in_chans=in_channels, + channel_adaptive=channel_adaptive, + **kwargs, + ) diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__init__.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..91716e58ab6158d814df8c653644d9af4c7be65c --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +from .decode_heads import BNHead, DPTHead +from .encoder_decoder import DepthEncoderDecoder diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/__init__.cpython-311.pyc b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b93131503bac93ee642908a67fb46c50ce1910de Binary files /dev/null and b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/decode_heads.cpython-311.pyc b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/decode_heads.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..90feb3775beb4cad5581062a2556fbd232290183 Binary files /dev/null and b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/decode_heads.cpython-311.pyc differ diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/encoder_decoder.cpython-311.pyc b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/encoder_decoder.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0077d014512bce974f292ee624a0b640d0a7e873 Binary files /dev/null and b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/encoder_decoder.cpython-311.pyc differ diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/ops.cpython-311.pyc b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/ops.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c11ff79463d4c6526a47528ec24ce07db8eb52a9 Binary files /dev/null and b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/__pycache__/ops.cpython-311.pyc differ diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/decode_heads.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/decode_heads.py new file mode 100644 index 0000000000000000000000000000000000000000..f455accad38fec6ecdd53460233a564c34f434da --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/decode_heads.py @@ -0,0 +1,747 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import copy +from functools import partial +import math +import warnings + +import torch +import torch.nn as nn + +from .ops import resize + + +# XXX: (Untested) replacement for mmcv.imdenormalize() +def _imdenormalize(img, mean, std, to_bgr=True): + import numpy as np + + mean = mean.reshape(1, -1).astype(np.float64) + std = std.reshape(1, -1).astype(np.float64) + img = (img * std) + mean + if to_bgr: + img = img[::-1] + return img + + +class DepthBaseDecodeHead(nn.Module): + """Base class for BaseDecodeHead. + + Args: + in_channels (List): Input channels. + channels (int): Channels after modules, before conv_depth. + conv_layer (nn.Module): Conv layers. Default: None. + act_layer (nn.Module): Activation layers. Default: nn.ReLU. + loss_decode (dict): Config of decode loss. + Default: (). + sampler (dict|None): The config of depth map sampler. + Default: None. + align_corners (bool): align_corners argument of F.interpolate. + Default: False. + min_depth (int): Min depth in dataset setting. + Default: 1e-3. + max_depth (int): Max depth in dataset setting. + Default: None. + norm_layer (dict|None): Norm layers. + Default: None. + classify (bool): Whether predict depth in a cls.-reg. manner. + Default: False. + n_bins (int): The number of bins used in cls. step. + Default: 256. + bins_strategy (str): The discrete strategy used in cls. step. + Default: 'UD'. + norm_strategy (str): The norm strategy on cls. probability + distribution. Default: 'linear' + scale_up (str): Whether predict depth in a scale-up manner. + Default: False. + """ + + def __init__( + self, + in_channels, + conv_layer=None, + act_layer=nn.ReLU, + channels=96, + loss_decode=(), + sampler=None, + align_corners=False, + min_depth=1e-3, + max_depth=None, + norm_layer=None, + classify=False, + n_bins=256, + bins_strategy="UD", + norm_strategy="linear", + scale_up=False, + ): + super(DepthBaseDecodeHead, self).__init__() + + self.in_channels = in_channels + self.channels = channels + self.conf_layer = conv_layer + self.act_layer = act_layer + self.loss_decode = loss_decode + self.align_corners = align_corners + self.min_depth = min_depth + self.max_depth = max_depth + self.norm_layer = norm_layer + self.classify = classify + self.n_bins = n_bins + self.scale_up = scale_up + + if self.classify: + assert bins_strategy in ["UD", "SID"], "Support bins_strategy: UD, SID" + assert norm_strategy in ["linear", "softmax", "sigmoid"], "Support norm_strategy: linear, softmax, sigmoid" + + self.bins_strategy = bins_strategy + self.norm_strategy = norm_strategy + self.softmax = nn.Softmax(dim=1) + self.conv_depth = nn.Conv2d(channels, n_bins, kernel_size=3, padding=1, stride=1) + else: + self.conv_depth = nn.Conv2d(channels, 1, kernel_size=3, padding=1, stride=1) + + self.relu = nn.ReLU() + self.sigmoid = nn.Sigmoid() + + def forward(self, inputs, img_metas): + """Placeholder of forward function.""" + pass + + def forward_train(self, img, inputs, img_metas, depth_gt): + """Forward function for training. + Args: + inputs (list[Tensor]): List of multi-level img features. + img_metas (list[dict]): List of image info dict where each dict + has: 'img_shape', 'scale_factor', 'flip', and may also contain + 'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. + For details on the values of these keys see + `depth/datasets/pipelines/formatting.py:Collect`. + depth_gt (Tensor): GT depth + + Returns: + dict[str, Tensor]: a dictionary of loss components + """ + depth_pred = self.forward(inputs, img_metas) + losses = self.losses(depth_pred, depth_gt) + + log_imgs = self.log_images(img[0], depth_pred[0], depth_gt[0], img_metas[0]) + losses.update(**log_imgs) + + return losses + + def forward_test(self, inputs, img_metas): + """Forward function for testing. + Args: + inputs (list[Tensor]): List of multi-level img features. + img_metas (list[dict]): List of image info dict where each dict + has: 'img_shape', 'scale_factor', 'flip', and may also contain + 'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. + For details on the values of these keys see + `depth/datasets/pipelines/formatting.py:Collect`. + + Returns: + Tensor: Output depth map. + """ + return self.forward(inputs, img_metas) + + def depth_pred(self, feat): + """Prediction each pixel.""" + if self.classify: + logit = self.conv_depth(feat) + + if self.bins_strategy == "UD": + bins = torch.linspace(self.min_depth, self.max_depth, self.n_bins, device=feat.device) + elif self.bins_strategy == "SID": + bins = torch.logspace(self.min_depth, self.max_depth, self.n_bins, device=feat.device) + + # following Adabins, default linear + if self.norm_strategy == "linear": + logit = torch.relu(logit) + eps = 0.1 + logit = logit + eps + logit = logit / logit.sum(dim=1, keepdim=True) + elif self.norm_strategy == "softmax": + logit = torch.softmax(logit, dim=1) + elif self.norm_strategy == "sigmoid": + logit = torch.sigmoid(logit) + logit = logit / logit.sum(dim=1, keepdim=True) + + output = torch.einsum("ikmn,k->imn", [logit, bins]).unsqueeze(dim=1) + + else: + if self.scale_up: + output = self.sigmoid(self.conv_depth(feat)) * self.max_depth + else: + output = self.relu(self.conv_depth(feat)) + self.min_depth + return output + + def losses(self, depth_pred, depth_gt): + """Compute depth loss.""" + loss = dict() + depth_pred = resize( + input=depth_pred, size=depth_gt.shape[2:], mode="bilinear", align_corners=self.align_corners, warning=False + ) + if not isinstance(self.loss_decode, nn.ModuleList): + losses_decode = [self.loss_decode] + else: + losses_decode = self.loss_decode + for loss_decode in losses_decode: + if loss_decode.loss_name not in loss: + loss[loss_decode.loss_name] = loss_decode(depth_pred, depth_gt) + else: + loss[loss_decode.loss_name] += loss_decode(depth_pred, depth_gt) + return loss + + def log_images(self, img_path, depth_pred, depth_gt, img_meta): + import numpy as np + + show_img = copy.deepcopy(img_path.detach().cpu().permute(1, 2, 0)) + show_img = show_img.numpy().astype(np.float32) + show_img = _imdenormalize( + show_img, + img_meta["img_norm_cfg"]["mean"], + img_meta["img_norm_cfg"]["std"], + img_meta["img_norm_cfg"]["to_rgb"], + ) + show_img = np.clip(show_img, 0, 255) + show_img = show_img.astype(np.uint8) + show_img = show_img[:, :, ::-1] + show_img = show_img.transpose(0, 2, 1) + show_img = show_img.transpose(1, 0, 2) + + depth_pred = depth_pred / torch.max(depth_pred) + depth_gt = depth_gt / torch.max(depth_gt) + + depth_pred_color = copy.deepcopy(depth_pred.detach().cpu()) + depth_gt_color = copy.deepcopy(depth_gt.detach().cpu()) + + return {"img_rgb": show_img, "img_depth_pred": depth_pred_color, "img_depth_gt": depth_gt_color} + + +class BNHead(DepthBaseDecodeHead): + """Just a batchnorm.""" + + def __init__(self, input_transform="resize_concat", in_index=(0, 1, 2, 3), upsample=1, **kwargs): + super().__init__(**kwargs) + self.input_transform = input_transform + self.in_index = in_index + self.upsample = upsample + # self.bn = nn.SyncBatchNorm(self.in_channels) + if self.classify: + self.conv_depth = nn.Conv2d(self.channels, self.n_bins, kernel_size=1, padding=0, stride=1) + else: + self.conv_depth = nn.Conv2d(self.channels, 1, kernel_size=1, padding=0, stride=1) + + def _transform_inputs(self, inputs): + """Transform inputs for decoder. + Args: + inputs (list[Tensor]): List of multi-level img features. + Returns: + Tensor: The transformed inputs + """ + + if "concat" in self.input_transform: + inputs = [inputs[i] for i in self.in_index] + if "resize" in self.input_transform: + inputs = [ + resize( + input=x, + size=[s * self.upsample for s in inputs[0].shape[2:]], + mode="bilinear", + align_corners=self.align_corners, + ) + for x in inputs + ] + inputs = torch.cat(inputs, dim=1) + elif self.input_transform == "multiple_select": + inputs = [inputs[i] for i in self.in_index] + else: + inputs = inputs[self.in_index] + + return inputs + + def _forward_feature(self, inputs, img_metas=None, **kwargs): + """Forward function for feature maps before classifying each pixel with + ``self.cls_seg`` fc. + Args: + inputs (list[Tensor]): List of multi-level img features. + Returns: + feats (Tensor): A tensor of shape (batch_size, self.channels, + H, W) which is feature map for last layer of decoder head. + """ + # accept lists (for cls token) + inputs = list(inputs) + for i, x in enumerate(inputs): + if len(x) == 2: + x, cls_token = x[0], x[1] + if len(x.shape) == 2: + x = x[:, :, None, None] + cls_token = cls_token[:, :, None, None].expand_as(x) + inputs[i] = torch.cat((x, cls_token), 1) + else: + x = x[0] + if len(x.shape) == 2: + x = x[:, :, None, None] + inputs[i] = x + x = self._transform_inputs(inputs) + # feats = self.bn(x) + return x + + def forward(self, inputs, img_metas=None, **kwargs): + """Forward function.""" + output = self._forward_feature(inputs, img_metas=img_metas, **kwargs) + output = self.depth_pred(output) + return output + + +class ConvModule(nn.Module): + """A conv block that bundles conv/norm/activation layers. + + This block simplifies the usage of convolution layers, which are commonly + used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU). + It is based upon three build methods: `build_conv_layer()`, + `build_norm_layer()` and `build_activation_layer()`. + + Besides, we add some additional features in this module. + 1. Automatically set `bias` of the conv layer. + 2. Spectral norm is supported. + 3. More padding modes are supported. Before PyTorch 1.5, nn.Conv2d only + supports zero and circular padding, and we add "reflect" padding mode. + + Args: + in_channels (int): Number of channels in the input feature map. + Same as that in ``nn._ConvNd``. + out_channels (int): Number of channels produced by the convolution. + Same as that in ``nn._ConvNd``. + kernel_size (int | tuple[int]): Size of the convolving kernel. + Same as that in ``nn._ConvNd``. + stride (int | tuple[int]): Stride of the convolution. + Same as that in ``nn._ConvNd``. + padding (int | tuple[int]): Zero-padding added to both sides of + the input. Same as that in ``nn._ConvNd``. + dilation (int | tuple[int]): Spacing between kernel elements. + Same as that in ``nn._ConvNd``. + groups (int): Number of blocked connections from input channels to + output channels. Same as that in ``nn._ConvNd``. + bias (bool | str): If specified as `auto`, it will be decided by the + norm_layer. Bias will be set as True if `norm_layer` is None, otherwise + False. Default: "auto". + conv_layer (nn.Module): Convolution layer. Default: None, + which means using conv2d. + norm_layer (nn.Module): Normalization layer. Default: None. + act_layer (nn.Module): Activation layer. Default: nn.ReLU. + inplace (bool): Whether to use inplace mode for activation. + Default: True. + with_spectral_norm (bool): Whether use spectral norm in conv module. + Default: False. + padding_mode (str): If the `padding_mode` has not been supported by + current `Conv2d` in PyTorch, we will use our own padding layer + instead. Currently, we support ['zeros', 'circular'] with official + implementation and ['reflect'] with our own implementation. + Default: 'zeros'. + order (tuple[str]): The order of conv/norm/activation layers. It is a + sequence of "conv", "norm" and "act". Common examples are + ("conv", "norm", "act") and ("act", "conv", "norm"). + Default: ('conv', 'norm', 'act'). + """ + + _abbr_ = "conv_block" + + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + groups=1, + bias="auto", + conv_layer=nn.Conv2d, + norm_layer=None, + act_layer=nn.ReLU, + inplace=True, + with_spectral_norm=False, + padding_mode="zeros", + order=("conv", "norm", "act"), + ): + super(ConvModule, self).__init__() + official_padding_mode = ["zeros", "circular"] + self.conv_layer = conv_layer + self.norm_layer = norm_layer + self.act_layer = act_layer + self.inplace = inplace + self.with_spectral_norm = with_spectral_norm + self.with_explicit_padding = padding_mode not in official_padding_mode + self.order = order + assert isinstance(self.order, tuple) and len(self.order) == 3 + assert set(order) == set(["conv", "norm", "act"]) + + self.with_norm = norm_layer is not None + self.with_activation = act_layer is not None + # if the conv layer is before a norm layer, bias is unnecessary. + if bias == "auto": + bias = not self.with_norm + self.with_bias = bias + + if self.with_explicit_padding: + if padding_mode == "zeros": + padding_layer = nn.ZeroPad2d + else: + raise AssertionError(f"Unsupported padding mode: {padding_mode}") + self.pad = padding_layer(padding) + + # reset padding to 0 for conv module + conv_padding = 0 if self.with_explicit_padding else padding + # build convolution layer + self.conv = self.conv_layer( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=conv_padding, + dilation=dilation, + groups=groups, + bias=bias, + ) + # export the attributes of self.conv to a higher level for convenience + self.in_channels = self.conv.in_channels + self.out_channels = self.conv.out_channels + self.kernel_size = self.conv.kernel_size + self.stride = self.conv.stride + self.padding = padding + self.dilation = self.conv.dilation + self.transposed = self.conv.transposed + self.output_padding = self.conv.output_padding + self.groups = self.conv.groups + + if self.with_spectral_norm: + self.conv = nn.utils.spectral_norm(self.conv) + + # build normalization layers + if self.with_norm: + # norm layer is after conv layer + if order.index("norm") > order.index("conv"): + norm_channels = out_channels + else: + norm_channels = in_channels + norm = partial(norm_layer, num_features=norm_channels) + self.add_module("norm", norm) + if self.with_bias: + from torch.nnModules.batchnorm import _BatchNorm + from torch.nnModules.instancenorm import _InstanceNorm + + if isinstance(norm, (_BatchNorm, _InstanceNorm)): + warnings.warn("Unnecessary conv bias before batch/instance norm") + else: + self.norm_name = None + + # build activation layer + if self.with_activation: + # nn.Tanh has no 'inplace' argument + # (nn.Tanh, nn.PReLU, nn.Sigmoid, nn.HSigmoid, nn.Swish, nn.GELU) + if not isinstance(act_layer, (nn.Tanh, nn.PReLU, nn.Sigmoid, nn.GELU)): + act_layer = partial(act_layer, inplace=inplace) + self.activate = act_layer() + + # Use msra init by default + self.init_weights() + + @property + def norm(self): + if self.norm_name: + return getattr(self, self.norm_name) + else: + return None + + def init_weights(self): + # 1. It is mainly for customized conv layers with their own + # initialization manners by calling their own ``init_weights()``, + # and we do not want ConvModule to override the initialization. + # 2. For customized conv layers without their own initialization + # manners (that is, they don't have their own ``init_weights()``) + # and PyTorch's conv layers, they will be initialized by + # this method with default ``kaiming_init``. + # Note: For PyTorch's conv layers, they will be overwritten by our + # initialization implementation using default ``kaiming_init``. + if not hasattr(self.conv, "init_weights"): + if self.with_activation and isinstance(self.act_layer, nn.LeakyReLU): + nonlinearity = "leaky_relu" + a = 0.01 # XXX: default negative_slope + else: + nonlinearity = "relu" + a = 0 + if hasattr(self.conv, "weight") and self.conv.weight is not None: + nn.init.kaiming_normal_(self.conv.weight, a=a, mode="fan_out", nonlinearity=nonlinearity) + if hasattr(self.conv, "bias") and self.conv.bias is not None: + nn.init.constant_(self.conv.bias, 0) + if self.with_norm: + if hasattr(self.norm, "weight") and self.norm.weight is not None: + nn.init.constant_(self.norm.weight, 1) + if hasattr(self.norm, "bias") and self.norm.bias is not None: + nn.init.constant_(self.norm.bias, 0) + + def forward(self, x, activate=True, norm=True): + for layer in self.order: + if layer == "conv": + if self.with_explicit_padding: + x = self.pad(x) + x = self.conv(x) + elif layer == "norm" and norm and self.with_norm: + x = self.norm(x) + elif layer == "act" and activate and self.with_activation: + x = self.activate(x) + return x + + +class Interpolate(nn.Module): + def __init__(self, scale_factor, mode, align_corners=False): + super(Interpolate, self).__init__() + self.interp = nn.functional.interpolate + self.scale_factor = scale_factor + self.mode = mode + self.align_corners = align_corners + + def forward(self, x): + x = self.interp(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners) + return x + + +class HeadDepth(nn.Module): + def __init__(self, features): + super(HeadDepth, self).__init__() + self.head = nn.Sequential( + nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1), + Interpolate(scale_factor=2, mode="bilinear", align_corners=True), + nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1), + nn.ReLU(), + nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0), + ) + + def forward(self, x): + x = self.head(x) + return x + + +class ReassembleBlocks(nn.Module): + """ViTPostProcessBlock, process cls_token in ViT backbone output and + rearrange the feature vector to feature map. + Args: + in_channels (int): ViT feature channels. Default: 768. + out_channels (List): output channels of each stage. + Default: [96, 192, 384, 768]. + readout_type (str): Type of readout operation. Default: 'ignore'. + patch_size (int): The patch size. Default: 16. + """ + + def __init__(self, in_channels=768, out_channels=[96, 192, 384, 768], readout_type="ignore", patch_size=16): + super(ReassembleBlocks, self).__init__() + + assert readout_type in ["ignore", "add", "project"] + self.readout_type = readout_type + self.patch_size = patch_size + + self.projects = nn.ModuleList( + [ + ConvModule( + in_channels=in_channels, + out_channels=out_channel, + kernel_size=1, + act_layer=None, + ) + 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 self.readout_type == "project": + 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())) + + def forward(self, inputs): + assert isinstance(inputs, list) + out = [] + for i, x in enumerate(inputs): + assert len(x) == 2 + x, cls_token = x[0], x[1] + feature_shape = x.shape + if self.readout_type == "project": + x = x.flatten(2).permute((0, 2, 1)) + readout = cls_token.unsqueeze(1).expand_as(x) + x = self.readout_projects[i](torch.cat((x, readout), -1)) + x = x.permute(0, 2, 1).reshape(feature_shape) + elif self.readout_type == "add": + x = x.flatten(2) + cls_token.unsqueeze(-1) + x = x.reshape(feature_shape) + else: + pass + x = self.projects[i](x) + x = self.resize_layers[i](x) + out.append(x) + return out + + +class PreActResidualConvUnit(nn.Module): + """ResidualConvUnit, pre-activate residual unit. + Args: + in_channels (int): number of channels in the input feature map. + act_layer (nn.Module): activation layer. + norm_layer (nn.Module): norm layer. + stride (int): stride of the first block. Default: 1 + dilation (int): dilation rate for convs layers. Default: 1. + """ + + def __init__(self, in_channels, act_layer, norm_layer, stride=1, dilation=1): + super(PreActResidualConvUnit, self).__init__() + + self.conv1 = ConvModule( + in_channels, + in_channels, + 3, + stride=stride, + padding=dilation, + dilation=dilation, + norm_layer=norm_layer, + act_layer=act_layer, + bias=False, + order=("act", "conv", "norm"), + ) + + self.conv2 = ConvModule( + in_channels, + in_channels, + 3, + padding=1, + norm_layer=norm_layer, + act_layer=act_layer, + bias=False, + order=("act", "conv", "norm"), + ) + + def forward(self, inputs): + inputs_ = inputs.clone() + x = self.conv1(inputs) + x = self.conv2(x) + return x + inputs_ + + +class FeatureFusionBlock(nn.Module): + """FeatureFusionBlock, merge feature map from different stages. + Args: + in_channels (int): Input channels. + act_layer (nn.Module): activation layer for ResidualConvUnit. + norm_layer (nn.Module): normalization layer. + expand (bool): Whether expand the channels in post process block. + Default: False. + align_corners (bool): align_corner setting for bilinear upsample. + Default: True. + """ + + def __init__(self, in_channels, act_layer, norm_layer, expand=False, align_corners=True): + super(FeatureFusionBlock, self).__init__() + + self.in_channels = in_channels + self.expand = expand + self.align_corners = align_corners + + self.out_channels = in_channels + if self.expand: + self.out_channels = in_channels // 2 + + self.project = ConvModule(self.in_channels, self.out_channels, kernel_size=1, act_layer=None, bias=True) + + self.res_conv_unit1 = PreActResidualConvUnit( + in_channels=self.in_channels, act_layer=act_layer, norm_layer=norm_layer + ) + self.res_conv_unit2 = PreActResidualConvUnit( + in_channels=self.in_channels, act_layer=act_layer, norm_layer=norm_layer + ) + + def forward(self, *inputs): + x = inputs[0] + if len(inputs) == 2: + if x.shape != inputs[1].shape: + res = resize(inputs[1], size=(x.shape[2], x.shape[3]), mode="bilinear", align_corners=False) + else: + res = inputs[1] + x = x + self.res_conv_unit1(res) + x = self.res_conv_unit2(x) + x = resize(x, scale_factor=2, mode="bilinear", align_corners=self.align_corners) + x = self.project(x) + return x + + +class DPTHead(DepthBaseDecodeHead): + """Vision Transformers for Dense Prediction. + This head is implemented of `DPT `_. + Args: + embed_dims (int): The embed dimension of the ViT backbone. + Default: 768. + post_process_channels (List): Out channels of post process conv + layers. Default: [96, 192, 384, 768]. + readout_type (str): Type of readout operation. Default: 'ignore'. + patch_size (int): The patch size. Default: 16. + expand_channels (bool): Whether expand the channels in post process + block. Default: False. + """ + + def __init__( + self, + embed_dims=768, + post_process_channels=[96, 192, 384, 768], + readout_type="ignore", + patch_size=16, + expand_channels=False, + **kwargs, + ): + super(DPTHead, self).__init__(**kwargs) + + self.in_channels = self.in_channels + self.expand_channels = expand_channels + self.reassemble_blocks = ReassembleBlocks(embed_dims, post_process_channels, readout_type, patch_size) + + self.post_process_channels = [ + channel * math.pow(2, i) if expand_channels else channel for i, channel in enumerate(post_process_channels) + ] + self.convs = nn.ModuleList() + for channel in self.post_process_channels: + self.convs.append(ConvModule(channel, self.channels, kernel_size=3, padding=1, act_layer=None, bias=False)) + self.fusion_blocks = nn.ModuleList() + for _ in range(len(self.convs)): + self.fusion_blocks.append(FeatureFusionBlock(self.channels, self.act_layer, self.norm_layer)) + self.fusion_blocks[0].res_conv_unit1 = None + self.project = ConvModule(self.channels, self.channels, kernel_size=3, padding=1, norm_layer=self.norm_layer) + self.num_fusion_blocks = len(self.fusion_blocks) + self.num_reassemble_blocks = len(self.reassemble_blocks.resize_layers) + self.num_post_process_channels = len(self.post_process_channels) + assert self.num_fusion_blocks == self.num_reassemble_blocks + assert self.num_reassemble_blocks == self.num_post_process_channels + self.conv_depth = HeadDepth(self.channels) + + def forward(self, inputs, img_metas): + assert len(inputs) == self.num_reassemble_blocks + x = [inp for inp in inputs] + x = self.reassemble_blocks(x) + x = [self.convs[i](feature) for i, feature in enumerate(x)] + out = self.fusion_blocks[0](x[-1]) + for i in range(1, len(self.fusion_blocks)): + out = self.fusion_blocks[i](out, x[-(i + 1)]) + out = self.project(out) + out = self.depth_pred(out) + return out diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/encoder_decoder.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/encoder_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..eb29ced67957a336e763b0e7c90c0eeaea36fea8 --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/encoder_decoder.py @@ -0,0 +1,351 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .ops import resize + + +def add_prefix(inputs, prefix): + """Add prefix for dict. + + Args: + inputs (dict): The input dict with str keys. + prefix (str): The prefix to add. + + Returns: + + dict: The dict with keys updated with ``prefix``. + """ + + outputs = dict() + for name, value in inputs.items(): + outputs[f"{prefix}.{name}"] = value + + return outputs + + +class DepthEncoderDecoder(nn.Module): + """Encoder Decoder depther. + + EncoderDecoder typically consists of backbone and decode_head. + """ + + def __init__(self, backbone, decode_head): + super(DepthEncoderDecoder, self).__init__() + + self.backbone = backbone + self.decode_head = decode_head + self.align_corners = self.decode_head.align_corners + + def extract_feat(self, img): + """Extract features from images.""" + return self.backbone(img) + + def encode_decode(self, img, img_metas, rescale=True, size=None): + """Encode images with backbone and decode into a depth estimation + map of the same size as input.""" + x = self.extract_feat(img) + out = self._decode_head_forward_test(x, img_metas) + # crop the pred depth to the certain range. + out = torch.clamp(out, min=self.decode_head.min_depth, max=self.decode_head.max_depth) + if rescale: + if size is None: + if img_metas is not None: + size = img_metas[0]["ori_shape"][:2] + else: + size = img.shape[2:] + out = resize(input=out, size=size, mode="bilinear", align_corners=self.align_corners) + return out + + def _decode_head_forward_train(self, img, x, img_metas, depth_gt, **kwargs): + """Run forward function and calculate loss for decode head in + training.""" + losses = dict() + loss_decode = self.decode_head.forward_train(img, x, img_metas, depth_gt, **kwargs) + losses.update(add_prefix(loss_decode, "decode")) + return losses + + def _decode_head_forward_test(self, x, img_metas): + """Run forward function and calculate loss for decode head in + inference.""" + depth_pred = self.decode_head.forward_test(x, img_metas) + return depth_pred + + def forward_dummy(self, img): + """Dummy forward function.""" + depth = self.encode_decode(img, None) + + return depth + + def forward_train(self, img, img_metas, depth_gt, **kwargs): + """Forward function for training. + + Args: + img (Tensor): Input images. + img_metas (list[dict]): List of image info dict where each dict + has: 'img_shape', 'scale_factor', 'flip', and may also contain + 'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. + For details on the values of these keys see + `depth/datasets/pipelines/formatting.py:Collect`. + depth_gt (Tensor): Depth gt + used if the architecture supports depth estimation task. + + Returns: + dict[str, Tensor]: a dictionary of loss components + """ + + x = self.extract_feat(img) + + losses = dict() + + # the last of x saves the info from neck + loss_decode = self._decode_head_forward_train(img, x, img_metas, depth_gt, **kwargs) + + losses.update(loss_decode) + + return losses + + def whole_inference(self, img, img_meta, rescale, size=None): + """Inference with full image.""" + return self.encode_decode(img, img_meta, rescale, size=size) + + def slide_inference(self, img, img_meta, rescale, stride, crop_size): + """Inference by sliding-window with overlap. + + If h_crop > h_img or w_crop > w_img, the small patch will be used to + decode without padding. + """ + + h_stride, w_stride = stride + h_crop, w_crop = crop_size + batch_size, _, h_img, w_img = img.size() + h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1 + w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1 + preds = img.new_zeros((batch_size, 1, h_img, w_img)) + count_mat = img.new_zeros((batch_size, 1, h_img, w_img)) + for h_idx in range(h_grids): + for w_idx in range(w_grids): + y1 = h_idx * h_stride + x1 = w_idx * w_stride + y2 = min(y1 + h_crop, h_img) + x2 = min(x1 + w_crop, w_img) + y1 = max(y2 - h_crop, 0) + x1 = max(x2 - w_crop, 0) + crop_img = img[:, :, y1:y2, x1:x2] + depth_pred = self.encode_decode(crop_img, img_meta, rescale) + preds += F.pad(depth_pred, (int(x1), int(preds.shape[3] - x2), int(y1), int(preds.shape[2] - y2))) + + count_mat[:, :, y1:y2, x1:x2] += 1 + assert (count_mat == 0).sum() == 0 + if torch.onnx.is_in_onnx_export(): + # cast count_mat to constant while exporting to ONNX + count_mat = torch.from_numpy(count_mat.cpu().detach().numpy()).to(device=img.device) + preds = preds / count_mat + return preds + + def inference(self, img, img_meta, rescale, size=None, mode="whole"): + """Inference with slide/whole style. + + Args: + img (Tensor): The input image of shape (N, 3, H, W). + img_meta (dict): Image info dict where each dict has: 'img_shape', + 'scale_factor', 'flip', and may also contain + 'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. + For details on the values of these keys see + `depth/datasets/pipelines/formatting.py:Collect`. + rescale (bool): Whether rescale back to original shape. + + Returns: + Tensor: The output depth map. + """ + + assert mode in ["slide", "whole"] + ori_shape = img_meta[0]["ori_shape"] + assert all(_["ori_shape"] == ori_shape for _ in img_meta) + if mode == "slide": + depth_pred = self.slide_inference(img, img_meta, rescale) + else: + depth_pred = self.whole_inference(img, img_meta, rescale, size=size) + output = depth_pred + flip = img_meta[0]["flip"] + if flip: + flip_direction = img_meta[0]["flip_direction"] + assert flip_direction in ["horizontal", "vertical"] + if flip_direction == "horizontal": + output = output.flip(dims=(3,)) + elif flip_direction == "vertical": + output = output.flip(dims=(2,)) + + return output + + def simple_test(self, img, img_meta, rescale=True): + """Simple test with single image.""" + depth_pred = self.inference(img, img_meta, rescale) + if torch.onnx.is_in_onnx_export(): + # our inference backend only support 4D output + depth_pred = depth_pred.unsqueeze(0) + return depth_pred + depth_pred = depth_pred.cpu().numpy() + # unravel batch dim + depth_pred = list(depth_pred) + return depth_pred + + def aug_test(self, imgs, img_metas, rescale=True): + """Test with augmentations. + + Only rescale=True is supported. + """ + # aug_test rescale all imgs back to ori_shape for now + assert rescale + # to save memory, we get augmented depth logit inplace + depth_pred = self.inference(imgs[0], img_metas[0], rescale) + for i in range(1, len(imgs)): + cur_depth_pred = self.inference(imgs[i], img_metas[i], rescale, size=depth_pred.shape[-2:]) + depth_pred += cur_depth_pred + depth_pred /= len(imgs) + depth_pred = depth_pred.cpu().numpy() + # unravel batch dim + depth_pred = list(depth_pred) + return depth_pred + + def forward_test(self, imgs, img_metas, **kwargs): + """ + Args: + imgs (List[Tensor]): the outer list indicates test-time + augmentations and inner Tensor should have a shape NxCxHxW, + which contains all images in the batch. + img_metas (List[List[dict]]): the outer list indicates test-time + augs (multiscale, flip, etc.) and the inner list indicates + images in a batch. + """ + for var, name in [(imgs, "imgs"), (img_metas, "img_metas")]: + if not isinstance(var, list): + raise TypeError(f"{name} must be a list, but got " f"{type(var)}") + num_augs = len(imgs) + if num_augs != len(img_metas): + raise ValueError(f"num of augmentations ({len(imgs)}) != " f"num of image meta ({len(img_metas)})") + # all images in the same aug batch all of the same ori_shape and pad + # shape + for img_meta in img_metas: + ori_shapes = [_["ori_shape"] for _ in img_meta] + assert all(shape == ori_shapes[0] for shape in ori_shapes) + img_shapes = [_["img_shape"] for _ in img_meta] + assert all(shape == img_shapes[0] for shape in img_shapes) + pad_shapes = [_["pad_shape"] for _ in img_meta] + assert all(shape == pad_shapes[0] for shape in pad_shapes) + + if num_augs == 1: + return self.simple_test(imgs[0], img_metas[0], **kwargs) + else: + return self.aug_test(imgs, img_metas, **kwargs) + + def forward(self, img, img_metas, return_loss=True, **kwargs): + """Calls either :func:`forward_train` or :func:`forward_test` depending + on whether ``return_loss`` is ``True``. + + Note this setting will change the expected inputs. When + ``return_loss=True``, img and img_meta are single-nested (i.e. Tensor + and List[dict]), and when ``resturn_loss=False``, img and img_meta + should be double nested (i.e. List[Tensor], List[List[dict]]), with + the outer list indicating test time augmentations. + """ + if return_loss: + return self.forward_train(img, img_metas, **kwargs) + else: + return self.forward_test(img, img_metas, **kwargs) + + def train_step(self, data_batch, optimizer, **kwargs): + """The iteration step during training. + + This method defines an iteration step during training, except for the + back propagation and optimizer updating, which are done in an optimizer + hook. Note that in some complicated cases or models, the whole process + including back propagation and optimizer updating is also defined in + this method, such as GAN. + + Args: + data (dict): The output of dataloader. + optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of + runner is passed to ``train_step()``. This argument is unused + and reserved. + + Returns: + dict: It should contain at least 3 keys: ``loss``, ``log_vars``, + ``num_samples``. + ``loss`` is a tensor for back propagation, which can be a + weighted sum of multiple losses. + ``log_vars`` contains all the variables to be sent to the + logger. + ``num_samples`` indicates the batch size (when the model is + DDP, it means the batch size on each GPU), which is used for + averaging the logs. + """ + losses = self(**data_batch) + + # split losses and images + real_losses = {} + log_imgs = {} + for k, v in losses.items(): + if "img" in k: + log_imgs[k] = v + else: + real_losses[k] = v + + loss, log_vars = self._parse_losses(real_losses) + + outputs = dict(loss=loss, log_vars=log_vars, num_samples=len(data_batch["img_metas"]), log_imgs=log_imgs) + + return outputs + + def val_step(self, data_batch, **kwargs): + """The iteration step during validation. + + This method shares the same signature as :func:`train_step`, but used + during val epochs. Note that the evaluation after training epochs is + not implemented with this method, but an evaluation hook. + """ + output = self(**data_batch, **kwargs) + return output + + @staticmethod + def _parse_losses(losses): + import torch.distributed as dist + + """Parse the raw outputs (losses) of the network. + + Args: + losses (dict): Raw output of the network, which usually contain + losses and other necessary information. + + Returns: + tuple[Tensor, dict]: (loss, log_vars), loss is the loss tensor + which may be a weighted sum of all losses, log_vars contains + all the variables to be sent to the logger. + """ + log_vars = OrderedDict() + for loss_name, loss_value in losses.items(): + if isinstance(loss_value, torch.Tensor): + log_vars[loss_name] = loss_value.mean() + elif isinstance(loss_value, list): + log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value) + else: + raise TypeError(f"{loss_name} is not a tensor or list of tensors") + + loss = sum(_value for _key, _value in log_vars.items() if "loss" in _key) + + log_vars["loss"] = loss + for loss_name, loss_value in log_vars.items(): + # reduce loss when distributed training + if dist.is_available() and dist.is_initialized(): + loss_value = loss_value.data.clone() + dist.all_reduce(loss_value.div_(dist.get_world_size())) + log_vars[loss_name] = loss_value.item() + + return loss, log_vars diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/ops.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..15880ee0cb7652d4b41c489b927bf6a156b40e5e --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/depth/ops.py @@ -0,0 +1,28 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import warnings + +import torch.nn.functional as F + + +def resize(input, size=None, scale_factor=None, mode="nearest", align_corners=None, warning=False): + if warning: + if size is not None and align_corners: + input_h, input_w = tuple(int(x) for x in input.shape[2:]) + output_h, output_w = tuple(int(x) for x in size) + if output_h > input_h or output_w > output_h: + if ( + (output_h > 1 and output_w > 1 and input_h > 1 and input_w > 1) + and (output_h - 1) % (input_h - 1) + and (output_w - 1) % (input_w - 1) + ): + warnings.warn( + f"When align_corners={align_corners}, " + "the output would more aligned if " + f"input size {(input_h, input_w)} is `x+1` and " + f"out size {(output_h, output_w)} is `nx+1`" + ) + return F.interpolate(input, size, scale_factor, mode, align_corners) diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/dinotxt_model.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/dinotxt_model.py new file mode 100644 index 0000000000000000000000000000000000000000..e561e1915485d30aee3cd04fb65da0e3db55dd36 --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/dinotxt_model.py @@ -0,0 +1,135 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import math +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn, Tensor + +from .vision_tower import VisionTower +from .text_tower import TextTower + + +@dataclass +class DinoTxtConfig: + embed_dim: int + vision_model_freeze_backbone: bool = True + vision_model_train_img_size: int = 224 + vision_model_use_class_token: bool = True + vision_model_use_patch_tokens: bool = False + vision_model_num_head_blocks: int = 0 + vision_model_head_blocks_drop_path: float = 0.3 + vision_model_use_linear_projection: bool = False + vision_model_patch_tokens_pooler_type: str = "mean" + vision_model_patch_token_layer: int = 1 # which layer to take patch tokens from + # 1 - last layer, 2 - second last layer, etc. + text_model_freeze_backbone: bool = False + text_model_num_head_blocks: int = 0 + text_model_head_blocks_is_causal: bool = False + text_model_head_blocks_drop_prob: float = 0.0 + text_model_tokens_pooler_type: str = "first" + text_model_use_linear_projection: bool = False + init_logit_scale: float = math.log(1 / 0.07) + init_logit_bias: Optional[float] = None + freeze_logit_scale: bool = False + + +class DinoTxt(nn.Module): + def __init__( + self, + model_config: DinoTxtConfig, + vision_backbone: nn.Module, + text_backbone: nn.Module, + ): + super().__init__() + self.model_config = model_config + self.visual_model = VisionTower( + vision_backbone, + model_config.vision_model_freeze_backbone, + model_config.embed_dim, + model_config.vision_model_num_head_blocks, + model_config.vision_model_head_blocks_drop_path, + model_config.vision_model_use_class_token, + model_config.vision_model_use_patch_tokens, + model_config.vision_model_patch_token_layer, + model_config.vision_model_patch_tokens_pooler_type, + model_config.vision_model_use_linear_projection, + ) + self.text_model = TextTower( + text_backbone, + model_config.text_model_freeze_backbone, + model_config.embed_dim, + model_config.text_model_num_head_blocks, + model_config.text_model_head_blocks_is_causal, + model_config.text_model_head_blocks_drop_prob, + model_config.text_model_tokens_pooler_type, + model_config.text_model_use_linear_projection, + ) + self.logit_scale = nn.Parameter(torch.ones(1) * model_config.init_logit_scale) + if model_config.freeze_logit_scale: + self.logit_scale.requires_grad = False + + def init_weights(self): + self.visual_model.init_weights() + self.text_model.init_weights() + + def get_visual_class_and_patch_tokens(self, image: Tensor) -> Tuple[Tensor, Tensor]: + return self.visual_model.get_class_and_patch_tokens(image) + + def encode_image( + self, + image: Tensor, + normalize: bool = False, + ) -> Tensor: + """ + Encode an image into a vector descriptor containing both global and local features. + + Args: + image (Tensor): Tensor of shape `(batch_size, rgb, height, width)`, normalized using ImageNet mean and std. + normalize (bool, optional): Whether to normalize the output vectors. Default is False. + Image features should always be normalized before comparing them with text features: + Returns: + Tensor: Tensor of shape `(batch_size, embed_dim)` containing the image features. + The first half of the vector corresponds to the global features (class token), + and the second half corresponds to the pooled patch features. + """ + features = self.visual_model(image) + return F.normalize(features, dim=-1) if normalize else features + + def encode_text(self, text: Tensor, normalize: bool = False) -> Tensor: + """ + Encode a text input into a vector descriptor. + + Args: + text (Tensor): Tensor of shape `(batch_size, seq_len)` containing token indices. + normalize (bool, optional): Whether to normalize the output vectors. Default is False. + Text features should be normalized before comparing them with image features: + Returns: + Tensor: Tensor of shape `(batch_size, embed_dim)` containing the text features. + As a consequence of the training procedure, assume that the first half of the tensor corresponds + to global image features and the second half to pooled patch features. + """ + features = self.text_model(text) + return F.normalize(features, dim=-1) if normalize else features + + def get_logits(self, image: Tensor, text: Tensor) -> Tuple[Tensor, Tensor]: + text_features = self.encode_text(text, normalize=True) + image_features = self.encode_image(image, normalize=True) + image_logits = self.logit_scale.exp() * image_features @ text_features.T + text_logits = image_logits.T + return image_logits, text_logits + + def forward( + self, + image: Tensor, + text: Tensor, + ) -> Tuple[Tensor, Tensor, Tensor]: + + text_features = self.encode_text(text, normalize=True) + image_features = self.encode_image(image, normalize=True) + return image_features, text_features, self.logit_scale.exp() diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/dinov2_wrapper.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/dinov2_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..2689dc2fdd73c5be0b32d7f60abb6c11b760e3c2 --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/dinov2_wrapper.py @@ -0,0 +1,64 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +from typing import Sequence + +import torch + + +class DINOv2Wrapper(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + self.embed_dim = model.embed_dim + self.num_heads = model.num_heads + self.num_register_tokens = model.num_register_tokens + + # Same as the original forward, but assert is_training and rename x_norm_regtokens -> x_storage_tokens + def forward(self, img, is_training: bool): + assert is_training + H, W = img.shape[-2:] + P = self.model.patch_size + x_dict = self.model(img, is_training=True) + x_dict["h"] = h = H // P + x_dict["w"] = w = W // P + assert x_dict["x_norm_patchtokens"].shape[-2] == h * w + return x_dict + + # Same as the original get_intermediate_layers, but allow returining extra tokens (registers) + def get_intermediate_layers( + self, + x: torch.Tensor, + n: int | Sequence[int] = 1, # Layers or n last layers to take + reshape: bool = False, + return_class_token: bool = False, + return_register_tokens: bool = False, + norm=True, + ) -> tuple[torch.Tensor] | tuple[tuple[torch.Tensor, ...], ...]: + if self.model.chunked_blocks: + outputs = self.model._get_intermediate_layers_chunked(x, n) + else: + outputs = self.model._get_intermediate_layers_not_chunked(x, n) + if norm: + outputs = [self.model.norm(out) for out in outputs] + class_tokens = [out[:, 0] for out in outputs] + register_tokens = [out[:, 1 : 1 + self.model.num_register_tokens] for out in outputs] + outputs = [out[:, 1 + self.model.num_register_tokens :] for out in outputs] + if reshape: + B, _, h, w = x.shape + outputs = [ + out.reshape(B, h // self.model.patch_size, w // self.model.patch_size, -1) + .permute(0, 3, 1, 2) + .contiguous() + for out in outputs + ] + + if not return_class_token and not return_register_tokens: + return tuple(outputs) + if return_class_token and not return_register_tokens: + return tuple(zip(outputs, class_tokens)) + if not return_class_token and return_register_tokens: + return tuple(zip(outputs, register_tokens)) + return tuple(zip(outputs, class_tokens, register_tokens)) diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/text_tower.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/text_tower.py new file mode 100644 index 0000000000000000000000000000000000000000..a74c15f0eb64092728e245bbce08958e75ddc82c --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/text_tower.py @@ -0,0 +1,104 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import torch +from torch import nn, Tensor + +from dinov2.layers import ( + CausalAttentionBlock, +) + + +class TextHead(nn.Module): + def __init__( + self, + input_dim: int, + embed_dim: int, + num_heads: int, + num_blocks: int, + block_drop_prob: float, + is_causal: bool, + use_linear_projection: bool, + ): + super().__init__() + block_list = [nn.Identity()] + self.ln_final = nn.Identity() + if num_blocks > 0: + block_list = [ + CausalAttentionBlock( + dim=input_dim, + num_heads=num_heads, + is_causal=is_causal, + dropout_prob=block_drop_prob, + ) + for _ in range(num_blocks) + ] + self.ln_final = nn.LayerNorm(input_dim) + self.block_list = nn.ModuleList(block_list) + self.num_blocks = num_blocks + self.linear_projection = nn.Identity() + if input_dim != embed_dim or use_linear_projection: + self.linear_projection = nn.Linear(input_dim, embed_dim, bias=False) + + def init_weights(self): + if self.num_blocks > 0: + for i in range(self.num_blocks): + self.block_list[i].init_weights() + self.ln_final.reset_parameters() + if isinstance(self.linear_projection, nn.Linear): + nn.init.normal_(self.linear_projection.weight, std=self.linear_projection.in_features**-0.5) + + def forward(self, text_tokens: Tensor) -> Tensor: + for block in self.block_list: + text_tokens = block(text_tokens) + text_tokens = self.ln_final(text_tokens) + return self.linear_projection(text_tokens) + + +class TextTower(nn.Module): + def __init__( + self, + backbone: nn.Module, + freeze_backbone: bool, + embed_dim: int, + num_head_blocks: int, + head_blocks_is_causal: bool, + head_blocks_block_drop_prob: float, + tokens_pooler_type: str, + use_linear_projection: bool, + ): + super().__init__() + self.backbone = backbone + self.freeze_backbone = freeze_backbone + backbone_out_dim = backbone.embed_dim + self.backbone = backbone + self.head = TextHead( + backbone_out_dim, + embed_dim, + self.backbone.num_heads, + num_head_blocks, + head_blocks_block_drop_prob, + head_blocks_is_causal, + use_linear_projection, + ) + self.tokens_pooler_type = tokens_pooler_type + + def init_weights(self): + self.backbone.init_weights() + self.head.init_weights() + + def forward(self, token_indices: Tensor) -> Tensor: + text_tokens = self.backbone(token_indices) + text_tokens = self.head(text_tokens) + if self.tokens_pooler_type == "first": + features = text_tokens[:, 0] + elif self.tokens_pooler_type == "last": + features = text_tokens[:, -1] + elif self.tokens_pooler_type == "argmax": + assert token_indices is not None + features = text_tokens[torch.arange(text_tokens.shape[0]), token_indices.argmax(dim=-1)] + else: + raise ValueError(f"Unknown text tokens pooler type: {self.pooler_type}") + return features diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/text_transformer.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/text_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..7f836481b1a9fe98271d894423c4fd79644dd487 --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/text_transformer.py @@ -0,0 +1,72 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +from typing import Callable, Optional, Tuple + +import torch +from torch import nn, Tensor + + +from dinov2.layers import CausalAttentionBlock + + +class TextTransformer(nn.Module): + def __init__( + self, + context_length: int, + vocab_size: int, + dim: int, + num_heads: int, + num_layers: int, + ffn_ratio: float, + is_causal: bool, + ls_init_value: Optional[float] = None, + act_layer: Callable = nn.GELU, + norm_layer: Callable = nn.LayerNorm, + dropout_prob: float = 0.0, + ): + super().__init__() + self.vocab_size = vocab_size + self.embed_dim = dim + self.num_heads = num_heads + + self.token_embedding = nn.Embedding(vocab_size, dim) + self.positional_embedding = nn.Parameter(torch.empty(context_length, dim)) + self.dropout = nn.Dropout(dropout_prob) + self.num_layers = num_layers + block_list = [ + CausalAttentionBlock( + dim=dim, + num_heads=num_heads, + ffn_ratio=ffn_ratio, + ls_init_value=ls_init_value, + is_causal=is_causal, + act_layer=act_layer, + norm_layer=norm_layer, + dropout_prob=dropout_prob, + ) + for _ in range(num_layers) + ] + self.blocks = nn.ModuleList(block_list) + self.ln_final = norm_layer(dim) + + def init_weights(self): + nn.init.normal_(self.token_embedding.weight, std=0.02) + nn.init.normal_(self.positional_embedding, std=0.01) + init_attn_std = self.embed_dim**-0.5 + init_proj_std = (self.embed_dim**-0.5) * ((2 * self.num_layers) ** -0.5) + init_fc_std = (2 * self.embed_dim) ** -0.5 + for block in self.blocks: + block.init_weights(init_attn_std, init_proj_std, init_fc_std) + self.ln_final.reset_parameters() + + def forward(self, token_indices: Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + _, N = token_indices.size() + x = self.token_embedding(token_indices) + self.positional_embedding[:N] + x = self.dropout(x) + for block in self.blocks: + x = block(x) + x = self.ln_final(x) + return x diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/tokenizer.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..97a6b32eae76e5a1fcf63964d6ff9bf482dfeb0e --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/tokenizer.py @@ -0,0 +1,45 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import torch +from typing import List, Union + + +from dinov2.thirdparty.CLIP.clip.simple_tokenizer import SimpleTokenizer + + +class Tokenizer(SimpleTokenizer): + def __init__(self, vocab_path: str): + SimpleTokenizer.__init__(self, bpe_path=vocab_path) + + def tokenize(self, texts: Union[str, List[str]], context_length: int = 77) -> torch.LongTensor: + """ + Returns the tokenized representation of given input string(s) + + Parameters + ---------- + texts : Union[str, List[str]] + An input string or a list of input strings to tokenize + context_length : int + The context length to use; all CLIP models use 77 as the context length + + Returns + ------- + A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length] + """ + if isinstance(texts, str): + texts = [texts] + sot_token = self.encoder["<|startoftext|>"] + eot_token = self.encoder["<|endoftext|>"] + all_tokens = [[sot_token] + self.encode(text) + [eot_token] for text in texts] + result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) + + for i, tokens in enumerate(all_tokens): + if len(tokens) > context_length: + tokens = tokens[:context_length] # Truncate + tokens[-1] = eot_token + result[i, : len(tokens)] = torch.tensor(tokens) + + return result diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/vision_tower.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/vision_tower.py new file mode 100644 index 0000000000000000000000000000000000000000..8777aa9d27546da5c31cf2c995c6c1ed6a24ac37 --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/text/vision_tower.py @@ -0,0 +1,179 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +from functools import partial +from typing import Callable, Tuple + +import torch +from torch import nn, Tensor + +from dinov2.layers import ( + LayerScale, + NestedTensorBlock as AttentionBlock, + SwiGLUFFNAligned as SwiGLUFFN, +) + + +def init_weights_vit_timm(module: nn.Module, name: str = ""): + """ViT weight initialization, original timm impl (for reproducibility)""" + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + if isinstance(module, nn.LayerNorm): + module.reset_parameters() + if isinstance(module, LayerScale): + module.reset_parameters() + if isinstance(module, nn.Conv2d): + module.reset_parameters() + + +def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module: + if not depth_first and include_root: + fn(module=module, name=name) + for child_name, child_module in module.named_children(): + child_name = ".".join((name, child_name)) if name else child_name + named_apply( + fn=fn, + module=child_module, + name=child_name, + depth_first=depth_first, + include_root=True, + ) + if depth_first and include_root: + fn(module=module, name=name) + return module + + +class VisionHead(nn.Module): + def __init__( + self, + input_dim: int, + embed_dim: int, + num_heads: int, + num_blocks: int, + blocks_drop_path: float, + use_class_token: bool, + use_patch_tokens: bool, + use_linear_projection: bool, + ): + super().__init__() + block_list = [nn.Identity()] + self.ln_final = nn.Identity() + if num_blocks > 0: + block_list = [ + AttentionBlock( + input_dim, + num_heads, + ffn_layer=partial(SwiGLUFFN, align_to=64), + init_values=1e-5, + drop_path=blocks_drop_path, + ) + for _ in range(num_blocks) + ] + self.ln_final = nn.LayerNorm(input_dim) + self.block_list = nn.ModuleList(block_list) + self.num_blocks = num_blocks + multiplier = 2 if use_class_token and use_patch_tokens else 1 + self.linear_projection = nn.Identity() + if multiplier * input_dim != embed_dim or use_linear_projection: + assert embed_dim % multiplier == 0, f"Expects {embed_dim} to be divisible by {multiplier}" + self.linear_projection = nn.Linear(input_dim, embed_dim // multiplier, bias=False) + + def init_weights(self): + if self.num_blocks > 0: + for i in range(self.num_blocks): + block = self.block_list[i] + named_apply(init_weights_vit_timm, block) + self.ln_final.reset_parameters() + if isinstance(self.linear_projection, nn.Linear): + nn.init.normal_(self.linear_projection.weight, std=self.linear_projection.in_features**-0.5) + + def forward(self, image_tokens: Tensor) -> Tensor: + for block in self.block_list: + image_tokens = block(image_tokens) + image_tokens = self.ln_final(image_tokens) + return self.linear_projection(image_tokens) + + +class VisionTower(nn.Module): + def __init__( + self, + backbone: nn.Module, + freeze_backbone: bool, + embed_dim: int, + num_head_blocks: int, + head_blocks_block_drop_path: float, + use_class_token: bool, + use_patch_tokens: bool, + patch_token_layer: int, + patch_tokens_pooler_type: str, + use_linear_projection: bool, + ): + super().__init__() + self.backbone = backbone + self.freeze_backbone = freeze_backbone + self.use_class_token = use_class_token + self.use_patch_tokens = use_patch_tokens + self.patch_token_layer = patch_token_layer + self.patch_tokens_pooler_type = patch_tokens_pooler_type + self.num_register_tokens = 0 + if hasattr(self.backbone, "num_register_tokens"): + self.num_register_tokens = self.backbone.num_register_tokens + elif hasattr(self.backbone, "n_storage_tokens"): + self.num_register_tokens = self.backbone.n_storage_tokens + backbone_out_dim = self.backbone.embed_dim + self.head = VisionHead( + backbone_out_dim, + embed_dim, + self.backbone.num_heads, + num_head_blocks, + head_blocks_block_drop_path, + use_class_token, + use_patch_tokens, + use_linear_projection, + ) + + def init_weights(self): + if not self.freeze_backbone: + self.backbone.init_weights() + self.head.init_weights() + + def get_backbone_features(self, images: Tensor) -> Tuple[Tensor, Tensor, Tensor]: + tokens = self.backbone.get_intermediate_layers( + images, + n=self.patch_token_layer, + return_class_token=True, + return_register_tokens=True, + ) + class_token = tokens[-1][1] + patch_tokens = tokens[0][0] + register_tokens = tokens[0][2] + return class_token, patch_tokens, register_tokens + + def get_class_and_patch_tokens(self, images: Tensor) -> Tuple[Tensor, Tensor]: + class_token, patch_tokens, register_tokens = self.get_backbone_features(images) + image_tokens = self.head(torch.cat([class_token.unsqueeze(1), register_tokens, patch_tokens], dim=1)) + class_token, patch_tokens = image_tokens[:, 0], image_tokens[:, self.num_register_tokens + 1 :] + return class_token, patch_tokens + + def forward(self, images: Tensor) -> Tensor: + class_token, patch_tokens = self.get_class_and_patch_tokens(images) + features = [] + if self.use_class_token: + features.append(class_token) + if self.use_patch_tokens: + if self.patch_tokens_pooler_type == "mean": + features.append(torch.mean(patch_tokens, dim=1)) + elif self.patch_tokens_pooler_type == "max": + features.append(torch.max(patch_tokens, dim=1).values) + elif self.patch_tokens_pooler_type == "gem": + power = 3 + eps = 1e-6 + patch_tokens_power = patch_tokens.clamp(min=eps).pow(power) + features.append(torch.mean(patch_tokens_power, dim=1).pow(1 / power)) + else: + raise ValueError(f"Unknown patch tokens pooler type: {self.patch_tokens_pooler_type}") + return torch.cat(features, dim=-1) diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/xray_dino/__pycache__/backbones.cpython-311.pyc b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/xray_dino/__pycache__/backbones.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ceac559bc950395c38e6eb6f0e53c422ad91850a Binary files /dev/null and b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/xray_dino/__pycache__/backbones.cpython-311.pyc differ diff --git a/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/xray_dino/backbones.py b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/xray_dino/backbones.py new file mode 100644 index 0000000000000000000000000000000000000000..f49b54be1c3aa5a0fbe17427b86f9277598aa925 --- /dev/null +++ b/migrator/checkpoints/hub/facebookresearch_dinov2_main/dinov2/hub/xray_dino/backbones.py @@ -0,0 +1,28 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the licence +# found in the LICENSE_XRAY_DINO_MODEL file in the root directory of this source tree. + +from typing import Union + +from ..backbones import Weights, _make_dinov2_model + + +def xray_dino_vitl16(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.XRAY_DINO, **kwargs): + """ + XRay-DINO ViT-L/16 model (optionally) pretrained on the XRay-DINO dataset. + """ + return _make_dinov2_model( + arch_name="vit_large", + patch_size=16, + img_size=512, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + block_chunks=4, + pretrained=pretrained, + weights=weights, + hash="ad31c2b0", + check_hash=True, + **kwargs, + ) diff --git a/migrator/code/mv-sam3d-for-6d-v2/IMPLEMENTATION.md b/migrator/code/mv-sam3d-for-6d-v2/IMPLEMENTATION.md new file mode 100644 index 0000000000000000000000000000000000000000..dc3e1efac0bb3430c97b1dddd9fddadbd138bffe --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/IMPLEMENTATION.md @@ -0,0 +1,208 @@ +# mv-sam3d-for-6d-v2 — Implementation + +A clean, self-contained rebuild of the multi-view (multi-camera) DexYCB 6D-pose +flow model. Conceptually: *take SAM-3D (`sam3d-for-6d`), drop the shape-K/V +forcing, change the conditioning fed to the flow transformer, for multi-view +DexYCB 6D pose.* Everything that already exists in the fork is **vendored +verbatim** and warm-started from pretrained weights; only the genuinely-new +pieces (DexYCB multi-view dataset, calibration, training loop, latent-cache +builder, and the one per-view embedding) are written from scratch. + +All paths are absolute. Python env: +`/NHNHOME/data/jonghoon/miniconda3/envs/sam3d-objects/bin/python`. +GPU work is **sbatch-only**; nothing sets `CUDA_VISIBLE_DEVICES`. The repo is +self-contained: `PYTHONPATH=` only — no dependency on the fork's path, +no TRELLIS. Every vendored-package `__init__.py` is empty (the fork's lidra init +chain is NOT pulled in), so imports are clean. + +--- + +## Model design (locked) + +* **Backbone:** `SparseStructureFlowTdfyWrapper` (MoT DiT), `mode="default"` + (fully-symmetric joint self-attention), **`is_shortcut_model=False`** = plain + rectified flow, v-prediction, `sigma_min=0`. `model_channels=cond_channels=1024`, + `num_blocks=24`, `num_heads=16`, `mlp_ratio=4`, `pe_mode=ape`, + `qk_rms_norm=True`, `time_scale=1000`. +* **Two MoT streams:** + * `shape` — SS latent `[B,8,16,16,16]` → `view(B,8,-1).permute(0,2,1)` → + `[B,4096,8]` (row-major to match `ShapePositionEmbedder` meshgrid `ij`, + res16/patch1). + * pose — one shared transformer over `{6drotation_normalized[1,6], + translation[1,3], scale[1,1]}` via `latent_share_transformer`. The shared + group is **keyed `6drotation_normalized`** (matches the pretrained block-stream + name so all per-block weights warm-start; see warm-start note below). + * **NO shape-K/V injection.** +* **Loss** = weighted v-pred MSE, all stream weights **1.0** (shape / + 6drotation_normalized / translation / scale) — matches the pretrained backbone's + training regime (SAM-3D fork uses rot=trans=scale=1.0). + +## Conditioning / multi-view (the genuinely-new part) + +One timestamp, multiple synchronized DexYCB cameras sharing one object pose. The +**anchor** camera defines the reference frame; others are **aux**. Everything is +in the anchor camera's PyTorch3D frame. The SAME `EmbedderFuser` instance runs +per view, and a per-view **camera embedding** `mlp_cam(extrinsics_v)` is added to +that view's tokens (replacing the old content-free `view_emb` slot table). A +2-layer MLP maps each view's relative aux→anchor extrinsics — rot6d(6) + +SSI-normalized camera centre(3), 9-D, anchor = identity — to `model_channels`, +with the output layer zero-initialized so step-0 conditioning equals the no-tag +baseline. This binds the camera frame to every token of the view (including the +geometry-free RGB/mask tokens) and is camera-consistent across the random aux +subset. Anchor +passes all 6 fuser kwargs (cropped+full of image/mask/pointmap); each aux passes +the 3 cropped kwargs. All views' tokens are concatenated into one condition +sequence the DiT cross-attends to. Pointmaps KEEP NaN at invalid pixels so the +pretrained `invalid_xyz_token` pathway fires; pointmaps/masks use nearest-mode +(NaN-preserving) interpolation. + +--- + +## File map + +``` +mvsam3d/ + data/ + canonical.py # CAD -> canonical voxel (vendored voxelize_mesh) + GT layout math + dexycb_calib.py # extrinsics/intrinsics; cross-camera P3D transform (aux->anchor) + dexycb_dataset.py # MultiViewDexYCB + collate_multiview [from scratch] + model/ + condition.py # MultiViewConditionEmbedder (real EmbedderFuser + per-view view_emb) + flow_model.py # MVSam3DFlow: MoT joint flow + flow_step + loss + warm-start + ss_vae.py # REAL SAM-3D SS encoder/decoder loaders (ss_encoder.ckpt / ss_decoder.ckpt) + train/ + train.py # YAML-driven DDP training loop (grad-accum + periodic val) [from scratch] + validate.py # sampler-based validation: decode->64^3 IoU + pose metrics [from scratch] +sam3d_objects/ # VENDORED closure (byte-identical to the fork; empty __init__ chain) + data/utils.py + model/layers/llama3/ff.py + model/backbone/dit/embedder/{dino,embedder_fuser,pointmap,point_remapper}.py + model/backbone/tdfy_dit/models/{mm_latent,mot_sparse_structure_flow,timestep_embedder, + sparse_structure_vae}.py + model/backbone/tdfy_dit/modules/{spatial,utils,norm}.py + transformer/ + attention/ + model/backbone/generator/flow_matching/solver.py # Euler ODE solver (reused for sampling) + training/stage1_flow_lora/distributed.py # DDP/DeepSpeed primitives (reused) + geometry/{geometry,rotation_constants}.py # column-convention rot6d utils + pipeline/inference_utils.py # voxelize_mesh / normalize_mesh_verts / trimesh2o3d_mesh (verbatim) +tools/ + build_canonical_latents.py # 21 objects -> cache/canonical/*.pt (GPU) + validate_model.py # pretrained-load + fwd/bwd + grad audit (GPU) + verify_rotation_roundtrip.py # rot6d encode/decode inverse check (CPU) + verify_gt_reproject.py # GT modal-reproject IoU vs seg (CPU) + amodal_reproject.py # GT amodal IoU vs independent mesh render (CPU) + diagnose_reproject.py # GT reproject recall vs seg (CPU) +configs/train.yaml +sbatch/train.sh # training launcher +cache/canonical/.pt # {latent[8,16,16,16], occupancy[64^3] uint8, max_extent, center_cad} +``` + +## Inherited (vendored, byte-identical to fork) vs from-scratch + +| Component | Source | Notes | +|---|---|---| +| MoT backbone, `Latent`, position embedders | vendored | `mot_sparse_structure_flow.py`, `mm_latent.py` | +| `EmbedderFuser`, `Dino`, `PointPatchEmbed`, `FeedForward`, `PointRemapper` | vendored | embedder closure | +| SS VAE `SparseStructureEncoder/Decoder` | vendored | `sparse_structure_vae.py` | +| rot6d utils (column convention) + constants | vendored | `geometry.py`, `rotation_constants.py` | +| `voxelize_mesh` / `normalize_mesh_verts` / `trimesh2o3d_mesh` | vendored | verbatim functions only | +| `Euler` ODE solver (`flow_matching/solver.py`) | vendored | reused by `MVSam3DFlow.sample` | +| DDP/DeepSpeed primitives (`distributed.py`) | vendored | reused by `train.py` | +| `MultiViewConditionEmbedder` | from scratch | thin wrapper; only `mlp_cam` (per-view camera embedding from extrinsics) is a new module | +| `MVSam3DFlow` (config/flow_step/loss/warm-start/sample/predict) | from scratch | plain flow, no shape-K/V | +| DexYCB dataset / calibration / canonical GT | from scratch | no fork equivalent | +| latent-cache builder, DDP training loop, validation, config | from scratch | | + +## Pretrained weights + +* **SS encoder/decoder = the REAL SAM-3D ones** (`checkpoints/hf/ss_encoder.ckpt`, + `ss_decoder.ckpt`), loaded with **0 missing / 0 unexpected** keys through the + vendored `SparseStructureEncoder`/`SparseStructureDecoder`. GT latents use the + posterior MEAN (`sample_posterior=False`). GT latents are validated by + roundtrip through the **standalone** `ss_decoder.ckpt` (expect IoU≈1.0). The + decoder embedded in `ss_generator.ckpt` is TRELLIS-identical and gives IoU≈0 on + these latents — reported for contrast, never used to validate. +* **Warm-start from `ss_generator.ckpt`:** the condition embedder + (`_base_models.condition_embedder.*`, 722 tensors) loads fully into the frozen + `EmbedderFuser`; the torso (`_base_models.generator.reverse_fn.backbone.*`) + loads **936 / 936** tensors — `t_embedder`, both `latent_mapping` streams, and + **all** per-block transformer weights for both MoT streams (`shape` and the + shared pose stream). **The shared pose group is named `6drotation_normalized` + (NOT `pose`)** so its per-block keys match the pretrained block-stream naming + (the fork names a shared MoT group after one of its member latents); naming it + `pose` previously left all 432 pose-stream block tensors at random init. Only + the shortcut `d_embedder` (4 tensors) and the unused `translation_scale` stream + have no home here and are skipped cleanly (never forced). After this, **the + ONLY from-scratch module in the whole model is `mlp_cam`** (the per-view camera + embedding; no pretrained source, output layer zero-init), which injects the + relative camera extrinsics that tag the new multi-view axis. + +## Data / GT + +* `dexycb_root=/NHNHOME/data/jonghoon/pose6daug/dataset/dexycb`. `ycb_ids`/`seg` + are 1-based class indices; object mask = `(seg==ycb_id)`. Depth uint16 mm /1000, + clip to [0.1,3.0] m, invalid→NaN. Pointmap = unproject(aligned_depth) in OpenCV + → PyTorch3D via `*[-1,-1,1]`. Aux pointmaps transformed into the anchor frame + via DexYCB extrinsics BEFORE one shared SSI on the anchor. +* **SSI** = base/default isotropic: `shift=[0,0,median_z]`, + `scale=mean(|p-shift|)` over finite points (matches the fork's BASE + `ScaleShiftInvariant.get_scale_and_shift`). Isotropic is required because the + GT pose is a rigid similarity; a per-dim anisotropic SSI would shear the + rotation cue. One SSI computed on the anchor is applied to all fused views. +* **GT layout (anchor cam, P3D):** with `M=[[1,0,0],[0,0,-1],[0,1,0]]`, + `D=diag(-1,-1,1)`: `R_p3d=D·R_cv`, `t_p3d=D·t_cv`; then `R=R_dex·Mᵀ`, + `t=R_dex·center_cad+t_dex`, `s=max_extent`. `center_cad`/`max_extent` are + derived from the vendored `voxelize_mesh` output (`max_extent=1/scale`, + `center_cad=Mᵀ·center_rot`). +* **Rotation 6D:** encode `normalize_rot6d(matrix_to_rot6d(R))`; decode + `rot6d_to_matrix(denormalize_rot6d(x))` — the vendored COLUMN-convention + functions used consistently in BOTH dataset and model. (The prior attempt's + interleaved `R[:,:2].reshape(6)` encode was inconsistent with the + column-convention decode — a real bug, fixed here.) + +## Distributed + validation + +* **DDP.** `sbatch/train.sh` launches `torch.distributed.run --nproc_per_node=$NPROC` + (`NPROC` defaults to `SLURM_GPUS_ON_NODE`). The generic DDP/DeepSpeed primitives + are the VENDORED `distributed.py` (`init_distributed`/`setup_runtime`/`barrier`/ + `reduce_scalar_dict`/`cleanup_distributed`); `train.py` is the only from-scratch + glue. Device selection is by `LOCAL_RANK` (never `CUDA_VISIBLE_DEVICES`). + `WORLD_SIZE=1` transparently skips the process group + DDP wrap. Training routes + through `model(batch)` (== `flow_step`) so DDP all-reduces gradients; supports + `train.grad_accum_steps` (micro-step `no_sync` until the boundary). +* **Sampling.** `MVSam3DFlow.sample` integrates the rectified-flow ODE from noise + (t=0) to data (t=1) with the VENDORED `Euler` solver, exact inverse of `flow_step` + (same `time_scale`, `mode="default"`). `predict` decodes the streams to metric + `{shape_latent, rotation[3,3], translation[3], scale}`. +* **Validation** (`validate.py`, every `train.val_every` steps on a disjoint, + deterministic held-out split): runs `predict` then reports `shape_iou` (decode + pred latent through the standalone `ss_decoder.ckpt` vs GT canonical occupancy), + `rot_deg` (geodesic), `trans_cm`, `scale_rel` — all summed-with-count and + all-reduced across ranks for an exact mean. Val uses a fixed `max_aux` aux views + so `view_emb` is never indexed past its trained range. +* **Verified (job 4331, 2× B200, world=2):** DDP process group + gradient reduce + + barrier, 4 steps finite loss, validation at steps 2 & 4 (sampler→decode→metrics, + reduced n=2), checkpoint saved, clean exit (COMPLETED). + +## How to run + +```sh +# multi-GPU: request N GPUs; one process per GPU is launched automatically +sbatch --gres=gpu:b200:4 sbatch/train.sh --config configs/train.yaml +# single GPU (default sbatch header) or quick smoke: +sbatch sbatch/train.sh --config configs/train.yaml --set train.steps=3 --set data.subjects=[subject-01] +``` + +View selection lives under `data:`: `random_views: true` with +`min_aux_views`/`max_aux_views` (total = 1 + k aux), or `random_views: false` +for the first `num_aux_views` cameras. `model.max_views` is derived in `train.py` +(now API-only; `mlp_cam` has no per-slot table, so the view count is unbounded). +Camera embedding: `condition.cam_hidden_dim` sets the `mlp_cam` hidden width. +Validation knobs: `data.val_fraction`/`split_seed`, `train.val_every`/ +`val_num_steps`/`val_max_batches`. Distributed knobs: `train.backend` +(`ddp`|`deepspeed`), `grad_accum_steps`, `amp`/`amp_dtype`. LR schedule: +`train.scheduler` (`cosine`|`constant`), `warmup_steps`, `min_lr` — REUSES the +vendored `train_utils.build_scheduler` (warmup→cosine→min_lr). Resume: +`train.resume_from=` restores trainables + optimizer + scheduler + scaler + +step via the vendored `checkpointing.{save,load}_checkpoint` (frozen DINO/fuser are +reconstructed from `pretrained_ckpt`, not stored). Verified: jobs 4385 (0→8) + +4386 (resume@3) — LR curve continuous across resume, optimizer state restored. diff --git a/migrator/code/mv-sam3d-for-6d-v2/configs/heldout.yaml b/migrator/code/mv-sam3d-for-6d-v2/configs/heldout.yaml new file mode 100644 index 0000000000000000000000000000000000000000..19a4e2d64fbbfa7adb65b9ee880dc87c01e0ec43 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/configs/heldout.yaml @@ -0,0 +1,70 @@ +# Held-out manifest for integrated training. Two kinds (distinction matters): +# ood_datasets : whole-dataset OOD eval; objects MAY overlap training +# (HO-3D is YCB) -> NO object-disjointness enforced. +# unseen_object_datasets : whole dataset held out AND its object identities +# reserved (removed from every train dataset). Objects +# are non-YCB so this is safe. +ood_datasets: [ho3d, ycbm, hope] +unseen_object_datasets: [housecat6d] +# GraspNet test_novel object holdout (M1 resolved). GraspNet stays a `role:train` +# dataset built on the full `graspnet` root, but the 42 objects that appear in the +# official test_novel scenes (160-189) are reserved here so they are removed from +# EVERY train dataset (leak-proof unseen-object holdout for GraspNet-novel eval). +# Derivation (authoritative, matches graspnetAPI): test_novel = sceneIds 160-189 +# (graspnetAPI/graspnet.py); getObjIds() = union of each scene's object_id_list.txt. +# We took that union over scenes 160-189 (42 distinct 0-based object ids). The +# adapter key is `graspnet/{cls_value-1:03d}` and cls_value (meta.mat cls_indexes) +# == object_id+1 (verified on scenes 160/165/170/180/189), so cls_value-1 == +# object_id_list value; keys below are `graspnet/` for the 42 novel ids. +# NOTE: 9 of these (017,020,026,027,030,051,062,063,066) also appear in +# train/seen/similar scenes — the global object-disjointness filter removes them +# from training everywhere, which is exactly why this reservation is required. +unseen_object_keys_extra: +- graspnet/003 +- graspnet/004 +- graspnet/006 +- graspnet/017 +- graspnet/019 +- graspnet/020 +- graspnet/024 +- graspnet/026 +- graspnet/027 +- graspnet/028 +- graspnet/030 +- graspnet/031 +- graspnet/032 +- graspnet/033 +- graspnet/035 +- graspnet/045 +- graspnet/047 +- graspnet/049 +- graspnet/051 +- graspnet/055 +- graspnet/062 +- graspnet/063 +- graspnet/065 +- graspnet/066 +- graspnet/067 +- graspnet/071 +- graspnet/072 +- graspnet/073 +- graspnet/074 +- graspnet/075 +- graspnet/076 +- graspnet/077 +- graspnet/078 +- graspnet/079 +- graspnet/080 +- graspnet/081 +- graspnet/082 +- graspnet/083 +- graspnet/084 +- graspnet/085 +- graspnet/086 +- graspnet/087 +never_train: [ho3d, housecat6d] # adapters refused for split:train + +roots: + housecat6d: ${MIGRATOR_DATA_ROOT}/housecat6d/extracted +adapter_kwargs: + housecat6d: {split: test} diff --git a/migrator/code/mv-sam3d-for-6d-v2/configs/train.yaml b/migrator/code/mv-sam3d-for-6d-v2/configs/train.yaml new file mode 100644 index 0000000000000000000000000000000000000000..066a454438ab0b8671ca3e66eec6acae3c9bbd05 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/configs/train.yaml @@ -0,0 +1,141 @@ +# Training config for the multi-view 6D-pose flow model (v2). +# ALL training-related hyperparameters live here; train.py reads them and nothing +# training-related is hardcoded in Python. Launch via: +# sbatch sbatch/train.sh --config configs/train.yaml +# Per-key overrides: --set key.subkey=value (repeatable), e.g. --set train.steps=3 + +# --------------------------------------------------------------------------- # +# data -> MultiViewDexYCB +# --------------------------------------------------------------------------- # +data: + dexycb_root: ${MIGRATOR_DATA_ROOT}/dexycb + canonical_cache: ${MIGRATOR_CACHE}/canonical + ssi_mode: scene # scene | object_centric — SSI convention for pointmaps + # + t/s targets + cam features. "scene" = base SSI + # (shift=[0,0,median z], all pre-2026-07-17 ckpts). + # "object_centric" = vendored ObjectCentricSSI (object- + # median shift, scene scale; release-pipeline pairing). + # A ckpt MUST be evaluated with its training mode. + frame_stride: 5 # sample every Nth frame in each sequence + anchor_sweep: true # every camera serves as anchor (else only the first) + subjects: null # null = all subjects; or a list of name substrings + image_size: 518 # DINOv2 input resolution (= IMG_SIZE) + # --- view selection --- + num_aux_views: 7 # legacy fixed count, used only when random_views=false + random_views: true # draw a random count + random subset of aux cameras per sample + min_aux_views: 1 # inclusive lower bound on aux count when random_views=true + max_aux_views: 3 # inclusive upper bound -> total views = 1 + [1..3] = 2..4 + # --- depth preprocessing constants --- + depth_scale: 1000.0 # DexYCB depth PNG is uint16 millimetres + depth_near: 0.1 # metres; clip implausibly-near depth + depth_far: 3.0 # metres; clip implausibly-far depth + # --- train/val split (deterministic; disjoint) --- + val_fraction: 0.02 # fraction of samples held out for validation (0 = no val) + split_seed: 0 # fixed seed for the train/val partition + +# --------------------------------------------------------------------------- # +# model -> MVSam3DFlow +# NOTE: max_views is DERIVED in train.py as +# (max_aux_views if random_views else num_aux_views) + 1 +# SHAPE_RES (16) and SHAPE_CHANNELS (8) are locked by the SAM-3D SS VAE and +# remain code constants in flow_model.py. +# --------------------------------------------------------------------------- # +model: + model_channels: 1024 + cond_channels: 1024 + num_blocks: 24 + num_heads: 16 + sigma_min: 0.0 # rectified flow; 0 = straight OT path + time_scale: 1000.0 # flow time scaling fed to the backbone + # Timestep sampling for the flow loss. Default follows the ORIGINAL SAM-3D + # generator pretraining schedule (ss_generator.yaml training_time_sampler_fn = + # lognorm_sampler mean=-1, std=1). v2's t convention (t=0 noise -> t=1 data) is + # identical to the FlowMatching module's, so mean maps directly (no sign flip). + # Set timestep_sampler=uniform to reproduce the stage1_flow_lora harness. + timestep_sampler: logit_normal # logit_normal | uniform + timestep_lognorm_mean: -1.0 + timestep_lognorm_std: 1.0 + # Warm-start from SAM-3D's pretrained ss_generator. Loads BOTH the condition + # embedder (722 tensors -> frozen EmbedderFuser) AND the DiT torso (shape / + # 6drotation_normalized / translation streams + blocks + t_embedder). Set to + # null to train from scratch. + pretrained_ckpt: ${MIGRATOR_CHECKPOINTS}/hf/ss_generator.ckpt + # Standalone SAM-3D SS decoder (authoritative latent space) — used ONLY by + # validation to decode predicted shape latents to 64^3 occupancy for IoU. + ss_decoder_ckpt: ${MIGRATOR_CHECKPOINTS}/hf/ss_decoder.ckpt + +# --------------------------------------------------------------------------- # +# condition -> MultiViewConditionEmbedder (thin multi-view wrapper around +# SAM-3D's real EmbedderFuser). Mirrors ss_generator.yaml's +# module.condition_embedder.backbone exactly: 3 embedders (2x Dino + +# PointPatchEmbed), learned pos-embedding, projection mult 4.0, dropout 0.1 +# with the pointmap/rgb_pointmap drop weight, freeze: true. +# --------------------------------------------------------------------------- # +condition: + dino_model: dinov2_vitl14_reg + dino_input_size: 518 + pointmap_input_size: 256 + pointmap_patch_size: 8 + pointmap_embed_dim: 512 + projection_hidden_mult: 4.0 # = projection_net_hidden_dim_multiplier + dropout_prob: 0.1 # fuser modality dropout (pointmap/rgb_pointmap) + # Per-view camera embedding (mlp_cam): maps each view's relative aux->anchor + # extrinsics (rot6d + SSI-normalized camera centre, 9-D; anchor = identity) to + # model_channels and adds it to that view's tokens. Replaces the old + # content-free view_emb index table. Output layer is zero-init (no perturbation + # at step 0). This is the ONLY from-scratch module in the model. + cam_hidden_dim: 256 # hidden width of the 2-layer mlp_cam + # SAM-3D freezes its ENTIRE condition embedder (freeze: true). We mirror that: + # the whole EmbedderFuser (both DINOs, PointPatchEmbed, all projection_nets, + # idx_emb) is frozen after warm-start. The per-view camera embedding (mlp_cam, + # the extrinsics-derived module that replaced view_emb) is always trainable + # (no pretrained source). freeze: false would fine-tune the entire condition + # pathway. + freeze: true + freeze_dino: true # DINO backbone freeze at construction time + +# --------------------------------------------------------------------------- # +# loss_weights -> MVSam3DFlow.loss_weights (per-stream v-prediction MSE weights) +# --------------------------------------------------------------------------- # +loss_weights: + # All 1.0 — matches the pretrained backbone's training regime (SAM-3D fork uses + # rot=trans=scale=1.0) so warm-started weights are fine-tuned under the same + # objective balance. Rotation has the largest raw loss (hardest stream), so 1.0 + # gives it the gradient share it needs rather than down-weighting it. + shape: 1.0 + 6drotation_normalized: 1.0 + translation: 1.0 + scale: 1.0 + +# --------------------------------------------------------------------------- # +# optim + train loop +# --------------------------------------------------------------------------- # +optim: + lr: 1.0e-4 + weight_decay: 0.0 # AdamW + grad_clip: 1.0 # max grad norm + +train: + batch_size: 1 # per-GPU micro-batch + steps: 100000 + num_workers: 8 + # --- LR schedule (reused fork build_scheduler: warmup -> cosine -> min_lr) --- + scheduler: cosine # cosine | constant + warmup_steps: 1000 # linear warmup steps before cosine decay + min_lr: 1.0e-6 # cosine decay floor + # --- resume --- + resume_from: "" # path to a checkpoint to resume (empty = fresh start) + log_every: 20 + ckpt_every: 2000 + ckpt_dir: ${MIGRATOR_CACHE}/ckpt + amp: true # autocast on cuda (dtype below) + amp_dtype: bf16 # bf16 | fp16 (fp16 enables GradScaler) + seed: 0 # seeds torch/numpy/random at startup (offset by rank) + # --- distributed (DDP via torchrun; world_size=1 transparently skips it) --- + backend: ddp # ddp | deepspeed + deepspeed_config: "" # path to a DeepSpeed JSON (only used when backend=deepspeed) + grad_accum_steps: 1 # micro-steps per optimizer step + # --- validation (model's own Euler sampler + decode) --- + val_every: 2000 # run validation every N optimizer steps (0 = never) + val_num_steps: 50 # Euler ODE steps for sampling during validation + val_max_batches: 16 # val batches per rank per validation pass diff --git a/migrator/code/mv-sam3d-for-6d-v2/configs/train_integrated.yaml b/migrator/code/mv-sam3d-for-6d-v2/configs/train_integrated.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b8634e6f374e387ab0f0086800ea16866e0f5791 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/configs/train_integrated.yaml @@ -0,0 +1,127 @@ +data: + dexycb_root: ${MIGRATOR_DATA_ROOT}/dexycb + canonical_cache: ${MIGRATOR_CACHE}/canonical + frame_stride: 5 + anchor_sweep: true + subjects: null + image_size: 518 + num_aux_views: 7 + random_views: true + min_aux_views: 1 + max_aux_views: 3 + depth_scale: 1000.0 + depth_near: 0.1 + depth_far: 3.0 + val_fraction: 0.02 + split_seed: 0 + heldout_manifest: ${MIGRATOR_MV_REPO}/configs/heldout.yaml + datasets: + - name: dexycb + adapter: dexycb + root: ${MIGRATOR_DATA_ROOT}/dexycb + role: train + weight: 1.0 + - name: pace + adapter: pace + root: ${MIGRATOR_DATA_ROOT}/pace + role: train + weight: 3.0 + data_split: test + - name: graspnet + adapter: graspnet + root: ${MIGRATOR_DATA_ROOT}/graspnet + role: train + weight: 2.0 + - name: ycbv + adapter: ycbv + root: ${MIGRATOR_DATA_ROOT}/ycbv + role: train + weight: 1.0 + data_split: test + # Adapters implemented + latent-complete + reprojection-gate >=0.8 (2026-07-17). + - name: hocap + adapter: hocap + root: ${MIGRATOR_DATA_ROOT}/hocap + role: train + weight: 1.0 + pseudo_gt: true # FoundationPose+SDF pseudo-GT (train-only; never eval) + - name: h2o + adapter: h2o + root: ${MIGRATOR_DATA_ROOT}/h2o + role: train + weight: 0.5 # only 8 objects + - name: hograspnet + adapter: hograspnet + root: ${MIGRATOR_DATA_ROOT}/hograspnet + role: train + weight: 1.0 + frame_stride: 15 # grasp-hold is quasi-static; keep enumeration tractable + - name: dexh2r + adapter: dexh2r + root: ${MIGRATOR_DATA_ROOT}/dexh2r + role: train + weight: 1.0 + pseudo_gt: true # global-registration+ICP pseudo-GT (train-only); enumerates extracted subjects + - name: contactpose + adapter: contactpose + root: ${MIGRATOR_DATA_ROOT}/contactpose + role: train + weight: 0.5 # static grasps; only bowl present on disk so far + pseudo_gt: true + - name: housecat6d + adapter: housecat6d + root: ${MIGRATOR_DATA_ROOT}/housecat6d/extracted + role: heldout + data_split: test +model: + model_channels: 1024 + cond_channels: 1024 + num_blocks: 24 + num_heads: 16 + sigma_min: 0.0 + time_scale: 1000.0 + timestep_sampler: logit_normal + timestep_lognorm_mean: -1.0 + timestep_lognorm_std: 1.0 + pretrained_ckpt: ${MIGRATOR_CHECKPOINTS}/hf/ss_generator.ckpt + ss_decoder_ckpt: ${MIGRATOR_CHECKPOINTS}/hf/ss_decoder.ckpt +condition: + dino_model: dinov2_vitl14_reg + dino_input_size: 518 + pointmap_input_size: 256 + pointmap_patch_size: 8 + pointmap_embed_dim: 512 + projection_hidden_mult: 4.0 + dropout_prob: 0.1 + cam_hidden_dim: 256 + freeze: true + freeze_dino: true +loss_weights: + shape: 1.0 + 6drotation_normalized: 1.0 + translation: 1.0 + scale: 1.0 +optim: + lr: 0.0001 + weight_decay: 0.0 + grad_clip: 1.0 +train: + batch_size: 1 + steps: 100000 + num_workers: 8 + scheduler: cosine + warmup_steps: 1000 + min_lr: 1.0e-06 + resume_from: '' + log_every: 20 + ckpt_every: 2000 + ckpt_dir: ${MIGRATOR_CACHE}/ckpt + amp: true + amp_dtype: bf16 + seed: 0 + backend: ddp + deepspeed_config: '' + grad_accum_steps: 1 + val_every: 2000 + val_num_steps: 50 + val_max_batches: 16 diff --git a/migrator/code/mv-sam3d-for-6d-v2/configs/train_integrated_object_ssi.yaml b/migrator/code/mv-sam3d-for-6d-v2/configs/train_integrated_object_ssi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b0531326de7c508bf4a0b5ecf13f1ce5eb8b84b3 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/configs/train_integrated_object_ssi.yaml @@ -0,0 +1,128 @@ +data: + ssi_mode: object_centric # SSI ablation + dexycb_root: ${MIGRATOR_DATA_ROOT}/dexycb + canonical_cache: ${MIGRATOR_CACHE}/canonical + frame_stride: 5 + anchor_sweep: true + subjects: null + image_size: 518 + num_aux_views: 7 + random_views: true + min_aux_views: 1 + max_aux_views: 3 + depth_scale: 1000.0 + depth_near: 0.1 + depth_far: 3.0 + val_fraction: 0.02 + split_seed: 0 + heldout_manifest: ${MIGRATOR_MV_REPO}/configs/heldout.yaml + datasets: + - name: dexycb + adapter: dexycb + root: ${MIGRATOR_DATA_ROOT}/dexycb + role: train + weight: 1.0 + - name: pace + adapter: pace + root: ${MIGRATOR_DATA_ROOT}/pace + role: train + weight: 3.0 + data_split: test + - name: graspnet + adapter: graspnet + root: ${MIGRATOR_DATA_ROOT}/graspnet + role: train + weight: 2.0 + - name: ycbv + adapter: ycbv + root: ${MIGRATOR_DATA_ROOT}/ycbv + role: train + weight: 1.0 + data_split: test + # Adapters implemented + latent-complete + reprojection-gate >=0.8 (2026-07-17). + - name: hocap + adapter: hocap + root: ${MIGRATOR_DATA_ROOT}/hocap + role: train + weight: 1.0 + pseudo_gt: true # FoundationPose+SDF pseudo-GT (train-only; never eval) + - name: h2o + adapter: h2o + root: ${MIGRATOR_DATA_ROOT}/h2o + role: train + weight: 0.5 # only 8 objects + - name: hograspnet + adapter: hograspnet + root: ${MIGRATOR_DATA_ROOT}/hograspnet + role: train + weight: 1.0 + frame_stride: 15 # grasp-hold is quasi-static; keep enumeration tractable + - name: dexh2r + adapter: dexh2r + root: ${MIGRATOR_DATA_ROOT}/dexh2r + role: train + weight: 1.0 + pseudo_gt: true # global-registration+ICP pseudo-GT (train-only); enumerates extracted subjects + - name: contactpose + adapter: contactpose + root: ${MIGRATOR_DATA_ROOT}/contactpose + role: train + weight: 0.5 # static grasps; only bowl present on disk so far + pseudo_gt: true + - name: housecat6d + adapter: housecat6d + root: ${MIGRATOR_DATA_ROOT}/housecat6d/extracted + role: heldout + data_split: test +model: + model_channels: 1024 + cond_channels: 1024 + num_blocks: 24 + num_heads: 16 + sigma_min: 0.0 + time_scale: 1000.0 + timestep_sampler: logit_normal + timestep_lognorm_mean: -1.0 + timestep_lognorm_std: 1.0 + pretrained_ckpt: ${MIGRATOR_CHECKPOINTS}/hf/ss_generator.ckpt + ss_decoder_ckpt: ${MIGRATOR_CHECKPOINTS}/hf/ss_decoder.ckpt +condition: + dino_model: dinov2_vitl14_reg + dino_input_size: 518 + pointmap_input_size: 256 + pointmap_patch_size: 8 + pointmap_embed_dim: 512 + projection_hidden_mult: 4.0 + dropout_prob: 0.1 + cam_hidden_dim: 256 + freeze: true + freeze_dino: true +loss_weights: + shape: 1.0 + 6drotation_normalized: 1.0 + translation: 1.0 + scale: 1.0 +optim: + lr: 0.0001 + weight_decay: 0.0 + grad_clip: 1.0 +train: + batch_size: 1 + steps: 100000 + num_workers: 8 + scheduler: cosine + warmup_steps: 1000 + min_lr: 1.0e-06 + resume_from: '' + log_every: 20 + ckpt_every: 2000 + ckpt_dir: ${MIGRATOR_CACHE}/ckpt_object_ssi + amp: true + amp_dtype: bf16 + seed: 0 + backend: ddp + deepspeed_config: '' + grad_accum_steps: 1 + val_every: 2000 + val_num_steps: 50 + val_max_batches: 16 diff --git a/migrator/code/mv-sam3d-for-6d-v2/configs/train_integrated_scene_ssi.yaml b/migrator/code/mv-sam3d-for-6d-v2/configs/train_integrated_scene_ssi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e0eefe3e0917a43aa434cfa1651cd515426dfcdb --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/configs/train_integrated_scene_ssi.yaml @@ -0,0 +1,128 @@ +data: + ssi_mode: scene # SSI ablation + dexycb_root: ${MIGRATOR_DATA_ROOT}/dexycb + canonical_cache: ${MIGRATOR_CACHE}/canonical + frame_stride: 5 + anchor_sweep: true + subjects: null + image_size: 518 + num_aux_views: 7 + random_views: true + min_aux_views: 1 + max_aux_views: 3 + depth_scale: 1000.0 + depth_near: 0.1 + depth_far: 3.0 + val_fraction: 0.02 + split_seed: 0 + heldout_manifest: ${MIGRATOR_MV_REPO}/configs/heldout.yaml + datasets: + - name: dexycb + adapter: dexycb + root: ${MIGRATOR_DATA_ROOT}/dexycb + role: train + weight: 1.0 + - name: pace + adapter: pace + root: ${MIGRATOR_DATA_ROOT}/pace + role: train + weight: 3.0 + data_split: test + - name: graspnet + adapter: graspnet + root: ${MIGRATOR_DATA_ROOT}/graspnet + role: train + weight: 2.0 + - name: ycbv + adapter: ycbv + root: ${MIGRATOR_DATA_ROOT}/ycbv + role: train + weight: 1.0 + data_split: test + # Adapters implemented + latent-complete + reprojection-gate >=0.8 (2026-07-17). + - name: hocap + adapter: hocap + root: ${MIGRATOR_DATA_ROOT}/hocap + role: train + weight: 1.0 + pseudo_gt: true # FoundationPose+SDF pseudo-GT (train-only; never eval) + - name: h2o + adapter: h2o + root: ${MIGRATOR_DATA_ROOT}/h2o + role: train + weight: 0.5 # only 8 objects + - name: hograspnet + adapter: hograspnet + root: ${MIGRATOR_DATA_ROOT}/hograspnet + role: train + weight: 1.0 + frame_stride: 15 # grasp-hold is quasi-static; keep enumeration tractable + - name: dexh2r + adapter: dexh2r + root: ${MIGRATOR_DATA_ROOT}/dexh2r + role: train + weight: 1.0 + pseudo_gt: true # global-registration+ICP pseudo-GT (train-only); enumerates extracted subjects + - name: contactpose + adapter: contactpose + root: ${MIGRATOR_DATA_ROOT}/contactpose + role: train + weight: 0.5 # static grasps; only bowl present on disk so far + pseudo_gt: true + - name: housecat6d + adapter: housecat6d + root: ${MIGRATOR_DATA_ROOT}/housecat6d/extracted + role: heldout + data_split: test +model: + model_channels: 1024 + cond_channels: 1024 + num_blocks: 24 + num_heads: 16 + sigma_min: 0.0 + time_scale: 1000.0 + timestep_sampler: logit_normal + timestep_lognorm_mean: -1.0 + timestep_lognorm_std: 1.0 + pretrained_ckpt: ${MIGRATOR_CHECKPOINTS}/hf/ss_generator.ckpt + ss_decoder_ckpt: ${MIGRATOR_CHECKPOINTS}/hf/ss_decoder.ckpt +condition: + dino_model: dinov2_vitl14_reg + dino_input_size: 518 + pointmap_input_size: 256 + pointmap_patch_size: 8 + pointmap_embed_dim: 512 + projection_hidden_mult: 4.0 + dropout_prob: 0.1 + cam_hidden_dim: 256 + freeze: true + freeze_dino: true +loss_weights: + shape: 1.0 + 6drotation_normalized: 1.0 + translation: 1.0 + scale: 1.0 +optim: + lr: 0.0001 + weight_decay: 0.0 + grad_clip: 1.0 +train: + batch_size: 1 + steps: 100000 + num_workers: 8 + scheduler: cosine + warmup_steps: 1000 + min_lr: 1.0e-06 + resume_from: '' + log_every: 20 + ckpt_every: 2000 + ckpt_dir: ${MIGRATOR_CACHE}/ckpt_scene_ssi + amp: true + amp_dtype: bf16 + seed: 0 + backend: ddp + deepspeed_config: '' + grad_accum_steps: 1 + val_every: 2000 + val_num_steps: 50 + val_max_batches: 16 diff --git a/migrator/code/mv-sam3d-for-6d-v2/docs/TRAINING_FEATURES.md b/migrator/code/mv-sam3d-for-6d-v2/docs/TRAINING_FEATURES.md new file mode 100644 index 0000000000000000000000000000000000000000..5a7b0f34e295e1ee36fe2777d85a411e21d47c7f --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/docs/TRAINING_FEATURES.md @@ -0,0 +1,122 @@ +# mv-sam3d-for-6d-v2 — Training Features + +> **Scope.** This will become the single doc that explains every training feature +> of the multi-view 6D-pose flow model. For now it documents **only the camera +> embedding (`mlp_cam`)** — the other features (joint flow objective, SSI pose +> targets, DDP/cosine/resume, held-out validation loss, modality dropout, …) will +> be added as further sections later. Everything below reflects the code as +> implemented (verified job 4526); paths are absolute. + +--- + +## 1. Camera embedding (`mlp_cam`) + +### One-line summary +Each view's tokens are tagged with a learned embedding **derived from that view's +camera pose** (relative aux→anchor extrinsics), instead of a content-free +per-slot index. This injects inter-camera geometry into the conditioning so the +appearance tokens (RGB/mask) know which camera they came from. + +### The gap it closes +The conditioning is multi-view: one **anchor** camera defines the reference frame, +the others are **aux** views. Before this feature, the only thing carrying +relative camera geometry into the conditions was the anchor-frame **pointmap** — +RGB and mask tokens carried no geometry, and the old per-view tag (`view_emb`) was +a learned `[max_views, C]` table indexed by **slot**, which was: +- **content-free** — it encoded only "which slot," not where the camera is; +- **order-arbitrary** — with random view sampling, "aux slot 1" is a *different + physical camera* every sample, so the table could never learn a stable + per-camera signal. + +`mlp_cam` replaces that table with a per-view bias computed from the actual +extrinsics, so the camera frame is bound to **every** token of the view — +including the otherwise geometry-free RGB/mask tokens — and is consistent across +the random aux subset. + +### Why this form (design decision) +Three options were considered for injecting camera geometry; we chose the first: + +1. **Additive per-view embedding (chosen).** `mlp_cam(extrinsics_v)` → one vector + added to all of view *v*'s tokens. It **binds** the camera pose to every + appearance token automatically (no ambiguity about which tokens belong to which + camera). Cheap, no token-count growth, condition-side only. +2. **Concatenated camera token (rejected).** A dedicated camera token per view in + the sequence has full capacity but *floats* — nothing links it to that view's + appearance tokens, so the model has to learn the association. +3. **Per-pixel ray positional encoding (rejected for this setting).** Plücker rays + through the embedder give per-token geometry, but it is **largely redundant** + with the anchor-frame pointmap (which already carries per-pixel 3D where depth + is valid) and cannot cleanly reach the frozen-DINO RGB/mask tokens. It would be + the right choice only in a depth-sparse / monocular regime. + +> Note: adopting `mlp_cam` deliberately **reverses** the earlier "no new modules" +> rule. It is now the *only* from-scratch module in the model. + +### What exactly is added +For each view, a **9-D camera feature** is built from the relative transform that +maps that camera's points into the anchor PyTorch3D frame +(`X_anchor = Rrel · X_aux + trel`): + +``` +cam_feat (9) = [ rot6d(Rrel) | (trel − shift) / scale ] + └── 6 ──┘ └──────── 3 ─────────┘ +``` + +- **`rot6d(Rrel)` (6)** — column-convention 6D of the aux→anchor rotation + (`matrix_to_rot6d`, the same convention as the object-rotation target). Raw 6D + columns are already unit-scaled, so **no** mean/std normalization is applied + here (that normalization is tuned to the *object* rotation distribution, not the + camera one). +- **`(trel − shift)/scale` (3)** — the aux camera **centre** in the anchor frame + (which is exactly `trel`, the image of the aux origin), expressed in the **same + SSI space as the conditioning pointmaps**. So the camera position lives in the + identical normalized coordinate system as the 3D points it observes. +- **Anchor view**: passes `Rrel = I`, `trel = 0`, computed by the *same* path — its + rot6d is the identity `[1,0,0,0,1,0]` and its centre normalizes to + `−shift/scale`. No special-casing. + +There is **no scale component in the transform itself** (camera↔camera is rigid +SE(3)); the SSI `scale`/`shift` appear only to normalize the camera *position* +into pointmap-consistent units — they are not a camera scale and not the object +scale `s`. The rotation never sees them. + +### Where it is computed and wired +| Stage | Location | What happens | +|---|---|---| +| build feature | `mvsam3d/data/dexycb_dataset.py` → `_cam_feature(Rrel, trel, shift, scale)` (`CAM_FEAT_DIM = 9`) | packs the 9-D vector; attached to every view dict as `view["cam"]` (anchor uses `I, 0`) | +| batch | `mvsam3d/data/dexycb_dataset.py` → `collate_multiview` | `cam` stacks per view position to `[B, 9]` with no special handling | +| consume | `mvsam3d/model/condition.py` → `MultiViewConditionEmbedder.forward` | per view: `tok = tok + self.mlp_cam(view["cam"])[:, None, :]` (broadcast over all tokens of the view) | + +### The module +- `self.mlp_cam = Sequential(Linear(9, cam_hidden_dim), SiLU, Linear(cam_hidden_dim, model_channels))`. +- **Output layer is zero-initialized** (weight & bias). At step 0 the per-view bias + is exactly **0**, so the conditioning equals the no-tag baseline and the + warm-started fuser/backbone are perturbed only as `mlp_cam` learns — i.e. + resuming/continuing from a prior run sees no init shock. +- **Trainable.** The whole `EmbedderFuser` (both DINOs, PointPatchEmbed, + projection nets, `idx_emb`) stays **frozen** after warm-start; `mlp_cam` is the + only trainable condition-side module (and the only from-scratch module in the + whole model). + +### Config +- `condition.cam_hidden_dim` (default `256`) — hidden width of the 2-layer MLP. +- `model.max_views` is now **API-only**: `mlp_cam` derives the per-view bias from + the extrinsics, so there is no per-slot table to size and **no upper bound on the + view count**. + +### Properties / invariants +- **Camera-consistent under random views**: the same physical camera yields the + same embedding regardless of which random slot it lands in. +- **Unbounded view count**: any number of cameras works (no table to overflow). +- **Warm-start neutral at step 0**: zero-init output ⇒ identical to the no-camera + baseline initially. +- **Intrinsics not included** (extrinsics only). Appending normalized + focal/principal-point would extend the feature to 13-D; deferred as minor for a + fixed-rig, pointmap-present DexYCB setup. + +### Verification +Job 4526 (B200): warm-start intact (backbone 936/936, fuser 722/722 frozen), +`condition trainable modules == ['mlp_cam']` (265,728 params), grad audit +`fuser=0 / mlp_cam=6.5e3 / backbone=1.75e5`, finite loss, train + validation +smoke clean. CPU smoke confirmed anchor `rot6d=[1,0,0,0,1,0]`, +`pos_ssi=[0,0,−median_z/scale]`, and zero-init forward = exactly 0. diff --git a/migrator/code/mv-sam3d-for-6d-v2/docs/dataset_convention_spec.md b/migrator/code/mv-sam3d-for-6d-v2/docs/dataset_convention_spec.md new file mode 100644 index 0000000000000000000000000000000000000000..e6f69f2ba323ced6a8d4162cdaed1b67e929d224 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/docs/dataset_convention_spec.md @@ -0,0 +1,44 @@ +# Dataset Coordinate-Convention Spec (verified 2026-07-16/17) + +TARGET (what every adapter must output, per view, per object): +- pose (R_oc, t_oc): `x_cam = R_oc @ p_cad + t_oc`, OpenCV axes (x-right, y-down, z-forward), t METERS, p_cad = mesh verts (m) in mesh frame. +- K OpenCV pinhole (fx,fy,cx,cy); depth meters registered to RGB; extrinsics world<->cam OpenCV meters; mesh meters in pose frame. +The base dataloader then applies D=diag(-1,-1,1) (OpenCV->PyTorch3D) and canonical M. Adapters MUST NOT pre-apply D or M. + +Legend: ✅=empirical reprojection IoU, 🔵=code/cross-check verified. + +## YCB-V 🔵 (train; YCB objs) — LOCAL(models); need BOP train_real +R_oc=cam_R_m2c.reshape(3,3); t_oc=cam_t_m2c*1e-3; mesh=models_eval/obj_{id:06d}.ply /1000. depth=raw*depth_scale/1000. K=cam_K. extr cam_R/t_w2c (mm->m) world->cam. No flip (native OpenCV). Masks mask_visib/. + +## HO-3D 🔵 (HELD-OUT, never train; YCB objs) — LOCAL +R=Rodrigues(meta['objRot']); t=meta['objTrans'](m); then R=diag(1,-1,-1)@R, t=diag(1,-1,-1)@t (OpenGL->OpenCV, applied to BOTH). mesh dexycb textured_simple.obj (m). depth=(R_ch+G_ch*256)*1.2499e-4. K=meta['camMat']. Eval-split masks rendered. + +## HOPE-Video 🔵 (HELD-OUT bonus OOD; non-YCB) — LOCAL +R=P[:3,:3]; t=P[:3,3]*1e-2 (cm->m). mesh meshes/eval/{class}.obj /100. depth raw*1e-3. extrinsics raw-as-meters world->cam (do NOT scale). K=camera.intrinsics. Masks precomputed masks_visib/. + +## YCB-M 🔵 (HELD-OUT; YCB objs) — LOCAL +R=quat_xyzw->R (F=I, no flip); t=location*1e-2 (cm->m), object in CAMERA frame. mesh models_aligned/ycb_models_aligned_cm /100 (NOT dexycb). depth raw/10000. camera_data=world->cam (invert->cam->world). ⚠️ rotation handedness taken as-is, calib validated translation only — render-verify rotation before any training use. + +## HouseCat6D ✅ IoU 0.854 (HELD-OUT #1 benchmark; 194 non-YCB) — LOCAL(sqsh) +R_oc,t_oc = label pkl rotations[i],translations[i] — ALREADY obj->cam OpenCV meters (identity conversion). mesh obj_models_small_size_final//.obj (m, origin-centered; cat=token before '-'). depth /1000. K=intrinsics.txt. extr camera_pose/NNNNNN.txt is CAM->WORLD (invert for world->cam); train/val have NO extrinsics (use pkl directly). ⚠️ masks: bg=255, use `mask!=255` (NOT >0). Index by model_list not class_ids. Do NOT rescale by scales/gt_scales. + +## GraspNet-1B ✅ IoU 0.981 (train + test_novel HELD-OUT; ~56 non-YCB of 88) — SAMPLE(scene_0017) +R_oc=meta['poses'][:3,:3,i]; t_oc=meta['poses'][:3,3,i] (m). mesh models/%03d/nontextured.ply (m). ⚠️ OBJECT-ID OFF-BY-ONE: label/cls_indexes value V -> mesh models/(V-1). depth raw/1000. K=camK.npy (Kinect≠RealSense, read per-scene). extr per-view cam->table=cam0_wrt_table@camera_poses[v] (not needed; per-view poses already in cam frame). CC-BY-NC. + +## PACE ✅ IoU 0.90-0.97 (train; 238 non-YCB) — SAMPLE(scene 000215) +BOP: R_oc=cam_R_m2c.reshape(3,3); t_oc=cam_t_m2c/1000; mesh models_eval/obj_XXXXXX.ply /1000. depth raw*depth_scale/1000 (real depth_scale=1.0; PBR /10000). K=cam_K. ⚠️ real split has NO cam_R/t_w2c (3-cam rig = separate scene folders; per-view poses self-contained). ⚠️ FILTER ARTICULATED: drop obj_id 545..692 (AKB-48); keep 0..544 rigid. masks {imgid:06d}_{gtidx:06d}.png (gtidx into scene_gt). + +## HO-Cap ✅ IoU 0.77-0.87 (TRAIN-ONLY pseudo-GT; 64 non-YCB) — SAMPLE(subject_1) +poses_o.npy [n_obj,n_frame,7]=[qx,qy,qz,qw,tx,ty,tz] scipy xyzw, obj->WORLD (world=tag_1). Per-cam: T_oc(s)=rs_RTs_inv[s]@quat_to_mat(poses_o[o,f]); OR read label_.npz['obj_poses'] (already per-cam obj->cam, verified 1e-7). rs_RTs[s]=inv(tag_1)@extr2master[s]; rs_RTs_inv=inv. mesh models//textured_mesh.obj (m, matches poses). depth /1000 with COLOR K. K=color intrinsics. ⚠️ seg_mask ids LOCAL 1..4 objects/5 rh/6 lh. poses/labels in separate zips. + +## DexH2R 🔵 round-trip 2.4e-16 (TRAIN-ONLY pseudo-GT, 6 of 18 views; 56 non-YCB) — MESHES+CALIB only (frames quota-blocked) +obj_pose.pt [T,4,4]=T_world<-obj (world=hand_arm base). Kinect i: T_oc=inv(cali0i)@hand_arm_mesh_to_kinect_pcd_0@obj_pose[t]. RealSense i(0,1): time-varying via ShadowHand FK(qpos.pt)@inv(realsense_to_forearm_i). mesh object_model/.obj (m). depth int16->uint16 /1000. K per-cam (Kinect 8-param distort; images undistorted). ⚠️ ONLY kinect0-3+realsense0-1 have depth (drop 12 ZCAM). No 2D masks (project real_obj_pcd or threshold depth). Use per-trial calibration/. + +## H2O 🔵 code-verified (train; 8 non-YCB) — NO DATA (registration wall) +obj_pose_RT/{f}.txt = class_id + 4x4 obj->cam OpenCV meters, USED AS-IS (⚠️ NO inversion — the "cam_to_obj" note refers to a DIFFERENT file, the keypoint obj_pose/). mesh object/{class}.ply (m). depth /1000. cam_pose/ = cam->world (invert for world->cam). K from cam_intrinsics.txt (fx fy cx cy w h); images undistorted. Per-camera RT already in that cam's frame. Masks must be rendered. Mesh folder names ≠ display names. + +## HOGraspNet 🔵 code-verified (train; 22 YCB+8, use scanned meshes; mocap GT) — MESHES only (annotations gated) +object_mat = 4x4 obj->WORLD (world=mas cam). extrinsic = (3,4) world->cam. T_oc=M4(extrinsic)@object_mat. R_oc=T_oc[:3,:3]; t_oc=T_oc[:3,3]/100 (cm->m). mesh obj_scanned_models/.obj: p_cad=(S*V_raw)/100 where S=_OBJECT_SCALE_FIXED[idx-1] (⚠️ MANDATORY per-obj scale; wine_glass .830, small_marker .1035, spatula .671, golf_ball .430, else 1.0). depth /1000. K OpenCV (full-res 1920x1080; cropped 640x480 -> shift cx,cy by bbox). ⚠️ use scanned meshes NOT dexycb CAD (different frame). obj idx = HOGraspNet 1..30 ordering. masks pseudo. + +## ContactPose ✅ IoU 0.979 (TRAIN-ONLY static; 25 non-YCB) — SAMPLE(p28/use/bowl) +cTo=object_pose(cam,frame)=obj->cam OpenCV meters (R_oc=cTo[:3,:3],t_oc=cTo[:3,3]). mesh object_models/.ply *1e-3 (PLY is mm). depth /1000 (Kinect v2, ~30mm systematic bias). ⚠️ projection needs affine A: pixels = A@K@x_cam (cp.P); left/right cams are PORTRAIT 540x960 while K is 960x540 landscape. extr wTc=cam->world (invert->cTw). ⚠️ static: sample 1-few frames per (subject,obj,intent) — 2306 grasps not 2.9M frames. Masks render from mesh+cTo. quats w-first if parsing raw. diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..edee3c1379e38c7ecfd53680b746165816760107 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/__init__.py @@ -0,0 +1,9 @@ +"""mvsam3d: multi-view 6D-pose flow model (migrator package). + +Importing the package registers the migrator path roots (MIGRATOR_ROOT, +MIGRATOR_DATA_ROOT, MIGRATOR_CHECKPOINTS, MIGRATOR_CACHE, MIGRATOR_MV_REPO) into +os.environ so that both YAML configs and helper scripts can reference paths via +``${MIGRATOR_*}`` placeholders regardless of where the package is installed. +User-set values are never overridden (setdefault).""" +from . import paths as paths # noqa: F401 +paths.register_env() diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cece9daaa7190b8cd18999ef5ce76adf8c1ba676 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/__pycache__/paths.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/__pycache__/paths.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f0f6065b096a42889013b1493c7336c71d16a7eb Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/__pycache__/paths.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..13d4f0b26030a21600d1e29ff0f2cac7120e9af5 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/base_dataset.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/base_dataset.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f4dc4c0056f3ff3a28518d9b06d06375048fed4 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/base_dataset.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/canonical.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/canonical.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..69ad91e26cc684c07010f25cf3003a65cae9b620 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/canonical.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/dexycb_calib.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/dexycb_calib.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..95c0902a2a40292cf24a878091018b2bef7bd2d5 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/dexycb_calib.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/dexycb_dataset.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/dexycb_dataset.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97d0a4f7d3d67ff2027ff532deb5fe34f457f749 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/dexycb_dataset.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/integrated.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/integrated.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9746874dcc89c7ee6b97fadbe9549764f352d000 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/__pycache__/integrated.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/adapters.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/adapters.py new file mode 100644 index 0000000000000000000000000000000000000000..edb74be024c797c6b1e04a6d78a8a5bd07426ebb --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/adapters.py @@ -0,0 +1,1402 @@ +"""Per-dataset adapters for MultiViewPoseDataset. + +Each adapter converts a dataset's native storage into the UNIFIED contract +(OpenCV obj->cam meters, depth-m registered to RGB, world/tag->cam extrinsics). +Conventions are the verified ones in docs/dataset_convention_spec.md. + +Design note on extrinsics for the two multi-view regimes: + (a) synced rig (DexYCB): extrinsic_cv returns the constant per-serial rig + calibration (tag->cam). + (b) static scene + moving cam (HouseCat6D, GraspNet, YCB-V, HOPE): the SAMPLE + fixes one object; we take the OBJECT'S canonical frame as the shared + "world", so extrinsic_cv returns that object's obj->cam pose at the given + frame. The object is static, so relative_p3d_from_ext(anchor,aux) yields the + exact relative camera transform with NO dependency on a camera-pose file + (works for train/val/test uniformly). +""" +from __future__ import annotations + +import glob +import os +from typing import Any, Dict, List, Tuple + +import numpy as np +import torch +from PIL import Image + +from .base_dataset import PoseDatasetAdapter, ViewSample +from . import dexycb_calib +from .dexycb_dataset import YCB_ID_TO_NAME, DEPTH_SCALE + + +# --------------------------------------------------------------------------- # +# canonical-latent cache access (shared) +# --------------------------------------------------------------------------- # +class _CanonCache: + def __init__(self, cache_dir: str): + self.cache_dir = cache_dir + self._c: Dict[str, dict] = {} + + def get(self, key: str) -> dict: + if key not in self._c: + # key may be namespaced "/"; file is /.pt + path = os.path.join(self.cache_dir, f"{key}.pt") + if not os.path.exists(path): + ds = key.split("/", 1)[0] if "/" in key else "" + raise FileNotFoundError( + f"canonical latent not built for key '{key}' (expected {path}). " + f"Run: python tools/build_canonical_latents_multi.py --dataset {ds} " + f"--root before training on this dataset." + ) + self._c[key] = torch.load(path, map_location="cpu", weights_only=False) + return self._c[key] + + +def _bop_gt_index(entries, oid, rank=0): + """Positional index of the rank-th entry with obj_id==oid in a BOP scene_gt + frame list; None if that instance is absent. Used so aux views resolve the + SAME physical object by (obj_id, occurrence-rank) instead of trusting the + anchor's positional index across frames (which breaks if scene_gt ordering + differs). Mirrors GraspNet's index-by-class approach, rank-aware for scenes + with duplicate instances of one obj_id.""" + cnt = 0 + for idx, o in enumerate(entries): + if int(o["obj_id"]) == oid: + if cnt == rank: + return idx + cnt += 1 + return None + + +# --------------------------------------------------------------------------- # +# DexYCB (reference adapter — reproduces MultiViewDexYCB exactly) +# --------------------------------------------------------------------------- # +class DexYCBAdapter(PoseDatasetAdapter): + name = "dexycb" + + def __init__(self, root: str, canonical_cache: str, *, frame_stride: int = 5, + anchor_sweep: bool = True, subjects=None, depth_scale: float = DEPTH_SCALE): + self.root = root + self.calib_root = os.path.join(root, "calibration") + self.canon = _CanonCache(canonical_cache) + self.frame_stride = frame_stride + self.anchor_sweep = anchor_sweep + self.subjects = subjects + self.depth_scale = depth_scale + self._ext: Dict[str, dict] = {} + + def _extrinsics(self, tag): + if tag not in self._ext: + self._ext[tag] = dexycb_calib.load_extrinsics(self.calib_root, tag) + return self._ext[tag] + + def list_samples(self) -> List[ViewSample]: + seq_dirs = sorted(glob.glob(os.path.join(self.root, "2020*-subject-*", "2020*"))) + if self.subjects is not None: + seq_dirs = [d for d in seq_dirs if any(s in d for s in self.subjects)] + out: List[ViewSample] = [] + for seq_dir in seq_dirs: + meta = dexycb_calib._load_yaml(os.path.join(seq_dir, "meta.yml")) + serials = [str(s) for s in meta["serials"]] + ycb_ids = list(meta["ycb_ids"]) + tag = str(meta["extrinsics"]) + n = int(meta["num_frames"]) + anchors = serials if self.anchor_sweep else serials[:1] + for frame in range(0, n, self.frame_stride): + for anchor in anchors: + aux = [s for s in serials if s != anchor] + for oi, oid in enumerate(ycb_ids): + if int(oid) not in YCB_ID_TO_NAME: + continue + obj_key = YCB_ID_TO_NAME[int(oid)] + aref = (seq_dir, anchor, frame, int(oid), oi, tag) + pool = [(seq_dir, s, frame, int(oid), oi, tag) for s in aux] + out.append(ViewSample(anchor=aref, aux_pool=pool, obj_key=obj_key)) + return out + + def load_view(self, ref: Any) -> Dict[str, Any]: + seq_dir, serial, frame, obj_id, obj_index, tag = ref + rgb = np.asarray(Image.open(os.path.join(seq_dir, serial, f"color_{frame:06d}.jpg")) + .convert("RGB"), dtype=np.float32) / 255.0 + seg = np.load(os.path.join(seq_dir, serial, f"labels_{frame:06d}.npz"))["seg"] + mask = (seg == obj_id).astype(np.float32) + depth = np.asarray(Image.open(os.path.join( + seq_dir, serial, f"aligned_depth_to_color_{frame:06d}.png")), + dtype=np.float32) / self.depth_scale + intr = dexycb_calib.load_intrinsics(self.calib_root, serial) + return {"rgb": rgb, "mask": mask, "depth_m": depth, "intr": intr} + + def extrinsic_cv(self, ref: Any) -> Dict[str, np.ndarray]: + seq_dir, serial, frame, obj_id, obj_index, tag = ref + e = self._extrinsics(tag)[serial] + return {"R": e["R"], "t": e["t"]} + + def object_pose_cv(self, sample: ViewSample) -> Tuple[np.ndarray, np.ndarray]: + seq_dir, serial, frame, obj_id, obj_index, tag = sample.anchor + py = np.load(os.path.join(seq_dir, serial, f"labels_{frame:06d}.npz"))["pose_y"][obj_index] + return py[:, :3].astype(np.float64), py[:, 3].astype(np.float64) + + def mesh_for(self, sample: ViewSample): + return _load_mesh(os.path.join(self.root, "models", sample.obj_key, "textured_simple.obj")) + + def canonical(self, obj_key: str) -> dict: + return self.canon.get(obj_key) + + +# --------------------------------------------------------------------------- # +# shared helpers for the new adapters +# --------------------------------------------------------------------------- # +import json as _json +import trimesh as _trimesh + +def _K_to_intr(K): + return {"fx": float(K[0, 0]), "fy": float(K[1, 1]), + "ppx": float(K[0, 2]), "ppy": float(K[1, 2])} + +def _load_mesh(path): + m = _trimesh.load(path, process=False, force="mesh") + if isinstance(m, _trimesh.Scene): + m = _trimesh.util.concatenate([g for g in m.geometry.values()]) + return m + + +# --------------------------------------------------------------------------- # +# HouseCat6D (type b, HELD-OUT; obj->cam OpenCV meters == identity conversion) +# --------------------------------------------------------------------------- # +class HouseCat6DAdapter(PoseDatasetAdapter): + name = "housecat6d" + + def __init__(self, root, canonical_cache, *, split="test", max_aux_frames=8, frame_stride=10): + # root = .../housecat6d/extracted ; split dir under it (test/train/val) + self.split_dir = os.path.join(root, split) + self.canon = _CanonCache(canonical_cache) + self.max_aux_frames = max_aux_frames + self.frame_stride = frame_stride + self._pkl_cache = {} + + def _pkl(self, scene, frame): + key = (scene, frame) + if key not in self._pkl_cache: + import pickle + with open(os.path.join(scene, "labels", f"{frame:06d}_label.pkl"), "rb") as f: + self._pkl_cache[key] = pickle.load(f) + return self._pkl_cache[key] + + def _frames(self, scene): + return sorted(int(os.path.splitext(f)[0]) for f in os.listdir(os.path.join(scene, "rgb"))) + + def list_samples(self): + out = [] + scenes = sorted(glob.glob(os.path.join(self.split_dir, "*"))) + for scene in scenes: + if not os.path.isdir(os.path.join(scene, "rgb")): + continue + frames = self._frames(scene) + if len(frames) < 2: + continue + models = self._pkl(scene, frames[0])["model_list"] + for oi, model in enumerate(models): + anchors = frames[:: self.frame_stride] or [frames[0]] + for af in anchors: + pool = [f for f in frames if f != af] + out.append(ViewSample( + anchor={"scene": scene, "frame": af, "oi": oi, "model": model}, + aux_pool=[{"scene": scene, "frame": f, "oi": oi, "model": model} for f in pool], + obj_key=f"housecat6d/{model}", + meta={"scene": scene, "model": model})) + return out + + def load_view(self, ref): + scene, frame, model = ref["scene"], ref["frame"], ref["model"] + rgb = np.asarray(Image.open(os.path.join(scene, "rgb", f"{frame:06d}.png")).convert("RGB"), + np.float32) / 255.0 + inst = np.asarray(Image.open(os.path.join(scene, "instance", f"{frame:06d}_{model}.png"))) + mask = (inst != 255).astype(np.float32) # bg=255 + depth = np.asarray(Image.open(os.path.join(scene, "depth", f"{frame:06d}.png")), np.float32) / 1000.0 + K = np.loadtxt(os.path.join(scene, "intrinsics.txt")).reshape(3, 3) + return {"rgb": rgb, "mask": mask, "depth_m": depth, "intr": _K_to_intr(K)} + + def extrinsic_cv(self, ref): # world := object frame -> obj->cam pose + d = self._pkl(ref["scene"], ref["frame"]) + return {"R": np.asarray(d["rotations"][ref["oi"]], np.float64), + "t": np.asarray(d["translations"][ref["oi"]], np.float64)} + + def object_pose_cv(self, sample): + e = self.extrinsic_cv(sample.anchor) + return e["R"], e["t"] + + def mesh_for(self, sample): + model = sample.meta["model"]; cat = model.split("-")[0] + base = os.path.dirname(os.path.dirname(self.split_dir)) # .../housecat6d + return _load_mesh(os.path.join(base, "extracted", "obj_models_small_size_final", + cat, f"{model}.obj")) + + def canonical(self, obj_key): + return self.canon.get(obj_key) + + +# --------------------------------------------------------------------------- # +# PACE (BOP; real split; obj->cam OpenCV mm). rigid-only (drop obj_id 545..692) +# NOTE: real multi-view = the 3 rig cams (separate scene folders per scene/video); +# a single-cam scene validates convention but not cross-cam grouping. +# --------------------------------------------------------------------------- # +class PACEAdapter(PoseDatasetAdapter): + name = "pace" + + def __init__(self, root, canonical_cache, *, split="test", depth_scale=1.0, frame_stride=10, max_aux_frames=8): + self.split_dir = os.path.join(root, split) + self.models_dir = os.path.join(root, "models_eval") + self.canon = _CanonCache(canonical_cache) + self.depth_scale = depth_scale + self.frame_stride = frame_stride + self._gt = {}; self._cam = {} + + def _scene_json(self, scene, which): + cache = self._gt if which == "gt" else self._cam + if scene not in cache: + cache[scene] = _json.load(open(os.path.join(scene, f"scene_{'gt' if which=='gt' else 'camera'}.json"))) + return cache[scene] + + def list_samples(self): + out = [] + for scene in sorted(glob.glob(os.path.join(self.split_dir, "*"))): + if not os.path.isdir(os.path.join(scene, "rgb")): + continue + gt = self._scene_json(scene, "gt") + imids = sorted(int(k) for k in gt.keys()) + for imid in imids[:: self.frame_stride] or imids[:1]: + entries = gt[str(imid)] + for gi, o in enumerate(entries): + oid = int(o["obj_id"]) + if 545 <= oid <= 692: # articulated AKB-48 -> drop + continue + # occurrence-rank of this instance within its frame + rank = sum(1 for o2 in entries[:gi] if int(o2["obj_id"]) == oid) + # aux = only frames that actually contain this (oid,rank) instance + pool = [{"scene": scene, "imid": j, "oid": oid, "rank": rank} + for j in imids if j != imid + and _bop_gt_index(gt[str(j)], oid, rank) is not None] + ref = {"scene": scene, "imid": imid, "oid": oid, "rank": rank} + out.append(ViewSample(anchor=ref, aux_pool=pool, obj_key=f"pace/{oid:06d}", + meta={"oid": oid})) + return out + + def _gi(self, ref): + idx = _bop_gt_index(self._scene_json(ref["scene"], "gt")[str(ref["imid"])], + ref["oid"], ref.get("rank", 0)) + if idx is None: + raise KeyError(f"pace: obj {ref['oid']} rank {ref.get('rank',0)} not in " + f"{ref['scene']} frame {ref['imid']}") + return idx + + def load_view(self, ref): + scene, imid = ref["scene"], ref["imid"] + rgb = np.asarray(Image.open(os.path.join(scene, "rgb", f"{imid:06d}.png")).convert("RGB"), np.float32) / 255.0 + depth = np.asarray(Image.open(os.path.join(scene, "depth", f"{imid:06d}.png")), np.float32) + cam = self._scene_json(scene, "camera")[str(imid)] + ds = float(cam.get("depth_scale", self.depth_scale)) + depth_m = depth * ds / 1000.0 + gi = self._gi(ref) + mp = os.path.join(scene, "mask_visib", f"{imid:06d}_{gi:06d}.png") + mask = (np.asarray(Image.open(mp)) > 0).astype(np.float32) if os.path.exists(mp) else np.zeros(depth.shape, np.float32) + K = np.asarray(cam["cam_K"], np.float64).reshape(3, 3) + return {"rgb": rgb, "mask": mask, "depth_m": depth_m, "intr": _K_to_intr(K)} + + def _pose(self, ref): + o = self._scene_json(ref["scene"], "gt")[str(ref["imid"])][self._gi(ref)] + R = np.asarray(o["cam_R_m2c"], np.float64).reshape(3, 3) + t = np.asarray(o["cam_t_m2c"], np.float64) / 1000.0 + return R, t + + def extrinsic_cv(self, ref): + R, t = self._pose(ref); return {"R": R, "t": t} + + def object_pose_cv(self, sample): + return self._pose(sample.anchor) + + def mesh_for(self, sample): + m = _load_mesh(os.path.join(self.models_dir, f"obj_{sample.meta['oid']:06d}.ply")) + m.apply_scale(1e-3) # mm -> m + return m + + def canonical(self, obj_key): + return self.canon.get(obj_key) + + +# --------------------------------------------------------------------------- # +# GraspNet-1B (type b; obj->cam OpenCV meters; mesh id = cls_value - 1) +# --------------------------------------------------------------------------- # +class GraspNetAdapter(PoseDatasetAdapter): + name = "graspnet" + + def __init__(self, root, canonical_cache, *, camera="kinect", frame_stride=20, max_aux_frames=8): + self.scenes_dir = os.path.join(root, "scenes") + self.models_dir = os.path.join(root, "models") + self.canon = _CanonCache(canonical_cache) + self.camera = camera + self.frame_stride = frame_stride + self._meta = {} + + def _load_meta(self, cdir, v): + key = (cdir, v) + if key not in self._meta: + import scipy.io as sio + self._meta[key] = sio.loadmat(os.path.join(cdir, "meta", f"{v:04d}.mat")) + return self._meta[key] + + def list_samples(self): + out = [] + for scene in sorted(glob.glob(os.path.join(self.scenes_dir, "scene_*"))): + cdir = os.path.join(scene, self.camera) + if not os.path.isdir(os.path.join(cdir, "rgb")): + continue + views = sorted(int(os.path.splitext(f)[0]) for f in os.listdir(os.path.join(cdir, "rgb"))) + m0 = self._load_meta(cdir, views[0]) + cls = [int(c) for c in m0["cls_indexes"].ravel()] + anchors = views[:: self.frame_stride] or [views[0]] + for av in anchors: + mv = self._load_meta(cdir, av) + clsv = [int(c) for c in mv["cls_indexes"].ravel()] + for i, V in enumerate(clsv): + pool = [{"cdir": cdir, "v": w, "cls": V} for w in views if w != av] + out.append(ViewSample( + anchor={"cdir": cdir, "v": av, "cls": V, "i": i}, + aux_pool=pool, obj_key=f"graspnet/{V-1:03d}", meta={"V": V})) + return out + + def load_view(self, ref): + cdir, v = ref["cdir"], ref["v"] + rgb = np.asarray(Image.open(os.path.join(cdir, "rgb", f"{v:04d}.png")).convert("RGB"), np.float32) / 255.0 + depth = np.asarray(Image.open(os.path.join(cdir, "depth", f"{v:04d}.png")), np.float32) / 1000.0 + label = np.asarray(Image.open(os.path.join(cdir, "label", f"{v:04d}.png"))) + mask = (label == ref["cls"]).astype(np.float32) + m = self._load_meta(cdir, v) + K = np.asarray(m["intrinsic_matrix"], np.float64).reshape(3, 3) + return {"rgb": rgb, "mask": mask, "depth_m": depth, "intr": _K_to_intr(K)} + + def _pose(self, cdir, v, cls): + m = self._load_meta(cdir, v) + clsv = [int(c) for c in m["cls_indexes"].ravel()] + i = clsv.index(cls) + P = m["poses"][:, :, i] + return P[:, :3].astype(np.float64), P[:, 3].astype(np.float64) + + def extrinsic_cv(self, ref): + R, t = self._pose(ref["cdir"], ref["v"], ref["cls"]); return {"R": R, "t": t} + + def object_pose_cv(self, sample): + a = sample.anchor + return self._pose(a["cdir"], a["v"], a["cls"]) + + def mesh_for(self, sample): + return _load_mesh(os.path.join(self.models_dir, f"{sample.meta['V']-1:03d}", "nontextured.ply")) + + def canonical(self, obj_key): + return self.canon.get(obj_key) + + +# --------------------------------------------------------------------------- # +# YCB-V (BOP; obj->cam OpenCV mm; frames-as-views static; YCB objs reuse latents) +# --------------------------------------------------------------------------- # +class YCBVAdapter(PoseDatasetAdapter): + name = "ycbv" + _BOP2NAME = None + + def __init__(self, root, canonical_cache, *, split="test", frame_stride=15, max_aux_frames=8): + self.split_dir = os.path.join(root, split) + self.models_dir = os.path.join(root, "models_eval") + self.canon = _CanonCache(canonical_cache) + self.frame_stride = frame_stride + self._gt = {}; self._cam = {} + if YCBVAdapter._BOP2NAME is None: + YCBVAdapter._BOP2NAME = {i + 1: n for i, n in enumerate([ + "002_master_chef_can","003_cracker_box","004_sugar_box","005_tomato_soup_can", + "006_mustard_bottle","007_tuna_fish_can","008_pudding_box","009_gelatin_box", + "010_potted_meat_can","011_banana","019_pitcher_base","021_bleach_cleanser", + "024_bowl","025_mug","035_power_drill","036_wood_block","037_scissors", + "040_large_marker","051_large_clamp","052_extra_large_clamp","061_foam_brick"])} + + def _sj(self, scene, which): + c = self._gt if which == "gt" else self._cam + if scene not in c: + c[scene] = _json.load(open(os.path.join(scene, f"scene_{'gt' if which=='gt' else 'camera'}.json"))) + return c[scene] + + def list_samples(self): + out = [] + for scene in sorted(glob.glob(os.path.join(self.split_dir, "*"))): + if not os.path.isdir(os.path.join(scene, "rgb")): + continue + gt = self._sj(scene, "gt") + imids = sorted(int(k) for k in gt.keys()) + for imid in imids[:: self.frame_stride] or imids[:1]: + entries = gt[str(imid)] + for gi, o in enumerate(entries): + oid = int(o["obj_id"]) + name = self._BOP2NAME.get(oid) + if name is None: + continue + rank = sum(1 for o2 in entries[:gi] if int(o2["obj_id"]) == oid) + pool = [{"scene": scene, "imid": j, "oid": oid, "rank": rank} + for j in imids if j != imid + and _bop_gt_index(gt[str(j)], oid, rank) is not None] + out.append(ViewSample( + anchor={"scene": scene, "imid": imid, "oid": oid, "rank": rank}, + aux_pool=pool, obj_key=name, meta={"oid": oid})) + return out + + def _gi(self, ref): + idx = _bop_gt_index(self._sj(ref["scene"], "gt")[str(ref["imid"])], + ref["oid"], ref.get("rank", 0)) + if idx is None: + raise KeyError(f"ycbv: obj {ref['oid']} rank {ref.get('rank',0)} not in " + f"{ref['scene']} frame {ref['imid']}") + return idx + + def load_view(self, ref): + scene, imid = ref["scene"], ref["imid"] + rgb = np.asarray(Image.open(os.path.join(scene, "rgb", f"{imid:06d}.png")).convert("RGB"), np.float32) / 255.0 + cam = self._sj(scene, "camera")[str(imid)] + ds = float(cam.get("depth_scale", 1.0)) + depth = np.asarray(Image.open(os.path.join(scene, "depth", f"{imid:06d}.png")), np.float32) * ds / 1000.0 + gi = self._gi(ref) + mp = os.path.join(scene, "mask_visib", f"{imid:06d}_{gi:06d}.png") + mask = (np.asarray(Image.open(mp)) > 0).astype(np.float32) if os.path.exists(mp) else np.zeros(depth.shape, np.float32) + K = np.asarray(cam["cam_K"], np.float64).reshape(3, 3) + return {"rgb": rgb, "mask": mask, "depth_m": depth, "intr": _K_to_intr(K)} + + def _pose(self, ref): + o = self._sj(ref["scene"], "gt")[str(ref["imid"])][self._gi(ref)] + return (np.asarray(o["cam_R_m2c"], np.float64).reshape(3, 3), + np.asarray(o["cam_t_m2c"], np.float64) / 1000.0) + + def extrinsic_cv(self, ref): + R, t = self._pose(ref); return {"R": R, "t": t} + + def object_pose_cv(self, sample): + return self._pose(sample.anchor) + + def mesh_for(self, sample): + m = _load_mesh(os.path.join(self.models_dir, f"obj_{sample.meta['oid']:06d}.ply")) + m.apply_scale(1e-3) + return m + + def canonical(self, obj_key): + return self.canon.get(obj_key) + + +# --------------------------------------------------------------------------- # +# STUBS / notes for adapters whose frame data is blocked or needs extra handling +# --------------------------------------------------------------------------- # +class _NotReadyAdapter(PoseDatasetAdapter): + """Placeholder: convention verified, adapter deferred. See reason.""" + def __init__(self, *a, reason="", **k): self._reason = reason + def list_samples(self): raise NotImplementedError(self._reason) + def load_view(self, ref): raise NotImplementedError(self._reason) + def extrinsic_cv(self, ref): raise NotImplementedError(self._reason) + def object_pose_cv(self, sample): raise NotImplementedError(self._reason) + def canonical(self, obj_key): raise NotImplementedError(self._reason) + +# --------------------------------------------------------------------------- # +# HO-Cap (type a, 8 synced RealSense; TRAIN-ONLY pseudo-GT). world := object +# frame (same pattern as the type-b adapters): each view's obj->cam comes from +# that camera's per-frame label npz (verified == official ope_gt to 1e-7), so no +# calibration-file compose is needed and relative_p3d_from_ext(anchor,aux) yields +# the exact cross-camera transform. object_ids[oi] <-> obj_poses[oi] <-> seg_mask +# id (oi+1), confirmed on disk. +# --------------------------------------------------------------------------- # +class HOCapAdapter(PoseDatasetAdapter): + name = "hocap" + + def __init__(self, root, canonical_cache, *, frame_stride=15, subjects=None): + self.root = root # dataset/hocap + self.canon = _CanonCache(canonical_cache) + self.frame_stride = frame_stride + self.subjects = subjects + self._lab = {} + + def _seqs(self): + out = [] + for subj in sorted(glob.glob(os.path.join(self.root, "subject_*"))): + if os.path.isfile(subj): + continue + if self.subjects and os.path.basename(subj) not in self.subjects: + continue + for seq in sorted(glob.glob(os.path.join(subj, "2*"))): + if os.path.exists(os.path.join(seq, "meta.yaml")): + out.append(seq) + return out + + def _meta(self, seq): + import yaml + with open(os.path.join(seq, "meta.yaml")) as f: + return yaml.safe_load(f) + + def _label(self, seq, serial, frame): + k = (seq, serial, frame) + if k not in self._lab: + self._lab[k] = np.load(os.path.join(seq, serial, f"label_{frame:06d}.npz")) + return self._lab[k] + + @staticmethod + def _present_frames(d): + """Integer frames under serial dir `d` that have color+depth+label ALL + present (so every needed file loads). Robust to sparse/non-contiguous + frame ids and partially-extracted data.""" + fr = None + try: + names = os.listdir(d) + except FileNotFoundError: + return set() + for pre, suf in (("color_", ".jpg"), ("depth_", ".png"), ("label_", ".npz")): + s = set() + for f in names: + if f.startswith(pre) and f.endswith(suf): + try: + s.add(int(f[len(pre):-len(suf)])) + except ValueError: + pass + fr = s if fr is None else (fr & s) + return fr or set() + + def list_samples(self): + out = [] + for seq in self._seqs(): + m = self._meta(seq) + serials = [str(s) for s in m["realsense"]["serials"]] + obj_ids = list(m["object_ids"]) + # Frames usable by EVERY serial (each view's color/depth/label exist). + common = None + for s in serials: + fs = self._present_frames(os.path.join(seq, s)) + common = fs if common is None else (common & fs) + frames = sorted(common or set())[:: self.frame_stride] + for frame in frames: + for oi, obj in enumerate(obj_ids): + for anchor in serials: + aux = [s for s in serials if s != anchor] + aref = (seq, anchor, frame, oi, obj) + pool = [(seq, s, frame, oi, obj) for s in aux] + out.append(ViewSample(anchor=aref, aux_pool=pool, + obj_key=f"hocap/{obj}", meta={"obj": obj, "oi": oi})) + return out + + def load_view(self, ref): + seq, serial, frame, oi, obj = ref + rgb = np.asarray(Image.open(os.path.join(seq, serial, f"color_{frame:06d}.jpg")).convert("RGB"), + np.float32) / 255.0 + depth = np.asarray(Image.open(os.path.join(seq, serial, f"depth_{frame:06d}.png")), np.float32) / 1000.0 + lab = self._label(seq, serial, frame) + mask = (lab["seg_mask"] == (oi + 1)).astype(np.float32) # local ids 1..n + K = np.asarray(lab["cam_K"], np.float64).reshape(3, 3) # color K (depth aligned to color) + return {"rgb": rgb, "mask": mask, "depth_m": depth, "intr": _K_to_intr(K)} + + def _pose(self, seq, serial, frame, oi): + T = np.asarray(self._label(seq, serial, frame)["obj_poses"][oi], np.float64) # obj->cam + return T[:3, :3], T[:3, 3] + + def extrinsic_cv(self, ref): + seq, serial, frame, oi, obj = ref + R, t = self._pose(seq, serial, frame, oi) + return {"R": R, "t": t} + + def object_pose_cv(self, sample): + seq, serial, frame, oi, obj = sample.anchor + return self._pose(seq, serial, frame, oi) + + def mesh_for(self, sample): + return _load_mesh(os.path.join(self.root, "models", sample.meta["obj"], "textured_mesh.obj")) + + def canonical(self, obj_key): + return self.canon.get(obj_key) + + +# --------------------------------------------------------------------------- # +# ContactPose (type a, 3 synced Kinect v2; TRAIN-ONLY, static grasps). cTo is +# obj->cam OpenCV meters; mesh PLY is mm (*1e-3). The stored images are in a +# per-camera orientation (left/right PORTRAIT 540x960, middle rot180) while the +# intrinsics K are always in the LANDSCAPE 960x540 frame with the affine A +# mapping landscape->stored (utilities/misc.get_A). We inverse-warp each stored +# rgb/depth back to the landscape frame so the pipeline's diagonal-K unprojection +# is exact; masks are rendered from mesh+cTo+K in that same landscape frame. +# world := object frame (per-view extrinsic = that view's cTo). Static -> sample +# a few frames per grasp (default), not all ~558. +# --------------------------------------------------------------------------- # +class ContactPoseAdapter(PoseDatasetAdapter): + name = "contactpose" + + def __init__(self, root, canonical_cache, *, frames_per_grasp=4, frame_stride=None): + self.root = root # dataset/contactpose + self.data_dir = os.path.join(root, "contactpose_data") + self.models_dir = os.path.join(root, "object_models") + self.canon = _CanonCache(canonical_cache) + self.frames_per_grasp = frames_per_grasp + self._ann = {} + self._meshc = {} + + def _annot(self, gdir): + if gdir not in self._ann: + self._ann[gdir] = _json.load(open(os.path.join(gdir, "annotations.json"))) + return self._ann[gdir] + + @staticmethod + def _q2m(q): # transforms3d wxyz -> R + from scipy.spatial.transform import Rotation + return Rotation.from_quat([q[1], q[2], q[3], q[0]]).as_matrix() + + @classmethod + def _pm(cls, p): + T = np.eye(4); T[:3, :3] = cls._q2m(p["rotation"]); T[:3, 3] = p["translation"]; return T + + @staticmethod + def _get_A(cam): # landscape(960x540) -> stored + def flipud(H): return np.array([[1, 0, 0], [0, -1, H], [0, 0, 1]], float) + def fliplr(W): return np.array([[-1, 0, W], [0, 1, 0], [0, 0, 1]], float) + def transpose(): return np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1]], float) + if cam == "kinect2_left": return fliplr(540) @ transpose() + if cam == "kinect2_right": return flipud(960) @ transpose() + if cam == "kinect2_middle": return fliplr(960) @ flipud(540) + raise NotImplementedError(cam) + + def list_samples(self): + out = [] + for gdir in sorted(glob.glob(os.path.join(self.data_dir, "*", "*"))): + if not os.path.exists(os.path.join(gdir, "annotations.json")): + continue + ann = self._annot(gdir); obj = ann["object"] + cams = [c for c, cv in ann["cameras"].items() if cv.get("valid", True)] + if len(cams) < 2: + continue + nfr = len(ann["frames"]) + if nfr == 0: + continue + idxs = np.unique(np.linspace(0, nfr - 1, self.frames_per_grasp).round().astype(int)) + for fi in idxs: + for anchor in cams: + aux = [c for c in cams if c != anchor] + aref = (gdir, anchor, int(fi), obj) + pool = [(gdir, c, int(fi), obj) for c in aux] + out.append(ViewSample(anchor=aref, aux_pool=pool, + obj_key=f"contactpose/{obj}", meta={"obj": obj})) + return out + + def _K(self, gdir, cam): + k = self._annot(gdir)["cameras"][cam]["K"] + return np.array([[k["fx"], 0, k["cx"]], [0, k["fy"], k["cy"]], [0, 0, 1]], float) + + def _cTo(self, gdir, cam, fi): + ann = self._annot(gdir) + wTc = self._pm(ann["cameras"][cam]["wTc"]) + oTw = self._pm(ann["frames"][fi]["oTw"]) + return np.linalg.inv(oTw @ wTc) # obj->cam + + @staticmethod + def _unrotate(img, A, Wl=960, Hl=540): + """Resample the stored image into the landscape (Wl x Hl) frame: + landscape[v_l,u_l] = stored[ round(A @ [u_l,v_l,1]) ]. Correct for any A + (rot180 / transpose+flip) by construction; nearest-neighbour (NaN/edge safe).""" + us, vs = np.meshgrid(np.arange(Wl), np.arange(Hl)) + land = np.stack([us.ravel(), vs.ravel(), np.ones(Wl * Hl)], 0) # 3xN + st = A @ land + su = np.round(st[0]).astype(int); sv = np.round(st[1]).astype(int) + Hs, Ws = img.shape[:2] + ok = (su >= 0) & (su < Ws) & (sv >= 0) & (sv < Hs) + out = np.zeros((Hl * Wl,) + img.shape[2:], img.dtype) + src = img.reshape((Hs * Ws,) + img.shape[2:]) + out[ok] = src[sv[ok] * Ws + su[ok]] + return out.reshape((Hl, Wl) + img.shape[2:]) + + def _mesh(self, obj): + if obj not in self._meshc: + m = _load_mesh(os.path.join(self.models_dir, f"{obj}.ply")) + m.apply_scale(1e-3) # mm -> m + self._meshc[obj] = m + return self._meshc[obj] + + def load_view(self, ref): + gdir, cam, fi, obj = ref + A = self._get_A(cam) + cdir = os.path.join(gdir, "images_full", cam) + rgb_s = np.asarray(Image.open(os.path.join(cdir, "color", f"frame{fi:03d}.png")).convert("RGB"), + np.float32) / 255.0 + dep_s = np.asarray(Image.open(os.path.join(cdir, "depth", f"frame{fi:03d}.png")), np.float32) / 1000.0 + rgb = self._unrotate(rgb_s, A); depth = self._unrotate(dep_s, A) + K = self._K(gdir, cam) + H, W = rgb.shape[:2] + # render object silhouette (landscape frame) from mesh + cTo + K + cTo = self._cTo(gdir, cam, fi) + V = np.asarray(self._mesh(obj).vertices) + Xc = (cTo[:3, :3] @ V.T + cTo[:3, 3:4]).T + z = Xc[:, 2]; f = z > 1e-3 + uv = (K @ Xc[f].T).T + u = np.round(uv[:, 0] / uv[:, 2]).astype(int); v = np.round(uv[:, 1] / uv[:, 2]).astype(int) + mask = np.zeros((H, W), np.float32) + ok = (u >= 0) & (u < W) & (v >= 0) & (v < H) + mask[v[ok], u[ok]] = 1.0 + try: + import scipy.ndimage as _ndi + mask = _ndi.binary_fill_holes(_ndi.binary_dilation(mask > 0, iterations=3)).astype(np.float32) + except Exception: + pass + return {"rgb": rgb, "mask": mask, "depth_m": depth, "intr": _K_to_intr(K)} + + def extrinsic_cv(self, ref): + gdir, cam, fi, obj = ref + cTo = self._cTo(gdir, cam, fi) + return {"R": cTo[:3, :3], "t": cTo[:3, 3]} + + def object_pose_cv(self, sample): + gdir, cam, fi, obj = sample.anchor + cTo = self._cTo(gdir, cam, fi) + return cTo[:3, :3], cTo[:3, 3] + + def mesh_for(self, sample): + return self._mesh(sample.meta["obj"]) + + def canonical(self, obj_key): + return self.canon.get(obj_key) +# --------------------------------------------------------------------------- # +# H2O (type a-like: 4 static synced RGB-D cams cam0..cam3 + 1 ego cam4 [skipped]; +# TRAIN, 8 non-YCB objects). Per README+h2oplayer: +# camN/obj_pose_RT/{f:06d}.txt = class_id + 4x4 obj->cam OpenCV METERS, USED AS-IS +# (NO inversion; the "cam_to_obj" note in README refers to the DIFFERENT +# keypoint file obj_pose/, not obj_pose_RT/). +# camN/cam_pose/{f:06d}.txt = 4x4 cam->world (invert -> world->cam extrinsic) +# camN/cam_intrinsics.txt = fx fy cx cy w h (images undistorted) +# camN/depth/{f:06d}.png = uint16 mm (/1000 -> meters, registered to rgb) +# mesh object//<*.obj> = meters (single .obj per folder; 'spray' folder's +# obj is 'lotion_spray.obj' -> glob, don't assume). +# Multi-view set = the 4 static synced cams at one frame f (anchor-sweep); each cam +# supplies its OWN obj->cam pose and its world->cam extrinsic, so the base's +# relative_p3d_from_ext gives the exact cross-camera transform (static world at +# instant f). Masks are NOT shipped -> rendered from mesh+pose (CPU vertex splat). +# canonical_key = h2o/. Mark TRAIN (not held-out). +# --------------------------------------------------------------------------- # +_H2O_CLASSES = {1: "book", 2: "espresso", 3: "lotion", 4: "spray", + 5: "milk", 6: "cocoa", 7: "chips", 8: "cappuccino"} + + +class H2OAdapter(PoseDatasetAdapter): + name = "h2o" + + def __init__(self, root, canonical_cache, *, frame_stride=15, num_static_cams=4, subjects=None): + self.root = root # dataset/h2o + self.obj_root = os.path.join(root, "object_extracted", "object") + self.canon = _CanonCache(canonical_cache) + self.frame_stride = frame_stride + self.ncam = int(num_static_cams) # cam0..cam3 (skip cam4 ego) + self.subjects = subjects + self._intr = {} + self._meshc = {} + + # --- layout helpers ------------------------------------------------------ # + @staticmethod + def _pose_dir(camdir): + """H2O ships 'obj_pose_RT' (README); some mirrors use lowercase. Detect.""" + for nm in ("obj_pose_RT", "obj_pose_rt"): + p = os.path.join(camdir, nm) + if os.path.isdir(p): + return p + return os.path.join(camdir, "obj_pose_RT") + + def _seq_dirs(self): + """A 'sequence' dir contains cam0..cam{ncam-1}. Native layout is + /subject{1..4}/{h1,h2,k1,k2,o1,o2}//camN. Discover robustly by + finding cam0 dirs that carry an obj_pose_RT folder.""" + out = [] + for cam0 in sorted(glob.glob(os.path.join(self.root, "subject*", "*", "*", "cam0"))): + seq = os.path.dirname(cam0) + if not os.path.isdir(os.path.join(cam0, "rgb")): + continue + if self.subjects is not None: + if not any(f"subject{s}" == seq.split(os.sep)[-3] or str(s) in seq.split(os.sep)[-3] + for s in self.subjects): + continue + out.append(seq) + return out + + def _frames(self, seq): + rgbdir = os.path.join(seq, "cam0", "rgb") + return sorted(int(os.path.splitext(f)[0]) for f in os.listdir(rgbdir) if f.endswith(".png")) + + def _common_frames(self, seq): + """Integer frames present (rgb + depth + obj_pose_RT) in EVERY of the + ncam static cams, so anchor and all aux views load. Robust to in-flight + extraction where some cams lag.""" + common = None + for c in range(self.ncam): + camdir = os.path.join(seq, f"cam{c}") + triples = (("rgb", ".png"), ("depth", ".png"), (self._pose_dir(camdir), ".txt")) + perc = None + for sub, suf in triples: + d = sub if os.path.isabs(sub) else os.path.join(camdir, sub) + s = set() + try: + for f in os.listdir(d): + if f.endswith(suf): + try: + s.add(int(os.path.splitext(f)[0])) + except ValueError: + pass + except FileNotFoundError: + s = set() + perc = s if perc is None else (perc & s) + common = (perc or set()) if common is None else (common & (perc or set())) + return sorted(common or set()) + + def _load_RT(self, seq, cam, frame): + """camN/obj_pose_RT/{f:06d}.txt -> (class_id, 4x4 obj->cam).""" + camdir = os.path.join(seq, f"cam{cam}") + arr = np.loadtxt(os.path.join(self._pose_dir(camdir), f"{frame:06d}.txt")) + cls = int(round(float(arr[0]))) + M = np.asarray(arr[1:], np.float64).reshape(4, 4) + return cls, M + + def _intrinsics(self, seq, cam): + key = (seq, cam) + if key not in self._intr: + p = os.path.join(seq, f"cam{cam}", "cam_intrinsics.txt") + fx, fy, cx, cy, w, h = np.loadtxt(p).reshape(6) + self._intr[key] = {"fx": float(fx), "fy": float(fy), "ppx": float(cx), "ppy": float(cy), + "w": int(w), "h": int(h)} + return self._intr[key] + + def _mesh(self, name): + if name not in self._meshc: + cand = sorted(glob.glob(os.path.join(self.obj_root, name, "*.obj"))) + if not cand: + raise FileNotFoundError(f"h2o: no .obj in {os.path.join(self.obj_root, name)}") + self._meshc[name] = _load_mesh(cand[0]) # meters, no scaling + return self._meshc[name] + + @staticmethod + def _zbuffer_cpu(V, faces, R, t, K, H, W): + """Exact CPU z-buffer of the mesh at obj->cam pose (R,t). Returns a HxW + float depth map (0 where empty) = the FRONT-surface depth. Per-face + barycentric fill over each face's 2D bbox, keeping the nearest z. Used to + (i) get the true object silhouette and (ii) derive an occlusion-aware + VISIBLE mask by comparing to the sensor depth.""" + Xc = (R @ V.T + t.reshape(3, 1)).T + z = Xc[:, 2] + uv = (K @ Xc.T).T + u = uv[:, 0] / np.maximum(uv[:, 2], 1e-9) + v = uv[:, 1] / np.maximum(uv[:, 2], 1e-9) + P = np.stack([u, v], 1) + zbuf = np.full((H, W), np.inf, np.float64) + tri = P[faces] # (m,3,2) + triz = z[faces] # (m,3) + keep = (triz > 1e-3).all(1) + tri = tri[keep]; triz = triz[keep] + for (p0, p1, p2), (z0, z1, z2) in zip(tri, triz): + x0, y0 = p0; x1, y1 = p1; x2, y2 = p2 + minx = int(max(np.floor(min(x0, x1, x2)), 0)); maxx = int(min(np.ceil(max(x0, x1, x2)), W - 1)) + miny = int(max(np.floor(min(y0, y1, y2)), 0)); maxy = int(min(np.ceil(max(y0, y1, y2)), H - 1)) + if maxx < minx or maxy < miny: + continue + xs, ys = np.meshgrid(np.arange(minx, maxx + 1), np.arange(miny, maxy + 1)) + det = (y1 - y2) * (x0 - x2) + (x2 - x1) * (y0 - y2) + if abs(det) < 1e-9: + continue + a = ((y1 - y2) * (xs - x2) + (x2 - x1) * (ys - y2)) / det + b = ((y2 - y0) * (xs - x2) + (x0 - x2) * (ys - y2)) / det + c = 1.0 - a - b + ins = (a >= 0) & (b >= 0) & (c >= 0) + if not ins.any(): + continue + zpix = a * z0 + b * z1 + c * z2 + yy = ys[ins]; xx = xs[ins]; zz = zpix[ins] + cur = zbuf[yy, xx] + upd = zz < cur + zbuf[yy[upd], xx[upd]] = zz[upd] + zbuf[~np.isfinite(zbuf)] = 0.0 + return zbuf + + # --- interface ----------------------------------------------------------- # + def list_samples(self): + out = [] + for seq in self._seq_dirs(): + frames = self._common_frames(seq) # present in ALL ncam cams + if not frames: + continue + for f in frames[:: self.frame_stride] or frames[:1]: + try: + cls, _ = self._load_RT(seq, 0, f) # scene-global object at frame f + except Exception: + continue + if cls <= 0 or cls not in _H2O_CLASSES: + continue # background / no object + obj = _H2O_CLASSES[cls] + cams = list(range(self.ncam)) + for anchor_ci in cams: + aux = [c for c in cams if c != anchor_ci] + aref = {"seq": seq, "cam": anchor_ci, "frame": f} + pool = [{"seq": seq, "cam": c, "frame": f} for c in aux] + out.append(ViewSample(anchor=aref, aux_pool=pool, + obj_key=f"h2o/{obj}", meta={"cls": cls, "obj": obj})) + return out + + def load_view(self, ref): + seq, cam, frame = ref["seq"], ref["cam"], ref["frame"] + camdir = os.path.join(seq, f"cam{cam}") + rgb = np.asarray(Image.open(os.path.join(camdir, "rgb", f"{frame:06d}.png")).convert("RGB"), + np.float32) / 255.0 + depth = np.asarray(Image.open(os.path.join(camdir, "depth", f"{frame:06d}.png")), + np.float32) / 1000.0 + intr = self._intrinsics(seq, cam) + K = np.array([[intr["fx"], 0, intr["ppx"]], [0, intr["fy"], intr["ppy"]], [0, 0, 1]], np.float64) + H, W = rgb.shape[:2] + # masks not shipped -> render the object silhouette (exact CPU z-buffer) + # from THIS cam's mesh + obj->cam pose, then derive an OCCLUSION-AWARE + # VISIBLE mask (consistent with the other adapters, which all supply + # occlusion-aware visible masks): drop pixels where the sensor sees a + # surface clearly IN FRONT of the object (a hand occluder). This is a + # physical visibility test (nearer surface only); pixels where the sensor + # is behind/absent are KEPT so a wrong pose cannot hide there. + cls, M = self._load_RT(seq, cam, frame) + mask = np.zeros((H, W), np.float32) + if cls in _H2O_CLASSES: + mesh = self._mesh(_H2O_CLASSES[cls]) + V = np.asarray(mesh.vertices, np.float64) + F = np.asarray(mesh.faces, np.int64) + zbuf = self._zbuffer_cpu(V, F, M[:3, :3], M[:3, 3], K, H, W) + sil = zbuf > 1e-6 + occ = sil & (depth > 1e-3) & (depth < zbuf - 0.02) # nearer surface => occluder + mask = (sil & ~occ).astype(np.float32) + return {"rgb": rgb, "mask": mask, "depth_m": depth, "intr": _K_to_intr(K)} + + def extrinsic_cv(self, ref): # world->cam (invert cam->world) + seq, cam, frame = ref["seq"], ref["cam"], ref["frame"] + c2w = np.loadtxt(os.path.join(seq, f"cam{cam}", "cam_pose", f"{frame:06d}.txt")).reshape(4, 4) + w2c = np.linalg.inv(c2w) + return {"R": w2c[:3, :3].astype(np.float64), "t": w2c[:3, 3].astype(np.float64)} + + def object_pose_cv(self, sample): # anchor cam's obj->cam, AS-IS + a = sample.anchor + _, M = self._load_RT(a["seq"], a["cam"], a["frame"]) + return M[:3, :3].astype(np.float64), M[:3, 3].astype(np.float64) + + def mesh_for(self, sample): + return self._mesh(sample.meta["obj"]) + + def canonical(self, obj_key): + return self.canon.get(obj_key) + + +# --------------------------------------------------------------------------- # +# HOGraspNet (type a, 4 synced Kinect cams; TRAIN; mocap GT; grasp = hand +# occlusion). world := the master ('mas') camera frame -> object_mat is +# obj->WORLD and each cam's calibration.extrinsic is WORLD->cam (OpenCV), so +# T_oc = M4(extrinsic) @ object_mat is obj->cam directly and +# relative_p3d_from_ext(anchor,aux) yields the exact cross-camera transform +# (same synced-rig pattern as DexYCB/HOCap). Translations are stored in cm +# (-> /100 for meters). The 4 cams at one frame-instant are the views +# (anchor-sweep). We use the SCANNED meshes (NOT dexycb CAD; different frame) +# with the MANDATORY per-object scale cfg._OBJECT_SCALE_FIXED[obj_idx-1] and +# obj_idx = HOGraspNet's own 1..30 ordering (config.OBJType), read from the +# sequence name (231023_S01_obj__grasp_). We consume the cropped +# 640x480 source_augmented RGB/depth, shifting cx,cy by the (hand-centred) +# crop bbox offset. obj_mask (pseudo GT) comes from extra_data. +# --------------------------------------------------------------------------- # +# HOGraspNet's own object ordering (config.OBJType, 1..30) -> mandatory per-object +# scale (config.cfg._OBJECT_SCALE_FIXED); index 0 == obj_idx 1. +_HOG_OBJECT_SCALE_FIXED = [ + 1.0, 1.0, 1.0, 1.0, 0.8296698468, 1.0, 1.0, 1.0, 1.0, 1.0, + 0.1035083229, 1.0, 0.6706711338, 1.0, 1.0, 0.43, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, +] + + +class HOGraspNetAdapter(PoseDatasetAdapter): + name = "hograspnet" + CAMS = ("mas", "sub1", "sub2", "sub3") # mas == world origin + _CROP_W, _CROP_H = 640, 480 # source_augmented crop size + _FULL_W, _FULL_H = 1920, 1080 # full-res intrinsics frame + + def __init__(self, root, canonical_cache, *, frame_stride=1, subjects=None): + self.root = root # dataset/hograspnet + self.anno_root = os.path.join(root, "labeling_data") + self.src_root = os.path.join(root, "source_augmented") + self.extra_root = os.path.join(root, "extra_data") + self.models_root = os.path.join(root, "obj_scanned_models") + self.canon = _CanonCache(canonical_cache) + self.frame_stride = frame_stride + self.subjects = set(subjects) if subjects is not None else None # e.g. {"S01",...} + self._anno = {} + self._objdir_map = None + + # -- object index (1..30) -> scanned-model dir name ("01_cracker_box") -- # + def _objdir(self): + if self._objdir_map is None: + m = {} + for d in os.listdir(self.models_root): + if os.path.isdir(os.path.join(self.models_root, d)): + try: + m[int(d.split("_")[0])] = d + except ValueError: + continue + self._objdir_map = m + return self._objdir_map + + def _anno_path(self, ref): + return os.path.join(self.anno_root, ref["seq"], ref["trial"], "annotation", + ref["cam"], f'{ref["cam"]}_{ref["frame"]}.json') + + def _read_anno(self, ref): + p = self._anno_path(ref) + if p not in self._anno: + with open(p, "r", encoding="UTF-8-SIG") as f: + self._anno[p] = _json.load(f) + return self._anno[p] + + def _crop_offset(self, hand_2d): + """Reproduce HOGraspNet util.utils.extractBbox: a fixed 640x480 window + centred on the hand 2D keypoints (full-res coords), clamped to image + bounds. Returns the integer (x0, y0) top-left offset used to crop the + source_augmented RGB/depth (so cx,cy shift by -x0,-y0).""" + x_min, x_max = float(hand_2d[:, 0].min()), float(hand_2d[:, 0].max()) + y_min, y_max = float(hand_2d[:, 1].min()), float(hand_2d[:, 1].max()) + x0 = max(0.0, (x_min + x_max) / 2 - self._CROP_W / 2) + y0 = max(0.0, (y_min + y_max) / 2 - self._CROP_H / 2) + if x0 + self._CROP_W > self._FULL_W: + x0 = self._FULL_W - self._CROP_W + if y0 + self._CROP_H > self._FULL_H: + y0 = self._FULL_H - self._CROP_H + return int(x0), int(y0) + + def list_samples(self): + out = [] + if not os.path.isdir(self.anno_root): + return out + objmap = self._objdir() + for seq in sorted(os.listdir(self.anno_root)): + seqdir = os.path.join(self.anno_root, seq) + if not os.path.isdir(seqdir): + continue + parts = seq.split("_") + if len(parts) < 6: + continue + subject = parts[1] # 'S01' + if self.subjects is not None and subject not in self.subjects: + continue + try: + obj_idx = int(parts[3]) + except ValueError: + continue + objdir = objmap.get(obj_idx) + if objdir is None: + continue + obj_key = f"hograspnet/{objdir}" + for trial in sorted(os.listdir(seqdir)): + annob = os.path.join(seqdir, trial, "annotation") + srcb = os.path.join(self.src_root, seq, trial, "rgb_crop") + if not os.path.isdir(annob) or not os.path.isdir(srcb): + continue # need cropped frames to load + cams = [c for c in self.CAMS if os.path.isdir(os.path.join(annob, c)) + and os.path.isdir(os.path.join(srcb, c))] + if not cams: + continue + framesets = {} + for c in cams: + anno_fs = set() + for fn in os.listdir(os.path.join(annob, c)): + if fn.endswith(".json"): + anno_fs.add(fn[:-5].split("_")[-1]) + + def _crop_frames(cropdir, suf, _c=c): + s = set() + cd = os.path.join(self.src_root, seq, trial, cropdir, _c) + try: + for fn in os.listdir(cd): + if fn.endswith(suf): + s.add(fn[:-len(suf)].split("_")[-1]) + except FileNotFoundError: + pass + return s + # frame usable only if anno json AND rgb_crop AND depth_crop exist + framesets[c] = anno_fs & _crop_frames("rgb_crop", ".jpg") & _crop_frames("depth_crop", ".png") + allframes = sorted(set().union(*framesets.values()), key=lambda x: int(x)) + for fi, frame in enumerate(allframes): + if self.frame_stride > 1 and (fi % self.frame_stride): + continue + present = [c for c in cams if frame in framesets[c]] + if not present: + continue + for anchor in present: + aux = [c for c in present if c != anchor] + aref = {"seq": seq, "trial": trial, "cam": anchor, "frame": frame} + pool = [{"seq": seq, "trial": trial, "cam": c, "frame": frame} for c in aux] + out.append(ViewSample(anchor=aref, aux_pool=pool, obj_key=obj_key, + meta={"obj_idx": obj_idx, "objdir": objdir})) + return out + + def load_view(self, ref): + anno = self._read_anno(ref) + K = np.squeeze(np.asarray(anno["calibration"]["intrinsic"], np.float64)).reshape(3, 3) + hand_2d = np.squeeze(np.asarray(anno["hand"]["projected_2D_pose_per_cam"], np.float64)) + x0, y0 = self._crop_offset(hand_2d) + Kc = K.copy() + Kc[0, 2] -= x0 # cropped-frame principal point + Kc[1, 2] -= y0 + base = f'{ref["cam"]}_{ref["frame"]}' + rgb_p = os.path.join(self.src_root, ref["seq"], ref["trial"], "rgb_crop", ref["cam"], base + ".jpg") + dep_p = os.path.join(self.src_root, ref["seq"], ref["trial"], "depth_crop", ref["cam"], base + ".png") + msk_p = os.path.join(self.extra_root, ref["seq"], ref["trial"], "obj_mask", ref["cam"], base + ".png") + rgb = np.asarray(Image.open(rgb_p).convert("RGB"), np.float32) / 255.0 + depth = np.asarray(Image.open(dep_p), np.float32) / 1000.0 # mm -> m + if os.path.exists(msk_p): + ma = np.asarray(Image.open(msk_p)) + if ma.ndim == 3: + ma = ma[..., 0] + mask = (ma >= 128).astype(np.float32) + else: + mask = np.zeros(depth.shape, np.float32) + return {"rgb": rgb, "mask": mask, "depth_m": depth, "intr": _K_to_intr(Kc)} + + def _extrinsic4(self, ref): + M = np.squeeze(np.asarray(self._read_anno(ref)["calibration"]["extrinsic"], np.float64)).reshape(3, 4) + M4 = np.eye(4, dtype=np.float64) + M4[:3, :] = M + return M4 + + def extrinsic_cv(self, ref): # world(mas)->cam, meters + M4 = self._extrinsic4(ref) + return {"R": M4[:3, :3], "t": M4[:3, 3] / 100.0} + + def object_pose_cv(self, sample): # obj->cam, anchor view, meters + a = sample.anchor + M4 = self._extrinsic4(a) + om = np.squeeze(np.asarray(self._read_anno(a)["Mesh"][0]["object_mat"], np.float64)).reshape(4, 4) + T = M4 @ om # obj->cam (translation in cm) + return T[:3, :3], T[:3, 3] / 100.0 + + def mesh_for(self, sample): + objdir = sample.meta["objdir"] + S = _HOG_OBJECT_SCALE_FIXED[sample.meta["obj_idx"] - 1] # mandatory per-object scale + m = _load_mesh(os.path.join(self.models_root, objdir, f"{objdir}.obj")) + m.apply_scale(S * 1e-2) # (S * V_raw) / 100 -> meters + return m + + def canonical(self, obj_key): + return self.canon.get(obj_key) + + +# --------------------------------------------------------------------------- # +# DexH2R (type a, 4 static synced Kinect RGB-D cams; TRAIN-ONLY pseudo-GT; +# 56 non-YCB objects; human->robot handover => heavy hand/robot occlusion). +# Per docs/dataset_convention_spec.md (DexH2R) + camera_calibration: +# obj_pose.pt = [T,4,4] = T_world<-obj (world == hand_arm base frame). +# Per Kinect i in {0,1,2,3} (STATIC extrinsics): +# world->cam_i W2C_i = inv(cali0i) @ hand_arm_mesh_to_kinect_pcd_0 +# obj->cam_i T_oc = W2C_i @ obj_pose[t] (OpenCV, meters) +# (equivalently T_oc = inv(cali0i) @ hand_arm @ obj_pose[t]; round-trips the +# repo's cam->world = inv(hand_arm) @ cali0i, verified in verify_conventions.py.) +# Extrinsics come from the PER-TRIAL calibration/ folder (kinect/cali0i.json + +# hand_arm_mesh_to_kinect_pcd_0.json). Intrinsics are NOT shipped per-trial -> +# read from the reference camera_calibration/ (Kinect 8-param distortion; the +# stored RGB/depth are already undistorted so K is used as a plain pinhole). +# depth = uint16 PNG in mm (/1000 -> meters), same res (1080x1920) & K as RGB +# (registered). Meshes object_model/.obj are meters in the pose frame. +# The 4 Kinect cams at one frame t are the multi-view set (anchor-sweep); each +# supplies its own W2C_i so the base's relative_p3d_from_ext gives the exact +# cross-camera transform (shared static world = hand_arm base at instant t). +# No 2D masks are shipped -> render an OCCLUSION-AWARE visible silhouette from +# mesh+pose (same CPU z-buffer + sensor-occluder test as the H2O adapter). +# RealSense 0/1 also carry depth but their world pose is time-varying via +# ShadowHand FK on qpos.pt (needs the shadow_asserts URDF) -> DEFERRED; the 4 +# Kinect views are sufficient for multi-view. ZCAM (12 RGB-only) dropped. +# canonical_key = dexh2r/. TRAIN, pseudo-GT. +# Layout (after extract, prefix stripped): /extracted//// +# {obj_pose.pt, calibration/, kinect//{rgb,depth}/.{jpg,png}, ...} +# --------------------------------------------------------------------------- # +class DexH2RAdapter(PoseDatasetAdapter): + name = "dexh2r" + KINECTS = (0, 1, 2, 3) # the 4 static Kinect RGB-D cams + + def __init__(self, root, canonical_cache, *, frame_stride=15, subjects=None, + split=None, extracted_subdir="extracted"): + self.root = root # dataset/dexh2r + self.ext_root = os.path.join(root, extracted_subdir) # extracted subjects live here + self.mesh_dir = os.path.join(root, "object_model", "object_model") + # reference intrinsics (per-trial calibration ships extrinsics only) + self.ref_calib = os.path.join(root, "camera_calibration", "camera_calibration") + self.canon = _CanonCache(canonical_cache) + self.frame_stride = int(frame_stride) + # `subjects`/`split`: restrict to a subset of subject-id strings; split reads + # dataset_split.json (_sub_names) when subjects is None. + if subjects is None and split is not None: + sp = _json.load(open(os.path.join(root, "dataset_split.json"))) + subjects = [str(s) for s in sp[f"{split}_sub_names"]] + self.subjects = set(str(s) for s in subjects) if subjects is not None else None + self._intr = {} + self._ext = {} + self._meshc = {} + self._nframe = {} + + # -- calibration ---------------------------------------------------------- # + @staticmethod + def _seven2mat(js): + from scipy.spatial.transform import Rotation + v = js["value0"]; r = v["rotation"]; t = v["translation"] + M = np.eye(4, dtype=np.float64) + M[:3, :3] = Rotation.from_quat([r["x"], r["y"], r["z"], r["w"]]).as_matrix() + M[:3, 3] = [t["x"], t["y"], t["z"]] + return M + + def _w2c(self, trial_dir, i): + """world(hand_arm)->cam_i (OpenCV, meters), from the PER-TRIAL calibration.""" + key = (trial_dir, i) + if key not in self._ext: + cdir = os.path.join(trial_dir, "calibration") + hand_arm = self._seven2mat(_json.load(open( + os.path.join(cdir, "hand_arm_mesh_to_kinect_pcd_0.json")))) + cali0i = self._seven2mat(_json.load(open( + os.path.join(cdir, "kinect", f"cali0{i}.json")))) + self._ext[key] = np.linalg.inv(cali0i) @ hand_arm # world->cam_i + return self._ext[key] + + def _K(self, i): + if i not in self._intr: + d = _json.load(open(os.path.join( + self.ref_calib, "camera_intrinsics", "kinect", f"camera_intrinsics_{i}.json"))) + self._intr[i] = np.asarray(d["camera_intrinsics"], np.float64).reshape(3, 3) + return self._intr[i] + + def _mesh(self, name): + if name not in self._meshc: + self._meshc[name] = _load_mesh(os.path.join(self.mesh_dir, f"{name}.obj")) + return self._meshc[name] + + def _obj_pose(self, trial_dir, frame): + """T_world<-obj at frame (4x4).""" + op = torch.load(os.path.join(trial_dir, "obj_pose.pt"), map_location="cpu") + return np.asarray(op[frame], np.float64) + + def _T_oc(self, trial_dir, i, frame): + return self._w2c(trial_dir, i) @ self._obj_pose(trial_dir, frame) # obj->cam_i + + def _oplen(self, trial_dir): + """Number of obj_pose timesteps (upper bound on valid frame index).""" + if trial_dir not in self._nframe: + try: + op = torch.load(os.path.join(trial_dir, "obj_pose.pt"), map_location="cpu") + self._nframe[trial_dir] = int(len(op)) + except Exception: + self._nframe[trial_dir] = 0 + return self._nframe[trial_dir] + + @staticmethod + def _cam_frames(trial_dir, i): + """Integer frames where kinect cam i has BOTH rgb .jpg and depth .png + (files are named '.jpg'/'.png'; robust to sparse ids).""" + cd = os.path.join(trial_dir, "kinect", str(i)) + out = None + for sub, suf in (("rgb", ".jpg"), ("depth", ".png")): + s = set() + try: + for f in os.listdir(os.path.join(cd, sub)): + if f.endswith(suf): + try: + s.add(int(f[:-len(suf)])) + except ValueError: + pass + except FileNotFoundError: + return set() + out = s if out is None else (out & s) + return out or set() + + # -- interface ------------------------------------------------------------ # + def list_samples(self): + out = [] + if not os.path.isdir(self.ext_root): + return out + for subj in sorted(os.listdir(self.ext_root)): + subj_dir = os.path.join(self.ext_root, subj) + if not os.path.isdir(subj_dir): + continue + if self.subjects is not None and subj not in self.subjects: + continue + for name in sorted(os.listdir(subj_dir)): + obj_dir = os.path.join(subj_dir, name) + if not os.path.isdir(obj_dir) or not os.path.exists( + os.path.join(self.mesh_dir, f"{name}.obj")): + continue + for seq in sorted(os.listdir(obj_dir)): + trial_dir = os.path.join(obj_dir, seq) + if not os.path.exists(os.path.join(trial_dir, "obj_pose.pt")): + continue + if not os.path.exists(os.path.join(trial_dir, "calibration", + "hand_arm_mesh_to_kinect_pcd_0.json")): + continue + cams = [i for i in self.KINECTS + if os.path.isdir(os.path.join(trial_dir, "kinect", str(i), "depth"))] + if len(cams) < 1: + continue + op_n = self._oplen(trial_dir) + if op_n < 1: + continue + # frames actually present (rgb+depth) per cam, bounded by obj_pose + per = {i: {f for f in self._cam_frames(trial_dir, i) if 0 <= f < op_n} + for i in cams} + allframes = sorted(set().union(*per.values()) if per else set()) + for fi, frame in enumerate(allframes): + if self.frame_stride > 1 and (fi % self.frame_stride): + continue + present = [i for i in cams if frame in per[i]] # cams with this frame + if not present: + continue + for anchor in present: + aux = [c for c in present if c != anchor] + aref = {"trial": trial_dir, "cam": anchor, "frame": frame, "name": name} + pool = [{"trial": trial_dir, "cam": c, "frame": frame, "name": name} + for c in aux] + out.append(ViewSample(anchor=aref, aux_pool=pool, + obj_key=f"dexh2r/{name}", + meta={"obj": name, "subject": subj, + "pseudo_gt": True, "train": True})) + return out + + def load_view(self, ref): + trial, i, frame, name = ref["trial"], ref["cam"], ref["frame"], ref["name"] + cdir = os.path.join(trial, "kinect", str(i)) + rgb = np.asarray(Image.open(os.path.join(cdir, "rgb", f"{frame}.jpg")).convert("RGB"), + np.float32) / 255.0 + depth = np.asarray(Image.open(os.path.join(cdir, "depth", f"{frame}.png")), + np.float32) / 1000.0 # uint16 mm -> m + K = self._K(i) + H, W = depth.shape[:2] + # No 2D masks shipped -> render the object silhouette (exact CPU z-buffer) + # at THIS cam's obj->cam pose, then derive an OCCLUSION-AWARE VISIBLE mask + # (identical policy to the H2O adapter): drop pixels where the sensor sees + # a surface clearly IN FRONT of the object (hand/robot occluder); KEEP + # pixels where the sensor is behind/absent so a wrong pose cannot hide. + T = self._T_oc(trial, i, frame) + mesh = self._mesh(name) + V = np.asarray(mesh.vertices, np.float64) + F = np.asarray(mesh.faces, np.int64) + zbuf = H2OAdapter._zbuffer_cpu(V, F, T[:3, :3], T[:3, 3], K, H, W) + sil = zbuf > 1e-6 + occ = sil & (depth > 1e-3) & (depth < zbuf - 0.02) # nearer surface => occluder + mask = (sil & ~occ).astype(np.float32) + return {"rgb": rgb, "mask": mask, "depth_m": depth, "intr": _K_to_intr(K)} + + def extrinsic_cv(self, ref): # world(hand_arm)->cam_i + W2C = self._w2c(ref["trial"], ref["cam"]) + return {"R": W2C[:3, :3], "t": W2C[:3, 3]} + + def object_pose_cv(self, sample): # obj->cam, anchor view + a = sample.anchor + T = self._T_oc(a["trial"], a["cam"], a["frame"]) + return T[:3, :3], T[:3, 3] + + def mesh_for(self, sample): + return self._mesh(sample.meta["obj"]) + + def canonical(self, obj_key): + return self.canon.get(obj_key) + + +ADAPTERS = { + "dexycb": DexYCBAdapter, "housecat6d": HouseCat6DAdapter, "pace": PACEAdapter, + "graspnet": GraspNetAdapter, "ycbv": YCBVAdapter, "contactpose": ContactPoseAdapter, + "hocap": HOCapAdapter, "h2o": H2OAdapter, "hograspnet": HOGraspNetAdapter, "dexh2r": DexH2RAdapter, +} diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/base_dataset.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/base_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..4f8692b6cc28950d1cead89e1a5295c0a39cc57e --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/base_dataset.py @@ -0,0 +1,272 @@ +"""Dataset-agnostic multi-view 6D-pose dataset (base + adapter protocol). + +This is the generalization of ``MultiViewDexYCB``: ALL geometry / SSI / GT-layout +/ view-construction math is reused VERBATIM from ``dexycb_dataset`` (the helper +functions are imported, not re-implemented), and the ONLY dataset-specific logic +lives behind a thin ``PoseDatasetAdapter``. + +One sample = one (object-state, anchor-view) with >=1 aux views. The two +multi-view regimes both funnel into the same anchor+aux construction: + (a) synced multi-cam rig -> aux views are OTHER cameras at the same instant; + per-view extrinsic = the constant rig calibration (tag->cam, OpenCV). + (b) static scene + moving cam -> aux views are OTHER FRAMES of the same static + sequence; per-view extrinsic = that frame's world->cam pose (OpenCV). +In both cases the adapter yields, per view ref, an OpenCV tag/world->cam +extrinsic ``{R,t}``; the base computes the anchor-relative rigid transform with +the SAME formula as ``dexycb_calib.relative_p3d`` and re-expresses every aux +pointmap in the anchor PyTorch3D frame under one shared anchor SSI. + +The adapter delivers the object pose as OpenCV **obj->cam, meters** for the +anchor view (== DexYCB pose_y semantics); the base applies D=diag(-1,-1,1) and +the canonical M-mapping. Adapters MUST NOT pre-apply D or M. + +Output dict is byte-identical in structure to ``MultiViewDexYCB.__getitem__`` and +feeds ``collate_multiview`` unchanged. +""" +from __future__ import annotations + +import random +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import torch + +from sam3d_objects.geometry.geometry import matrix_to_rot6d, normalize_rot6d + +from . import canonical +# Reuse the proven geometry helpers VERBATIM (imported, not re-implemented). +from .dexycb_dataset import ( + IMG_SIZE, + _apply_ssi, + _cam_feature, + _ssi_params, + _to_p3d, + _unproject, + _build_view_fn, +) + +# --- Robust data loading: a large multi-dataset loader WILL hit corrupt / +# truncated / missing files (e.g. a PNG half-written during a disk-full). One +# bad file must NEVER kill a run. We (1) tolerate mildly-truncated images, +# (2) force-decode on open so a hard-broken file raises WITH its path (logged +# once to _bad_files.log), and (3) retry another random sample in __getitem__. --- +import os as _os +from PIL import Image as _PILImage, ImageFile as _PILImageFile + +_PILImageFile.LOAD_TRUNCATED_IMAGES = True +from mvsam3d import paths as _mv_paths +_BAD_FILES_LOG = _mv_paths.BAD_FILES_LOG +_bad_seen: set = set() + + +def _log_bad_file(path, err): + k = str(path) + if k in _bad_seen: + return + _bad_seen.add(k) + try: + with open(_BAD_FILES_LOG, "a") as _fh: + _fh.write(f"{k}\t{type(err).__name__}: {err}\n") + except Exception: + pass + + +_orig_pil_open = _PILImage.open + + +def _safe_pil_open(fp, *a, **k): + """Force-decode on open so a broken/truncated image raises here (with its + path) rather than deep inside np.asarray/convert with no context.""" + try: + im = _orig_pil_open(fp, *a, **k) + im.load() + return im + except Exception as _e: + _log_bad_file(getattr(fp, "name", fp), _e) + raise + + +_PILImage.open = _safe_pil_open # every adapter uses the shared PIL.Image.open + +_D = np.diag([-1.0, -1.0, 1.0]).astype(np.float64) +MIN_VALID_PM_SCALE = 1e-6 + + +@dataclass +class ViewSample: + """One training unit. ``anchor`` and ``aux_pool`` are adapter-defined view + refs; ``obj_key`` is the canonical-latent cache key.""" + anchor: Any + aux_pool: List[Any] + obj_key: str + meta: dict = field(default_factory=dict) + + +class PoseDatasetAdapter(ABC): + """Minimal dataset-specific interface. Everything returned is in the UNIFIED + contract: OpenCV, meters, depth registered to RGB.""" + + name: str = "base" + + @abstractmethod + def list_samples(self) -> List[ViewSample]: + ... + + @abstractmethod + def load_view(self, ref: Any) -> Dict[str, Any]: + """Return {'rgb':[H,W,3] float[0,1], 'mask':[H,W] float{0,1}, + 'depth_m':[H,W] float meters (registered to rgb), 'intr':{fx,fy,ppx,ppy}}.""" + ... + + @abstractmethod + def extrinsic_cv(self, ref: Any) -> Dict[str, np.ndarray]: + """Return {'R':(3,3),'t':(3,)} tag/world->cam, OpenCV, meters.""" + ... + + @abstractmethod + def object_pose_cv(self, sample: ViewSample) -> Tuple[np.ndarray, np.ndarray]: + """Return (R_cv, t_cv): object(CAD)->camera in the ANCHOR view, OpenCV, + meters. x_cam = R_cv @ p_cad + t_cv.""" + ... + + @abstractmethod + def canonical(self, obj_key: str) -> dict: + """Return the cached canonical dict: {latent[8,16,16,16], + occupancy[64,64,64] or [1,...], max_extent float, center_cad (3,)}.""" + ... + + +def relative_p3d_from_ext(anchor_ext: Dict[str, np.ndarray], + aux_ext: Dict[str, np.ndarray]): + """(Rrel, trel) mapping a PyTorch3D point in aux cam -> anchor cam. + + IDENTICAL formula to dexycb_calib.relative_p3d, but taking the two OpenCV + tag/world->cam extrinsics directly (works for both synced-rig and + static-scene-moving-cam regimes).""" + Ra, ta = anchor_ext["R"], anchor_ext["t"] + Rb, tb = aux_ext["R"], aux_ext["t"] + RaRbT = Ra @ Rb.T + Rrel = _D @ RaRbT @ _D + trel = _D @ ta - _D @ (RaRbT @ tb) + return Rrel.astype(np.float64), trel.astype(np.float64) + + +def _transform_pointmap(pm: np.ndarray, Rrel: np.ndarray, trel: np.ndarray) -> np.ndarray: + H, W, _ = pm.shape + flat = pm.reshape(-1, 3) + valid = np.isfinite(flat).all(axis=1) + out = flat.copy() + out[valid] = flat[valid] @ Rrel.T + trel + return out.reshape(H, W, 3) + + +class MultiViewPoseDataset(torch.utils.data.Dataset): + """Adapter-driven multi-view 6D-pose dataset. Behaviourally identical to + MultiViewDexYCB when driven by DexYCBAdapter.""" + + def __init__(self, adapter: PoseDatasetAdapter, *, + image_size: int = IMG_SIZE, + random_views: bool = False, + num_aux_views: int = 7, + min_aux_views: int = 1, + max_aux_views: int = 3, + depth_near: float = 0.1, + depth_far: float = 3.0, + ssi_mode: str = "scene"): + self.adapter = adapter + self.image_size = image_size + self.random_views = random_views + self.num_aux = num_aux_views + self.min_aux = min_aux_views + self.max_aux = max_aux_views + self.depth_near = depth_near + self.depth_far = depth_far + if ssi_mode not in ("scene", "object_centric"): + raise ValueError(f"ssi_mode must be 'scene' or 'object_centric', got {ssi_mode!r}") + self.ssi_mode = ssi_mode + if random_views and not (1 <= min_aux_views <= max_aux_views): + raise ValueError(f"random_views requires 1<=min_aux<=max_aux, got " + f"{min_aux_views}..{max_aux_views}") + self.samples: List[ViewSample] = adapter.list_samples() + + def __len__(self): + return len(self.samples) + + def _anchor_pm(self, ref): + v = self.adapter.load_view(ref) + pm = _to_p3d(_unproject(v["depth_m"], v["intr"], self.depth_near, self.depth_far)) + return v, pm + + def __getitem__(self, idx: int) -> dict: + # Robust wrapper: if a sample references a corrupt/truncated/missing file, + # log it once and fall back to another random sample — never crash the run. + n = len(self.samples) + for _try in range(24): + j = idx if _try == 0 else random.randrange(n) + try: + return self._getitem_impl(j) + except Exception as _e: # noqa: BLE001 - any decode/IO failure -> skip sample + _log_bad_file(f"[sample idx={j} adapter={getattr(self.adapter, 'name', '?')}]", _e) + # Exhausted retries: surface the error (don't hide a genuine bug forever). + return self._getitem_impl(random.randrange(n)) + + def _getitem_impl(self, idx: int) -> dict: + s = self.samples[idx] + + # ---- anchor view + shared SSI from anchor pointmap ---- + a_view, a_pm = self._anchor_pm(s.anchor) + shift, scale = _ssi_params(a_pm, mask=a_view["mask"], mode=self.ssi_mode) + anchor_ext = self.adapter.extrinsic_cv(s.anchor) + + views = [_build_view_fn(a_view["rgb"], a_view["mask"], + _apply_ssi(a_pm, shift, scale), True, self.image_size)] + views[0]["cam"] = _cam_feature(np.eye(3), np.zeros(3), shift, scale) + + # ---- choose aux views ---- + pool = s.aux_pool + if not pool: + aux_refs = [] # anchor-only sample (no other views available) + elif self.random_views: + k = random.randint(min(self.min_aux, len(pool)), min(self.max_aux, len(pool))) + aux_refs = random.sample(pool, k) + else: + aux_refs = pool[: self.num_aux] + + for aux in aux_refs: + x = self.adapter.load_view(aux) + x_pm = _to_p3d(_unproject(x["depth_m"], x["intr"], self.depth_near, self.depth_far)) + Rrel, trel = relative_p3d_from_ext(anchor_ext, self.adapter.extrinsic_cv(aux)) + x_pm = _transform_pointmap(x_pm, Rrel, trel) + av = _build_view_fn(x["rgb"], x["mask"], _apply_ssi(x_pm, shift, scale), + False, self.image_size) + av["cam"] = _cam_feature(Rrel, trel, shift, scale) + views.append(av) + + # ---- GT layout (anchor cam, PyTorch3D) ---- + R_cv, t_cv = self.adapter.object_pose_cv(s) + R_p3d, t_p3d = _D @ np.asarray(R_cv, np.float64), _D @ np.asarray(t_cv, np.float64) + canon = self.adapter.canonical(s.obj_key) + R, t, sc = canonical.gt_layout_anchor_frame( + R_p3d, t_p3d, np.asarray(canon["center_cad"]), float(canon["max_extent"])) + rot6d = normalize_rot6d(matrix_to_rot6d(torch.from_numpy(R).float())).float() + + # ---- SSI-consistent pose targets (mirror dexycb_dataset exactly) ---- + shift_t = torch.from_numpy(np.asarray(shift, np.float64)).float() + pm_scale_safe = torch.tensor(float(scale), dtype=torch.float32).clamp_min(MIN_VALID_PM_SCALE) + t_ssi = (torch.from_numpy(t).float() - shift_t) / pm_scale_safe + gt_scale_m = torch.tensor(float(sc), dtype=torch.float32) + gt_scale_ratio = (gt_scale_m / pm_scale_safe).clamp_min(MIN_VALID_PM_SCALE) + s_log3 = torch.log(gt_scale_ratio).expand(3).contiguous() + + return { + "shape_latent": canon["latent"].float(), + "rot6d": rot6d, + "translation": t_ssi, + "scale": s_log3, + "pm_shift": shift_t, + "pm_scale": pm_scale_safe.reshape(1), + "views": views, + "object_name": s.obj_key, + } diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/canonical.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/canonical.py new file mode 100644 index 0000000000000000000000000000000000000000..6e0e8a594fa41ca09b2e848fddba9195989e1e10 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/canonical.py @@ -0,0 +1,76 @@ +"""Canonical-voxel construction and anchor-frame GT layout math. + +The occupancy is produced by the VENDORED ``voxelize_mesh`` (byte-identical to +SAM-3D's ``inference_utils.voxelize_mesh``), so the SS encoder sees exactly the +distribution it was trained on. ``voxelize_mesh`` operates in the ROTATED +(z-up→y-up) frame and returns that frame's ``scale = 1/max_extent`` and +``center`` (the rotated-frame AABB center). Per the locked design we recover +the CAD-frame quantities the GT layout consumes: + + max_extent = 1 / scale (rotation-invariant: M is axis-aligned) + center_cad = M^T @ center_rot (rotated-frame center → CAD frame) + +GT layout (anchor camera, PyTorch3D). A canonical voxel point ``v_canon`` maps +to the camera as ``X_cam = R @ (s · v_canon) + t``. Composing the voxelization +with the DexYCB object pose ``X_cam = R_dex @ p_cad + t_dex`` gives: + + s = max_extent (isotropic scalar) + R = R_dex @ M^T (absorbs the z-up→y-up grid rotation) + t = R_dex @ center_cad + t_dex + +with ``M = [[1,0,0],[0,0,-1],[0,1,0]]`` (the grid rotation baked into +``voxelize_mesh``) and the OpenCV→PyTorch3D flip ``D = diag(-1,-1,1)`` applied to +``(R_dex, t_dex)`` upstream in the dataset. +""" +from __future__ import annotations + +import numpy as np + +from sam3d_objects.pipeline.inference_utils import trimesh2o3d_mesh, voxelize_mesh + +# Fixed z-up → y-up rotation applied inside voxelize_mesh: +# verts are right-multiplied by M.T, i.e. v_rot = v_cad @ M.T <=> v_rot = M @ v_cad. +M_ZUP_TO_YUP = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]], dtype=np.float64) + + +def voxelize_canonical(mesh, resolution: int = 64): + """CAD trimesh → (occupancy ``[1,R,R,R]`` long, max_extent float, center_cad (3,)). + + ``occupancy`` is exactly the vendored ``voxelize_mesh`` output. ``max_extent`` + and ``center_cad`` are recovered in the ORIGINAL CAD frame from the rotated- + frame ``scale``/``center`` ``voxelize_mesh`` returns, and are what the GT + layout math consumes. + """ + # voxelize_mesh consumes an open3d TriangleMesh (it assigns + # o3d.utility.Vector3dVector to mesh.vertices and calls the o3d voxelizer), + # so convert via the vendored trimesh2o3d_mesh first — exactly as the fork's + # pipeline does. trimesh2o3d_mesh builds a fresh o3d mesh, so the input + # trimesh is not mutated. + o3d_mesh = trimesh2o3d_mesh(mesh) + ss, scale, center_rot = voxelize_mesh(o3d_mesh, resolution=resolution) + if scale is None or center_rot is None: + # Mesh was already normalized to [-0.5, 0.5]; recover stats from raw verts. + verts = np.asarray(mesh.vertices, dtype=np.float64) + verts = verts @ M_ZUP_TO_YUP.T # voxelize_mesh works in the rotated frame + vmin, vmax = verts.min(axis=0), verts.max(axis=0) + center_rot = (vmax + vmin) / 2.0 + max_extent = float(np.max(vmax - vmin)) + else: + max_extent = 1.0 / float(scale) + center_rot = np.asarray(center_rot, dtype=np.float64) + # center_rot is in the rotated (y-up) frame; bring it back to the CAD frame. + center_cad = M_ZUP_TO_YUP.T @ center_rot + return ss, max_extent, center_cad.astype(np.float64) + + +def gt_layout_anchor_frame(R_dex: np.ndarray, t_dex: np.ndarray, + center_cad: np.ndarray, max_extent: float): + """DexYCB object pose (anchor cam frame) → canonical (R, t, s). + + R_dex (3,3), t_dex (3,): object→camera pose in the anchor camera frame + (already in the desired PyTorch3D convention). Returns R (3,3), t (3,), + s (float). + """ + R = R_dex @ M_ZUP_TO_YUP.T + t = R_dex @ np.asarray(center_cad, np.float64) + np.asarray(t_dex, np.float64) + return R.astype(np.float64), t.astype(np.float64), float(max_extent) diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/dexycb_calib.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/dexycb_calib.py new file mode 100644 index 0000000000000000000000000000000000000000..bb31c41fad30838f983ec88961c5a5835987fa32 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/dexycb_calib.py @@ -0,0 +1,75 @@ +"""DexYCB camera calibration: per-camera extrinsics/intrinsics and the +cross-camera transform that re-expresses an aux camera's pointmap and pose in +the ANCHOR camera frame. + +DexYCB stores, per extrinsics group, a 3x4 ``[R_c | t_c]`` (row-major, 12 +floats) mapping the shared calibration/tag frame -> camera ``c`` in the +**OpenCV** convention (x-right, y-down, z-fwd). Exactly one camera is identity +(the master/tag). The whole training pipeline works in **PyTorch3D** frame +(x-left, y-up, z-fwd), related by ``D = diag(-1,-1,1)`` (``D == D^{-1}``). + +For a 3D point in aux cam B (PyTorch3D) the anchor-frame point is: + + X_tag = R_B^T (D X_B^p3d - t_B) + X_A^p3d = D (R_A X_tag + t_A) + = (D R_A R_B^T D) X_B^p3d + (D t_A - D R_A R_B^T t_B) + +i.e. a rigid transform ``X_A = Rrel @ X_B + trel`` with +``Rrel = D R_A R_B^T D`` and ``trel = D t_A - D R_A R_B^T t_B``. +""" +from __future__ import annotations + +import os + +import numpy as np +import yaml + +_D = np.diag([-1.0, -1.0, 1.0]).astype(np.float64) + + +def _load_yaml(path): + with open(path) as f: + return yaml.load(f, Loader=yaml.UnsafeLoader) # files contain python/tuple tags + + +def load_extrinsics(calib_root: str, extrinsics_tag: str) -> dict[str, dict]: + """serial -> {"R": (3,3), "t": (3,)} (tag->cam, OpenCV).""" + path = os.path.join(calib_root, f"extrinsics_{extrinsics_tag}", "extrinsics.yml") + raw = _load_yaml(path)["extrinsics"] + out = {} + for serial, vals in raw.items(): + m = np.asarray(vals, dtype=np.float64).reshape(3, 4) + out[str(serial)] = {"R": m[:, :3], "t": m[:, 3]} + return out + + +def load_intrinsics(calib_root: str, serial: str) -> dict: + """color intrinsics fx,fy,ppx,ppy for a serial (640x480).""" + path = os.path.join(calib_root, "intrinsics", f"{serial}_640x480.yml") + return _load_yaml(path)["color"] + + +def relative_p3d(ext: dict, anchor_serial: str, aux_serial: str): + """Return (Rrel, trel) mapping a PyTorch3D point in ``aux`` cam -> anchor cam. + + X_anchor = Rrel @ X_aux + trel. + """ + Ra, ta = ext[anchor_serial]["R"], ext[anchor_serial]["t"] + Rb, tb = ext[aux_serial]["R"], ext[aux_serial]["t"] + RaRbT = Ra @ Rb.T + Rrel = _D @ RaRbT @ _D + trel = _D @ ta - _D @ (RaRbT @ tb) + return Rrel.astype(np.float64), trel.astype(np.float64) + + +def transform_pointmap_to_anchor(pm: np.ndarray, Rrel: np.ndarray, trel: np.ndarray) -> np.ndarray: + """Apply the rigid transform to an (H,W,3) pointmap, preserving NaN/invalid. + + pm is in the aux camera's PyTorch3D frame; returns it in the anchor frame. + """ + H, W, _ = pm.shape + flat = pm.reshape(-1, 3) + valid = np.isfinite(flat).all(axis=1) + out = flat.copy() + out[valid] = flat[valid] @ Rrel.T + trel + return out.reshape(H, W, 3) diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/dexycb_dataset.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/dexycb_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..0e81afe09e3530905fb9f4367deb2041eb592e66 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/dexycb_dataset.py @@ -0,0 +1,501 @@ +"""Multi-view (multi-camera) DexYCB dataset for the 6D-pose flow model. + +One sample = one (sequence, timestamp, object, anchor-camera). The anchor +camera defines the reference frame; the other synchronized cameras are aux +views. Everything is expressed in the **anchor camera's PyTorch3D frame**: + + * Pointmaps are unprojected from each camera's aligned depth (metric, OpenCV), + converted to PyTorch3D (X' = diag(-1,-1,1) X), and aux cameras are rigidly + transformed into the anchor frame via DexYCB calibration extrinsics BEFORE a + single shared SSI (scale-shift invariant) normalization computed on the + anchor view. Pointmaps deliberately KEEP NaN at invalid/background pixels. + * GT (R, t, s) is the anchor camera's object pose, OpenCV→PyTorch3D, mapped to + the canonical voxel frame (R_dex·Mᵀ, R_dex·c_cad + t_dex, max_extent) — see + canonical.gt_layout_anchor_frame. + * The 6D rotation is ``normalize_rot6d(matrix_to_rot6d(R))`` using the + VENDORED column-convention rotation utils (consistent with the model's + decode ``rot6d_to_matrix(denormalize_rot6d(·))``). + * The shape target is the precomputed canonical SS latent for the object + (cache/canonical/.pt). + +There is NO shape-K/V; this dataset only emits plain conditioning tokens + flow +targets. Output dict feeds MVSam3DFlow.flow_step directly. +""" +from __future__ import annotations + +import glob +import os +import random +from dataclasses import dataclass +from typing import Dict, List, Optional + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +from torch.utils.data import Dataset + +from sam3d_objects.geometry.geometry import ( + matrix_to_rot6d, + normalize_rot6d, +) + +from . import canonical, dexycb_calib + +_D = np.diag([-1.0, -1.0, 1.0]).astype(np.float64) +DEPTH_SCALE = 1000.0 # DexYCB depth PNG is uint16 millimetres +IMG_SIZE = 518 +DEPTH_NEAR, DEPTH_FAR = 0.1, 3.0 # metres; clip implausible depth + +# DexYCB ycb_ids / seg use 1-based class indices (NOT the YCB numeric prefix). +# This is the canonical DexYCB _YCB_CLASSES ordering; seg == ycb_id for the object. +YCB_ID_TO_NAME = { + 1: "002_master_chef_can", 2: "003_cracker_box", 3: "004_sugar_box", + 4: "005_tomato_soup_can", 5: "006_mustard_bottle", 6: "007_tuna_fish_can", + 7: "008_pudding_box", 8: "009_gelatin_box", 9: "010_potted_meat_can", + 10: "011_banana", 11: "019_pitcher_base", 12: "021_bleach_cleanser", + 13: "024_bowl", 14: "025_mug", 15: "035_power_drill", 16: "036_wood_block", + 17: "037_scissors", 18: "040_large_marker", 19: "051_large_clamp", + 20: "052_extra_large_clamp", 21: "061_foam_brick", +} + + +@dataclass +class Sample: + seq_dir: str + frame: int + obj_id: int + obj_index: int # index into the sequence's ycb_ids / pose_y + anchor_serial: str + aux_serials: List[str] # FULL pool of non-anchor cameras; subset chosen per __getitem__ + extrinsics_tag: str + + +# --------------------------------------------------------------------------- # +# geometry helpers +# --------------------------------------------------------------------------- # +def _unproject(depth_m: np.ndarray, intr: dict, + depth_near: float = DEPTH_NEAR, depth_far: float = DEPTH_FAR) -> np.ndarray: + """aligned depth (H,W) metres → (H,W,3) pointmap in OpenCV cam frame, NaN invalid.""" + H, W = depth_m.shape + fx, fy, cx, cy = intr["fx"], intr["fy"], intr["ppx"], intr["ppy"] + us, vs = np.meshgrid(np.arange(W), np.arange(H)) + z = depth_m + x = (us - cx) / fx * z + y = (vs - cy) / fy * z + pm = np.stack([x, y, z], axis=-1).astype(np.float64) + invalid = (z <= depth_near) | (z >= depth_far) + pm[invalid] = np.nan + return pm + + +def _to_p3d(pm_cv: np.ndarray) -> np.ndarray: + """OpenCV pointmap → PyTorch3D (flip x,y).""" + return pm_cv * np.array([-1.0, -1.0, 1.0]) + + +from mvsam3d import paths as _mv_paths +_FORK_ROOT = _mv_paths.FORK_ROOT + + +def _load_object_centric_ssi(): + """[raw SAM-3D] import the vendored ObjectCentricSSI from the fork. + + This repo vendors a PARTIAL ``sam3d_objects`` subset that shadows the fork's + full package on sys.path, so ``sam3d_objects.data.dataset`` fails to import. + Fix: extend the partial packages' ``__path__`` with the fork's directories so + the missing subpackages resolve there (pure reuse — never reimplement). + Additive only: names that already resolve keep resolving to this repo.""" + try: + from sam3d_objects.data.dataset.tdfy.img_and_mask_transforms import ObjectCentricSSI + return ObjectCentricSSI + except ModuleNotFoundError: + import sam3d_objects + import sam3d_objects.data + for pkg, sub in ((sam3d_objects, "sam3d_objects"), + (sam3d_objects.data, os.path.join("sam3d_objects", "data"))): + p = os.path.join(_FORK_ROOT, sub) + if p not in pkg.__path__: + pkg.__path__.append(p) + from sam3d_objects.data.dataset.tdfy.img_and_mask_transforms import ObjectCentricSSI + return ObjectCentricSSI + + +def _ssi_params(pm: np.ndarray, mask: np.ndarray | None = None, mode: str = "scene"): + """SSI params from a pointmap. Two selectable conventions (``mode``): + + * ``"scene"`` (default, used by all pre-2026-07-17 checkpoints): + shift=[0,0,median_z], scale=mean|p-shift| over the whole scene. Matches the + fork's BASE ``ScaleShiftInvariant.get_scale_and_shift`` default + (NOT ObjectCentricSSI). + * ``"object_centric"``: [raw SAM-3D] vendored ``ObjectCentricSSI`` + (use_scene_scale=True) — shift = 3D median of the points INSIDE the object + ``mask`` (object recentred to the origin), scale = scene scale. This is the + convention the release HF checkpoint's inference pipeline.yaml pairs with + the frozen condition encoder. Requires ``mask``. + + A checkpoint must always be evaluated with the mode it was trained with — + t/s targets and the cam features live in this SSI space. + """ + if mode == "object_centric": + if mask is None: + raise ValueError("_ssi_params(mode='object_centric') requires the object mask") + ObjectCentricSSI = _load_object_centric_ssi() + pm_t = torch.from_numpy(np.ascontiguousarray(pm)).float().permute(2, 0, 1) # [3,H,W] + mask_t = torch.from_numpy(np.ascontiguousarray(mask)).float()[None] # [1,H,W] + scale_t, shift_t = ObjectCentricSSI(use_scene_scale=True)._compute_scale_and_shift(pm_t, mask_t) + shift = shift_t.numpy().astype(np.float64) # [3] + scale = max(float(scale_t.reshape(-1)[0]), 1e-6) # isotropic + return shift, scale + if mode != "scene": + raise ValueError(f"unknown ssi mode: {mode!r} (expected 'scene' or 'object_centric')") + flat = pm.reshape(-1, 3) + valid = np.isfinite(flat).all(axis=1) + if valid.sum() < 16: + return np.zeros(3), 1.0 + pts = flat[valid] + shift = np.array([0.0, 0.0, np.median(pts[:, 2])]) + scale = float(np.mean(np.abs(pts - shift))) + scale = max(scale, 1e-6) + return shift, scale + + +def _apply_ssi(pm: np.ndarray, shift, scale) -> np.ndarray: + return (pm - shift) / scale + + +# Per-view camera feature width fed to the model's mlp_cam: rot6d(6) + camera +# centre in the anchor frame, SSI-normalized (3). +CAM_FEAT_DIM = 9 + + +def _cam_feature(Rrel: np.ndarray, trel: np.ndarray, shift, scale) -> torch.Tensor: + """Pack a view's camera->anchor extrinsics into a 9-vector for ``mlp_cam``. + + ``[0:6]`` = column-convention rot6d of ``Rrel`` (the aux->anchor rotation), + using the SAME vendored ``matrix_to_rot6d`` as the object-rotation target so + the convention is consistent (NOT ``normalize_rot6d`` — that's tuned to the + object distribution; raw rot6d columns are already unit-scaled). + ``[6:9]`` = the camera centre in the anchor frame (which is exactly ``trel``, + the image of the aux origin under ``X_anchor = Rrel X_aux + trel``), expressed + in the SAME SSI space as the conditioning pointmaps -> ``(trel - shift)/scale``. + + Computed uniformly for every view: the anchor passes ``Rrel=I, trel=0`` (its + centre normalizes to ``-shift/scale``), so there is no special-casing. + """ + rot6d = matrix_to_rot6d(torch.from_numpy(np.asarray(Rrel, np.float64)).float()) # [6] + centre = torch.from_numpy(np.asarray(trel, np.float64)).float() # [3] + shift_t = torch.from_numpy(np.asarray(shift, np.float64)).float() # [3] + pos = (centre - shift_t) / float(scale) # [3] + return torch.cat([rot6d, pos]) # [9] + + +def _mask_bbox(mask: np.ndarray, pad: float = 0.1): + """Square bbox around mask with padding → (r0,c0,r1,c1). Full frame if empty.""" + H, W = mask.shape + ys, xs = np.where(mask > 0) + if ys.size == 0: + return 0, 0, H, W + r0, r1 = ys.min(), ys.max() + 1 + c0, c1 = xs.min(), xs.max() + 1 + h, w = r1 - r0, c1 - c0 + side = max(h, w) * (1.0 + pad) + cy, cx = (r0 + r1) / 2.0, (c0 + c1) / 2.0 + r0 = int(round(cy - side / 2)); r1 = int(round(cy + side / 2)) + c0 = int(round(cx - side / 2)); c1 = int(round(cx + side / 2)) + return r0, c0, r1, c1 + + +def _crop_resize(arr: torch.Tensor, box, size: int, mode: str, fill: float): + """arr [C,H,W] → [C,size,size]; crop to (possibly out-of-bounds) box then resize. + + Out-of-bounds region is filled with `fill` (NaN-safe for pointmaps). Nearest + interpolation is used for pointmaps/masks so NaN is preserved (never averaged). + """ + C, H, W = arr.shape + r0, c0, r1, c1 = box + bh, bw = r1 - r0, c1 - c0 + canvas = arr.new_full((C, bh, bw), fill) + sr0, sc0 = max(r0, 0), max(c0, 0) + sr1, sc1 = min(r1, H), min(c1, W) + if sr1 > sr0 and sc1 > sc0: + canvas[:, sr0 - r0:sr1 - r0, sc0 - c0:sc1 - c0] = arr[:, sr0:sr1, sc0:sc1] + interp = {"bilinear": "bilinear", "nearest": "nearest"}[mode] + kw = {"align_corners": False} if interp == "bilinear" else {} + out = F.interpolate(canvas[None], size=(size, size), mode=interp, **kw)[0] + return out + + +def _full_resize(arr: torch.Tensor, size: int, mode: str, fill: float): + """Pad [C,H,W] to square (centered) then resize to size.""" + C, H, W = arr.shape + s = max(H, W) + canvas = arr.new_full((C, s, s), fill) + r0, c0 = (s - H) // 2, (s - W) // 2 + canvas[:, r0:r0 + H, c0:c0 + W] = arr + interp = {"bilinear": "bilinear", "nearest": "nearest"}[mode] + kw = {"align_corners": False} if interp == "bilinear" else {} + return F.interpolate(canvas[None], size=(size, size), mode=interp, **kw)[0] + + +# --------------------------------------------------------------------------- # +# dataset +# --------------------------------------------------------------------------- # +def _build_view_fn(rgb, mask, pm_anchor_frame, anchor: bool, image_size: int): + """Module-level view construction (shared by MultiViewDexYCB._build_view and + the dataset-agnostic MultiViewPoseDataset). Byte-identical logic.""" + img_t = torch.from_numpy(rgb).permute(2, 0, 1).float() + msk_t = torch.from_numpy(mask)[None].float() + pm_t = torch.from_numpy(pm_anchor_frame).permute(2, 0, 1).float() + box = _mask_bbox(mask) + out = { + "image_cropped": _crop_resize(img_t, box, image_size, "bilinear", 0.0), + "mask_cropped": _crop_resize(msk_t, box, image_size, "nearest", 0.0), + "pointmap_cropped": _crop_resize(pm_t, box, image_size, "nearest", float("nan")), + } + if anchor: + out["image_full"] = _full_resize(img_t, image_size, "bilinear", 0.0) + out["mask_full"] = _full_resize(msk_t, image_size, "nearest", 0.0) + out["pointmap_full"] = _full_resize(pm_t, image_size, "nearest", float("nan")) + return out + + +class MultiViewDexYCB(Dataset): + def __init__( + self, + dexycb_root: str = os.path.join(_mv_paths.DATA_ROOT, "dexycb"), + canonical_cache: str = _mv_paths.CANONICAL_CACHE, + num_aux_views: int = 7, + frame_stride: int = 5, + anchor_sweep: bool = True, + subjects: Optional[List[str]] = None, + image_size: int = IMG_SIZE, + random_views: bool = False, + min_aux_views: int = 1, + max_aux_views: int = 3, + depth_scale: float = DEPTH_SCALE, + depth_near: float = DEPTH_NEAR, + depth_far: float = DEPTH_FAR, + split: str = "train", + val_fraction: float = 0.0, + split_seed: int = 0, + ssi_mode: str = "scene", + ): + """When ``random_views`` is True, each __getitem__ draws a random aux count + ``k ~ U[min_aux_views, max_aux_views]`` and a random subset of that size + from the non-anchor cameras (total views = 1 + k, i.e. 2..4 by default). + When False, the first ``num_aux_views`` aux cameras are used.""" + self.root = dexycb_root + self.calib_root = os.path.join(dexycb_root, "calibration") + self.cache = canonical_cache + self.num_aux = num_aux_views + self.image_size = image_size + self.random_views = random_views + self.min_aux = min_aux_views + self.max_aux = max_aux_views + self.depth_scale = depth_scale + self.depth_near = depth_near + self.depth_far = depth_far + # SSI convention for pointmaps + t/s targets + cam features ("scene" | + # "object_centric"); validated eagerly so a config typo fails at startup. + if ssi_mode not in ("scene", "object_centric"): + raise ValueError(f"ssi_mode must be 'scene' or 'object_centric', got {ssi_mode!r}") + self.ssi_mode = ssi_mode + if random_views and not (1 <= min_aux_views <= max_aux_views): + raise ValueError( + f"random_views requires 1 <= min_aux_views ({min_aux_views}) " + f"<= max_aux_views ({max_aux_views})" + ) + self._canon: Dict[str, dict] = {} + + seq_dirs = sorted(glob.glob(os.path.join(dexycb_root, "2020*-subject-*", "2020*"))) + if subjects is not None: + seq_dirs = [d for d in seq_dirs if any(s in d for s in subjects)] + + self.samples: List[Sample] = [] + for seq_dir in seq_dirs: + meta = dexycb_calib._load_yaml(os.path.join(seq_dir, "meta.yml")) + serials = [str(s) for s in meta["serials"]] + ycb_ids = list(meta["ycb_ids"]) + tag = str(meta["extrinsics"]) + n = int(meta["num_frames"]) + anchors = serials if anchor_sweep else serials[:1] + for frame in range(0, n, frame_stride): + for anchor in anchors: + aux = [s for s in serials if s != anchor] # full pool + for oi, oid in enumerate(ycb_ids): + if int(oid) not in YCB_ID_TO_NAME: + continue + self.samples.append( + Sample(seq_dir, frame, int(oid), oi, anchor, aux, tag) + ) + + # ---- deterministic train/val split ------------------------------- + # A fixed-seed shuffle partitions the full sample list once; the same + # (val_fraction, split_seed) yields the same disjoint train/val sets, so + # validation never touches a training sample. val_fraction=0 keeps all + # samples in "train" (no split). + if val_fraction > 0.0: + if not 0.0 < val_fraction < 1.0: + raise ValueError(f"val_fraction must be in (0,1), got {val_fraction}") + order = list(range(len(self.samples))) + random.Random(split_seed).shuffle(order) + n_val = max(1, int(round(len(order) * val_fraction))) + val_idx = set(order[:n_val]) + if split == "val": + self.samples = [self.samples[i] for i in sorted(val_idx)] + elif split == "train": + self.samples = [self.samples[i] for i in range(len(self.samples)) if i not in val_idx] + else: + raise ValueError(f"split must be 'train' or 'val', got {split!r}") + + def __len__(self): + return len(self.samples) + + # ---- per-object canonical cache ---- + def _canonical(self, obj_id: int) -> dict: + name = YCB_ID_TO_NAME[obj_id] + if name not in self._canon: + d = torch.load(os.path.join(self.cache, f"{name}.pt"), map_location="cpu", + weights_only=False) + self._canon[name] = d + return self._canon[name] + + # ---- raw IO ---- + def _load_color(self, seq_dir, serial, frame) -> np.ndarray: + p = os.path.join(seq_dir, serial, f"color_{frame:06d}.jpg") + return np.asarray(Image.open(p).convert("RGB"), dtype=np.float32) / 255.0 + + def _load_depth(self, seq_dir, serial, frame) -> np.ndarray: + p = os.path.join(seq_dir, serial, f"aligned_depth_to_color_{frame:06d}.png") + return np.asarray(Image.open(p), dtype=np.float32) / self.depth_scale + + def _load_seg(self, seq_dir, serial, frame, obj_id) -> np.ndarray: + p = os.path.join(seq_dir, serial, f"labels_{frame:06d}.npz") + seg = np.load(p)["seg"] + return (seg == obj_id).astype(np.float32) + + def _pose_y(self, seq_dir, serial, frame, obj_index): + p = os.path.join(seq_dir, serial, f"labels_{frame:06d}.npz") + py = np.load(p)["pose_y"][obj_index] # (3,4) OpenCV CAD->cam + return py[:, :3].astype(np.float64), py[:, 3].astype(np.float64) + + # ---- view tensor construction ---- + def _build_view(self, rgb, mask, pm_anchor_frame, anchor: bool): + """rgb [H,W,3] in [0,1]; mask [H,W]; pm [H,W,3] (anchor frame, already SSI). + + Returns a dict {image_cropped, mask_cropped, pointmap_cropped, [+ *_full]}. + """ + # Delegates to the module-level _build_view_fn (byte-identical). Pointmaps + # intentionally KEEP NaN at invalid/background pixels so the pretrained + # PointPatchEmbed's isfinite() valid_mask + invalid_xyz_token path works. + return _build_view_fn(rgb, mask, pm_anchor_frame, anchor, self.image_size) + + def __getitem__(self, idx: int) -> dict: + s = self.samples[idx] + ext = dexycb_calib.load_extrinsics(self.calib_root, s.extrinsics_tag) + + # ---- anchor view + shared SSI from anchor pointmap ---- + a_rgb = self._load_color(s.seq_dir, s.anchor_serial, s.frame) + a_mask = self._load_seg(s.seq_dir, s.anchor_serial, s.frame, s.obj_id) + a_intr = dexycb_calib.load_intrinsics(self.calib_root, s.anchor_serial) + a_pm = _to_p3d(_unproject(self._load_depth(s.seq_dir, s.anchor_serial, s.frame), a_intr, + self.depth_near, self.depth_far)) + shift, scale = _ssi_params(a_pm, mask=a_mask, mode=self.ssi_mode) + + views = [self._build_view(a_rgb, a_mask, _apply_ssi(a_pm, shift, scale), anchor=True)] + # Anchor camera == the reference frame: relative transform is identity. + views[0]["cam"] = _cam_feature(np.eye(3), np.zeros(3), shift, scale) + + # ---- choose which aux cameras to use (random count + subset, or fixed) ---- + pool = s.aux_serials + if self.random_views: + k = random.randint(self.min_aux, min(self.max_aux, len(pool))) + aux_serials = random.sample(pool, k) + else: + aux_serials = pool[: self.num_aux] + + # ---- aux views: transform pointmap into anchor frame, share SSI ---- + for aux in aux_serials: + x_rgb = self._load_color(s.seq_dir, aux, s.frame) + x_mask = self._load_seg(s.seq_dir, aux, s.frame, s.obj_id) + x_intr = dexycb_calib.load_intrinsics(self.calib_root, aux) + x_pm = _to_p3d(_unproject(self._load_depth(s.seq_dir, aux, s.frame), x_intr, + self.depth_near, self.depth_far)) + Rrel, trel = dexycb_calib.relative_p3d(ext, s.anchor_serial, aux) + x_pm = dexycb_calib.transform_pointmap_to_anchor(x_pm, Rrel, trel) + av = self._build_view(x_rgb, x_mask, _apply_ssi(x_pm, shift, scale), anchor=False) + # Bind this aux view to its camera pose (aux->anchor extrinsics). + av["cam"] = _cam_feature(Rrel, trel, shift, scale) + views.append(av) + + # ---- GT layout (anchor cam, PyTorch3D) ---- + R_cv, t_cv = self._pose_y(s.seq_dir, s.anchor_serial, s.frame, s.obj_index) + R_p3d, t_p3d = _D @ R_cv, _D @ t_cv # OpenCV → PyTorch3D + canon = self._canonical(s.obj_id) + R, t, sc = canonical.gt_layout_anchor_frame( + R_p3d, t_p3d, np.asarray(canon["center_cad"]), float(canon["max_extent"]) + ) + # 6D rotation: VENDORED column-convention encode, normalized to the + # pretrained backbone's distribution. Consistent with the model's decode + # rot6d_to_matrix(denormalize_rot6d(·)). + R_t = torch.from_numpy(R).float() + rot6d = normalize_rot6d(matrix_to_rot6d(R_t)).float() # [6] + + # ---- SSI-consistent pose targets ---------------------------------- + # Mirror the fork stage1 model_wrapper._build_pose_latents EXACTLY + # (sam3d-for-6d stage1 model_wrapper.py lines 560-581); + # the conditioning pointmaps were normalized by the shared anchor SSI + # (shift [3], scale scalar), so the t/s targets share that SSI space. + # Guards use the same constants as the fork: + # min_valid_pointmap_scale = 1e-6 (clamps pm_scale and the scale ratio) + # No mean/std on t or s (fork normalize=False; NOT the ScaleShiftInvariant + # pose-target/inference decoder). 3-channel isotropic scale latent. + MIN_VALID_PM_SCALE = 1e-6 + shift_t = torch.from_numpy(np.asarray(shift, dtype=np.float64)).float() # [3] + # fork line 560: pm_scale_safe = pm_scale.float().clamp_min(min_valid_pointmap_scale) + pm_scale_safe = torch.tensor(float(scale), dtype=torch.float32).clamp_min(MIN_VALID_PM_SCALE) + # fork line 561: gt_trans_ssi = (gt_trans_metric - pm_shift) / pm_scale_safe + t_ssi = (torch.from_numpy(t).float() - shift_t) / pm_scale_safe # [3] + # fork lines 580-581: gt_scale_ratio = (gt_scale_m / pm_scale_scalar).clamp_min(min_valid_pointmap_scale) + # gt_scale_log = log(gt_scale_ratio).expand(...,3) + gt_scale_m = torch.tensor(float(sc), dtype=torch.float32) # scalar metric extent + gt_scale_ratio = (gt_scale_m / pm_scale_safe).clamp_min(MIN_VALID_PM_SCALE) + s_log3 = torch.log(gt_scale_ratio).expand(3).contiguous() # [3] isotropic + + return { + "shape_latent": canon["latent"].float(), # [8,16,16,16] + "rot6d": rot6d, # [6] + "translation": t_ssi, # [3] SSI space + "scale": s_log3, # [3] log-ratio, isotropic + "pm_shift": shift_t, # [3] SSI shift + "pm_scale": pm_scale_safe.reshape(1), # [1] SSI scale (clamped, fork-safe) + "views": views, + "object_name": YCB_ID_TO_NAME[s.obj_id], + } + + +def collate_multiview(batch: List[dict]) -> dict: + """Stack a list of samples. ``views`` becomes a list (len=n_views) of dicts of + batched [B,C,H,W] tensors; the rest stack on dim 0. + + Samples may carry different view counts (random-view training). All views at + a given position share the same group structure (position 0 = anchor with 6 + groups, positions >=1 = aux with 3 cropped groups), so the batch is truncated + to the minimum view count present. At batch_size=1 this is a no-op.""" + n_views = min(len(b["views"]) for b in batch) + out = { + "shape_latent": torch.stack([b["shape_latent"] for b in batch]), + "rot6d": torch.stack([b["rot6d"] for b in batch]), + "translation": torch.stack([b["translation"] for b in batch]), + "scale": torch.stack([b["scale"] for b in batch]), + "pm_shift": torch.stack([b["pm_shift"] for b in batch]), + "pm_scale": torch.stack([b["pm_scale"] for b in batch]), + "object_name": [b["object_name"] for b in batch], + } + views = [] + for v in range(n_views): + keys = batch[0]["views"][v].keys() + views.append({k: torch.stack([b["views"][v][k] for b in batch]) for k in keys}) + out["views"] = views + return out diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/integrated.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/integrated.py new file mode 100644 index 0000000000000000000000000000000000000000..0c3d8dc6761ce15ee6c2981f485dbd11d91e0dfd --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/data/integrated.py @@ -0,0 +1,120 @@ +"""Integrated multi-dataset training set builder + held-out enforcement. + +Reads a ``data.datasets`` list from the training config, instantiates one +MultiViewPoseDataset per entry (via its adapter), enforces the held-out manifest +(configs/heldout.yaml), concatenates, and provides a diversity-weighted sampler. + +HELD-OUT SEMANTICS (two distinct kinds — the distinction is load-bearing): + * OOD datasets (ho3d, ycbm, hope): whole-dataset held-out for evaluation. Their + objects MAY overlap training objects (e.g. HO-3D is YCB). We do NOT enforce + object-disjointness for these — doing so would delete every YCB object from + DexYCB/YCB-V training. They are simply never built with split:train. + * unseen-object datasets (housecat6d, graspnet test_novel): held out AND their + object identities reserved — removed from EVERY train dataset (global + object-disjointness). Their objects are non-YCB, so this is safe. +""" +from __future__ import annotations + +import os +from typing import Dict, List + +import numpy as np +import torch +import yaml +from torch.utils.data import ConcatDataset, WeightedRandomSampler + +from .adapters import ADAPTERS +from .base_dataset import MultiViewPoseDataset + + +import inspect + + +def _adapter_kwargs(adapter_cls, entry: dict) -> dict: + """Pass through only the entry keys the adapter's __init__ actually accepts. + `role` is ROUTING (train vs heldout) and is never passed to the adapter; + `data_split` maps to the adapter's on-disk `split` (test/train/val).""" + accepted = set(inspect.signature(adapter_cls.__init__).parameters) - {"self", "root", "canonical_cache"} + kw = {k: entry[k] for k in accepted if k in entry and k not in ("role", "data_split")} + if "split" in accepted and "data_split" in entry: + kw["split"] = entry["data_split"] + return kw + + +def load_heldout(path: str | None) -> dict: + if not path or not os.path.exists(path): + return {"ood_datasets": [], "unseen_object_datasets": [], + "unseen_object_keys_extra": [], "never_train": []} + with open(path) as f: + return yaml.safe_load(f) or {} + + +def reserved_object_keys(heldout: dict, canonical_cache: str) -> set: + """Union of obj_keys from unseen-object datasets + explicit extras. + These keys are removed from ALL training datasets.""" + keys = set(heldout.get("unseen_object_keys_extra", []) or []) + for name in heldout.get("unseen_object_datasets", []) or []: + ds_root = heldout["roots"][name] + ad = ADAPTERS[name](ds_root, canonical_cache, **heldout.get("adapter_kwargs", {}).get(name, {})) + for s in ad.list_samples(): + keys.add(s.obj_key) + return keys + + +def build_integrated_dataset(data_cfg: dict, split: str, *, image_size: int, + random_views: bool, min_aux: int, max_aux: int, + num_aux: int, depth_near: float, depth_far: float, + ssi_mode: str = "scene", + heldout_path: str | None = None): + """Build a ConcatDataset over the configured datasets for `split`, applying + held-out rules. Returns (dataset, per_sample_weights or None).""" + cache = data_cfg["canonical_cache"] + heldout = load_heldout(heldout_path) + never_train = set(heldout.get("never_train", []) or []) + reserved = reserved_object_keys(heldout, cache) if split == "train" else set() + + parts: List[MultiViewPoseDataset] = [] + weights: List[float] = [] + report = [] + for entry in data_cfg["datasets"]: + name = entry["name"] + adapter_name = entry.get("adapter", name) + role = entry.get("role", "train") # routing: "train" | "heldout" + # only build datasets whose role matches the request + if split == "train": + if role != "train" or name in never_train: + continue + else: # val / held-out + if role == "train": + continue + ad = ADAPTERS[adapter_name](entry["root"], cache, **_adapter_kwargs(ADAPTERS[adapter_name], entry)) + ds = MultiViewPoseDataset(ad, image_size=image_size, + random_views=random_views if split == "train" else False, + num_aux_views=num_aux, min_aux_views=min_aux, + max_aux_views=max_aux, depth_near=depth_near, depth_far=depth_far, + ssi_mode=ssi_mode) + # object-disjointness filter (train only) + if split == "train" and reserved: + kept = [i for i, s in enumerate(ds.samples) if s.obj_key not in reserved] + ds.samples = [ds.samples[i] for i in kept] + if len(ds) == 0: + continue + parts.append(ds) + w = float(entry.get("weight", 1.0)) + weights.extend([w / len(ds)] * len(ds)) # diversity weight: per-dataset mixing ratio + report.append(f"{name}:{len(ds)}(w={w})") + + if not parts: + raise RuntimeError(f"No datasets for split={split}") + concat = ConcatDataset(parts) + print(f"[integrated] {split}: " + " ".join(report) + f" -> total {len(concat)}", flush=True) + sample_weights = torch.tensor(weights, dtype=torch.double) if split == "train" else None + return concat, sample_weights + + +def make_weighted_sampler(sample_weights, *, world_size: int, rank: int, seed: int = 0): + """Diversity-weighted sampler. For DDP, each rank draws an independent + weighted stream (seeded by rank) — the mixing ratio is preserved per rank.""" + g = torch.Generator().manual_seed(seed + rank) + return WeightedRandomSampler(sample_weights, num_samples=len(sample_weights), + replacement=True, generator=g) diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/infer/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/infer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7ae400b756187aa596fd9058b4e29e8d839ae887 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/infer/__init__.py @@ -0,0 +1,17 @@ +"""Inference package for the multi-view 6D-pose flow model (v2). + +Reusable, data-source-decoupled inference that mirrors the EXACT inference path +used by ``mvsam3d/train/validate.run_validation``: + + * build the unwrapped ``MVSam3DFlow`` with the pretrained warm-start so the + FROZEN condition fuser + warm-started backbone are reconstructed, THEN + * restore the trained tensors (backbone + mlp_cam) with the vendored + ``load_checkpoint`` (v2 checkpoints store ONLY trainable params), THEN + * run ``inner.predict(batch, num_steps=...)`` under ``eval()`` + AMP autocast + + ``no_grad`` — byte-for-byte the same sampler/decoder validation uses. +""" +from .predict import ( # noqa: F401 + build_inference_model, + load_config, + predict_pose, +) diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/infer/predict.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/infer/predict.py new file mode 100644 index 0000000000000000000000000000000000000000..edc73e07d4ec0ebbc0e80bab755fe5ea3657467f --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/infer/predict.py @@ -0,0 +1,373 @@ +"""Reusable eval-time inference for the multi-view 6D-pose flow model (v2). + +This module is the single source of truth for "load a trained v2 checkpoint and +predict (R, t, s, shape) for a batch". It reproduces the inference path of +``mvsam3d/train/validate.run_validation`` EXACTLY: + + * model construction mirrors ``mvsam3d/train/train.main`` — + ``MVSam3DFlow(... pretrained_ckpt= ...)`` first (which + reconstructs the FROZEN condition fuser + warm-starts the backbone), then the + vendored ``load_checkpoint`` overwrites ONLY the trained tensors (backbone + + mlp_cam) on top of that warm-start. v2 checkpoints under cache/ckpt store + trainable params only, so the warm-start MUST happen first or the frozen + condition pathway would be left uninitialized. + * prediction mirrors ``run_validation``: ``inner.eval()``, the SAME AMP + autocast (bf16 on cuda by default), ``torch.no_grad()``, and a single + ``inner.predict(batch, num_steps=...)`` call — same vendored Euler sampler, + same ``num_steps`` semantics, same ``decode_pose_*`` helpers. + +The pose-prediction entrypoint ``predict_pose(model, batch, num_steps)`` is +deliberately DECOUPLED from any data source: it accepts an already-collated batch +dict (``views`` + ``pm_shift`` / ``pm_scale`` + ...), so an HO3D relative-pose +eval harness can import and call it with its own batches. The CLI below wires a +DexYCB batch source around it purely for the smoke test / standalone runs. +""" +from __future__ import annotations + +import argparse +import json +import os +from typing import Dict, List, Optional, Tuple + +import torch +import yaml + +from sam3d_objects.training.stage1_flow_lora.checkpointing import load_checkpoint +from mvsam3d.model.flow_model import MVSam3DFlow +from mvsam3d import paths as _paths + +DEFAULT_CONFIG = _paths.DEFAULT_CONFIG + + +# --------------------------------------------------------------------------- # +# config + small utils +# --------------------------------------------------------------------------- # +def load_config(path: str) -> dict: + """Load the training YAML (same keys train.py reads).""" + with open(path) as f: + return _paths.expand(yaml.safe_load(f)) + + +def _move(obj, device): + """Recursively move tensors in a (possibly nested) batch to ``device``.""" + if torch.is_tensor(obj): + return obj.to(device, non_blocking=True) + if isinstance(obj, dict): + return {k: _move(v, device) for k, v in obj.items()} + if isinstance(obj, list): + return [_move(v, device) for v in obj] + return obj + + +# --------------------------------------------------------------------------- # +# model construction + checkpoint restore (mirrors train.main order) +# --------------------------------------------------------------------------- # +def build_inference_model( + cfg: dict, + checkpoint_path: str, + device, + *, + verbose: bool = True, +) -> Tuple[MVSam3DFlow, Dict]: + """Build ``MVSam3DFlow`` with the pretrained warm-start, then restore the + trained tensors from ``checkpoint_path`` with the vendored ``load_checkpoint``. + + Returns ``(model_in_eval_mode, load_report)`` where ``load_report`` carries the + restored ``step`` plus the ``load_state_dict`` missing/unexpected key report. + + Construction order is IDENTICAL to ``train.main``: + 1. ``MVSam3DFlow(..., pretrained_ckpt=model_cfg['pretrained_ckpt'], ...)`` + — reconstructs the frozen fuser + warm-starts the backbone. + 2. vendored ``load_checkpoint(checkpoint_path, model)`` — overwrites the + trainable tensors (backbone + mlp_cam) on top of the warm-start. + """ + model_cfg = cfg["model"] + cond_cfg = cfg["condition"] + loss_cfg = cfg["loss_weights"] + data_cfg = cfg["data"] + + # max_views derivation copied verbatim from train.main (API-compat only; the + # per-view bias now comes from mlp_cam, so this is not a hard upper bound). + max_aux = (data_cfg["max_aux_views"] if data_cfg["random_views"] + else data_cfg["num_aux_views"]) + max_views = max_aux + 1 + + # (1) warm-start: pretrained_ckpt reconstructs frozen fuser + backbone torso. + model = MVSam3DFlow( + model_channels=model_cfg["model_channels"], + cond_channels=model_cfg["cond_channels"], + num_blocks=model_cfg["num_blocks"], + num_heads=model_cfg["num_heads"], + max_views=max_views, + loss_weights=loss_cfg, + sigma_min=model_cfg["sigma_min"], + time_scale=model_cfg["time_scale"], + timestep_sampler=model_cfg["timestep_sampler"], + timestep_lognorm_mean=model_cfg["timestep_lognorm_mean"], + timestep_lognorm_std=model_cfg["timestep_lognorm_std"], + pretrained_ckpt=model_cfg["pretrained_ckpt"], + condition_kwargs=cond_cfg, + ).to(device) + + # (2) restore TRAINABLES on top of the warm-start, via the vendored loader. + # We transiently wrap the model's load_state_dict to capture the + # IncompatibleKeys (missing/unexpected) that load_checkpoint discards, so we + # can report them. The actual restore is still done by the vendored + # load_checkpoint (single 1x checkpoint read; its m.load_state_dict call picks + # up the instance-level override below). + inner = model.module if hasattr(model, "module") else model + _orig_load = inner.load_state_dict + captured: Dict = {} + + def _capturing_load(state_dict, strict=False, **kw): + res = _orig_load(state_dict, strict=strict, **kw) + captured["restored_tensors"] = len(state_dict) + captured["missing_keys"] = list(getattr(res, "missing_keys", [])) + captured["unexpected_keys"] = list(getattr(res, "unexpected_keys", [])) + return res + + inner.load_state_dict = _capturing_load # instance attr shadows bound method + try: + step, history = load_checkpoint( + checkpoint_path=checkpoint_path, model=model, map_location="cpu", + ) + finally: + inner.load_state_dict = _orig_load + + model.eval() # same eval() mode validation uses for the sampler + + # Sanity: every trainable param name should be present in the restored set + # (i.e. NOT in missing_keys). missing_keys here are the frozen fuser / DINO + # tensors that the trainable-only checkpoint legitimately does not carry. + trainable_names = {n for n, p in inner.named_parameters() if p.requires_grad} + missing_trainable = sorted(trainable_names & set(captured.get("missing_keys", []))) + + report = { + "step": int(step), + "restored_tensors": captured.get("restored_tensors", 0), + "n_missing_keys": len(captured.get("missing_keys", [])), + "n_unexpected_keys": len(captured.get("unexpected_keys", [])), + "unexpected_keys": captured.get("unexpected_keys", []), + "n_trainable_params": len(trainable_names), + "missing_trainable_params": missing_trainable, + } + if verbose: + print(f"[infer] built MVSam3DFlow + warm-start from " + f"{model_cfg['pretrained_ckpt']}") + print(f"[infer] load_checkpoint({checkpoint_path}) -> step={report['step']}; " + f"restored {report['restored_tensors']} trainable tensors; " + f"{report['n_missing_keys']} keys not in ckpt (frozen fuser/DINO, " + f"expected); {report['n_unexpected_keys']} unexpected.") + if missing_trainable: + print(f"[infer][WARN] {len(missing_trainable)} TRAINABLE params were " + f"NOT restored (left at warm-start/init): {missing_trainable[:8]}" + + (" ..." if len(missing_trainable) > 8 else "")) + else: + print("[infer] all trainable params restored from the checkpoint: OK") + return model, report + + +# --------------------------------------------------------------------------- # +# the decoupled pose-prediction entrypoint +# --------------------------------------------------------------------------- # +@torch.no_grad() +def predict_pose( + model, + batch: Dict, + num_steps: int, + *, + amp_dtype: str = "bf16", + device=None, +) -> Dict[str, torch.Tensor]: + """Run the model's own rectified-flow sampler on one collated ``batch`` and + return the decoded metric-space pose + shape latent. + + Wraps ``inner.predict(batch, num_steps=...)`` EXACTLY as + ``validate.run_validation`` does: unwrap to ``inner``, ``inner.eval()``, the + same AMP autocast (bf16 on cuda by default; matches train/val ``amp_dtype``), + and ``no_grad``. Outputs are cast to float32 to match the metric space the + validation metrics are computed in (``pred[...].float()``). + + Args: + model: an ``MVSam3DFlow`` (or DDP-wrapped); unwrapped internally. + batch: an already-collated batch dict with at least ``views``, ``pm_shift``, + ``pm_scale`` (the data source is irrelevant — HO3D, DexYCB, ...). + num_steps: Euler ODE steps (same semantics as ``val_num_steps``). + amp_dtype: "bf16" | "fp16" | "float16" (autocast dtype on cuda). + + Returns dict: + rotation [B,3,3] float32 (decode_pose_rotation; orthonormal, det +1), + translation [B,3] float32 (metric, decode_pose_translation), + scale [B] float32 (metric, decode_pose_scale), + shape_latent [B,8,16,16,16] float32. + """ + inner = model.module if hasattr(model, "module") else model + was_training = inner.training + inner.eval() + if device is None: + device = next(inner.parameters()).device + batch = _move(batch, device) + + use_amp = (device.type == "cuda" + and amp_dtype.lower() in {"bf16", "fp16", "float16"}) + amp_dt = torch.bfloat16 if amp_dtype.lower() == "bf16" else torch.float16 + autocast = (torch.autocast(device_type="cuda", dtype=amp_dt) + if use_amp else torch.autocast(device_type="cpu", enabled=False)) + + try: + with autocast: + pred = inner.predict(batch, num_steps=num_steps) + out = { + "rotation": pred["rotation"].float(), # [B,3,3] + "translation": pred["translation"].float(), # [B,3] metric + "scale": pred["scale"].float(), # [B] metric + "shape_latent": pred["shape_latent"].float(), # [B,8,16,16,16] + } + finally: + if was_training: + inner.train() + return out + + +# --------------------------------------------------------------------------- # +# CLI batch source (DexYCB) — used for the smoke / standalone runs only. +# The pose entrypoint above does NOT depend on any of this. +# --------------------------------------------------------------------------- # +def _build_dexycb(data_cfg: dict, *, subjects, frame_stride, num_aux_views, + split: str = "val"): + """Construct MultiViewDexYCB with deterministic views (random_views=False), + mirroring train.build_dataset for the val split. ``subjects`` / ``frame_stride`` + / ``num_aux_views`` may be overridden for a fast smoke.""" + from mvsam3d.data.dexycb_dataset import MultiViewDexYCB + return MultiViewDexYCB( + dexycb_root=data_cfg["dexycb_root"], + canonical_cache=data_cfg["canonical_cache"], + num_aux_views=num_aux_views, + frame_stride=frame_stride, + anchor_sweep=data_cfg["anchor_sweep"], + subjects=subjects, + image_size=data_cfg["image_size"], + random_views=False, # deterministic views (like val) + min_aux_views=data_cfg["min_aux_views"], + max_aux_views=data_cfg["max_aux_views"], + depth_scale=data_cfg["depth_scale"], + depth_near=data_cfg["depth_near"], + depth_far=data_cfg["depth_far"], + split=split, + val_fraction=data_cfg["val_fraction"], + split_seed=data_cfg["split_seed"], + ssi_mode=data_cfg.get("ssi_mode", "scene"), + ) + + +def _orthonormality(R: torch.Tensor) -> Tuple[float, float]: + """(det, ||RᵀR - I||) for a [3,3] rotation, for a finite/orthonormal check.""" + det = float(torch.det(R)) + err = float((R.transpose(-1, -2) @ R - torch.eye(3, device=R.device, dtype=R.dtype)).abs().max()) + return det, err + + +def main(): + ap = argparse.ArgumentParser( + description="Predict (R,t,s,shape) for DexYCB samples and write to disk.") + ap.add_argument("--config", default=DEFAULT_CONFIG, help="training YAML config") + ap.add_argument("--checkpoint", required=True, help="v2 checkpoint (cache/ckpt/step_*.pt)") + ap.add_argument("--out", required=True, help="output .pt path (a .json sibling is also written)") + ap.add_argument("--num-steps", type=int, default=None, + help="Euler ODE steps (default: config train.val_num_steps)") + ap.add_argument("--num-samples", type=int, default=4, help="number of samples to predict") + ap.add_argument("--idx", type=int, nargs="*", default=None, + help="explicit sample indices (overrides --num-samples)") + ap.add_argument("--subjects", type=str, nargs="*", default=None, + help="subject substrings to filter (e.g. subject-01); default = config") + ap.add_argument("--frame-stride", type=int, default=None, + help="override data.frame_stride (larger = fewer/faster samples)") + ap.add_argument("--num-aux-views", type=int, default=None, + help="override fixed aux-view count (smaller = faster)") + ap.add_argument("--split", default="val", choices=["train", "val"]) + ap.add_argument("--seed", type=int, default=1234, help="RNG seed for the sampler noise z0") + args = ap.parse_args() + + cfg = load_config(args.config) + data_cfg = cfg["data"] + num_steps = args.num_steps if args.num_steps is not None else cfg["train"]["val_num_steps"] + amp_dtype = cfg["train"]["amp_dtype"] + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Deterministic sampler noise (validate.py seeds VAL_SEED before sampling). + torch.manual_seed(args.seed) + if device.type == "cuda": + torch.cuda.manual_seed_all(args.seed) + + model, report = build_inference_model(cfg, args.checkpoint, device, verbose=True) + + from mvsam3d.data.dexycb_dataset import collate_multiview + subjects = args.subjects if args.subjects else data_cfg["subjects"] + frame_stride = args.frame_stride if args.frame_stride is not None else data_cfg["frame_stride"] + num_aux = (args.num_aux_views if args.num_aux_views is not None + else (data_cfg["max_aux_views"] if data_cfg["random_views"] + else data_cfg["num_aux_views"])) + ds = _build_dexycb(data_cfg, subjects=subjects, frame_stride=frame_stride, + num_aux_views=num_aux, split=args.split) + print(f"[infer] dataset: {len(ds)} samples (subjects={subjects} " + f"frame_stride={frame_stride} num_aux_views={num_aux} split={args.split})") + if len(ds) == 0: + raise RuntimeError("dataset is empty for the given filters") + + idxs: List[int] = (args.idx if args.idx is not None + else list(range(min(args.num_samples, len(ds))))) + + records = [] + for i in idxs: + batch = collate_multiview([ds[i]]) # B=1 + pred = predict_pose(model, batch, num_steps=num_steps, amp_dtype=amp_dtype, + device=device) + R = pred["rotation"][0] + det, ortho_err = _orthonormality(R) + finite = bool(torch.isfinite(R).all() and torch.isfinite(pred["translation"]).all() + and torch.isfinite(pred["scale"]).all()) + rec = { + "idx": int(i), + "object_name": batch["object_name"][0], + "rotation": R.cpu(), + "translation": pred["translation"][0].cpu(), + "scale": float(pred["scale"][0]), + "shape_latent_shape": list(pred["shape_latent"].shape), + "det": det, + "ortho_err": ortho_err, + "finite": finite, + } + records.append(rec) + print(f"[infer] idx={i:>5d} obj={rec['object_name']:<22s} " + f"det={det:+.4f} ortho_err={ortho_err:.2e} " + f"t={rec['translation'].tolist()} s={rec['scale']:.4f} finite={finite}") + + # ---- write .pt (full tensors) + .json (human-readable summary) ---- + os.makedirs(os.path.dirname(os.path.abspath(args.out)) or ".", exist_ok=True) + torch.save({"load_report": report, "num_steps": num_steps, + "checkpoint": os.path.abspath(args.checkpoint), + "records": records}, args.out) + json_path = os.path.splitext(args.out)[0] + ".json" + with open(json_path, "w") as f: + json.dump({ + "checkpoint": os.path.abspath(args.checkpoint), + "step": report["step"], + "num_steps": num_steps, + "load_report": {k: report[k] for k in + ("step", "restored_tensors", "n_missing_keys", + "n_unexpected_keys", "n_trainable_params", + "missing_trainable_params")}, + "predictions": [{ + "idx": r["idx"], "object_name": r["object_name"], + "rotation": r["rotation"].tolist(), + "translation": r["translation"].tolist(), + "scale": r["scale"], "det": r["det"], "ortho_err": r["ortho_err"], + "finite": r["finite"], "shape_latent_shape": r["shape_latent_shape"], + } for r in records], + }, f, indent=2) + print(f"[infer] wrote {os.path.abspath(args.out)}") + print(f"[infer] wrote {os.path.abspath(json_path)}") + + +if __name__ == "__main__": + main() diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..06831f0cb623a217fadb6368d5a12cd0d03503ae Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/condition.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/condition.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b0b28283cebd9c24e8ee16cd69f16ea9ed7abb6 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/condition.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/flow_model.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/flow_model.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e72452197595c00cb432e8636ffbb8b2e8b7bd81 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/flow_model.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/ss_vae.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/ss_vae.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e7d80316e3764d44a4013c4bf1d50355f528cde Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/__pycache__/ss_vae.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/condition.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/condition.py new file mode 100644 index 0000000000000000000000000000000000000000..bd059a451a8d829cda918121d3a0acf14aaeab8c --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/condition.py @@ -0,0 +1,310 @@ +"""Multi-view (multi-camera) condition embedder for the 6D-pose flow model. + +Single timestamp, multiple synchronized DexYCB cameras that share one object +pose. One camera is the *anchor* (defines the reference frame); the rest are +*aux* views. + +This module is a THIN multi-view wrapper around SAM-3D's real, defined +``EmbedderFuser`` (``sam3d_objects.model.backbone.dit.embedder.embedder_fuser``). +The fuser is constructed from the SAME config SAM-3D ships in +``checkpoints/hf/ss_generator.yaml`` (``module.condition_embedder.backbone``): +three embedders (2x ``Dino`` + ``PointPatchEmbed``), each exposing two +modality kwargs with ``cropped`` / ``full`` positional-encoding groups, a learned +``idx_emb`` positional/modality embedding, per-embedder LayerNorm+FeedForward +``projection_nets``, and modality dropout. Its pretrained weights load from +``ss_generator.ckpt`` under prefix ``_base_models.condition_embedder.*``. + +The ONLY thing this wrapper adds on top of the defined fuser is multi-view: +the SAME fuser instance is run per camera, and a learned per-view *camera* +embedding (``mlp_cam`` -- the single module with NO pretrained source) is added +to each view's tokens. ``mlp_cam`` maps that view's relative aux->anchor +extrinsics (rot6d + SSI-normalized camera centre, 9-D; the anchor is identity) +to ``model_channels``, so the camera frame is bound to every token of the view -- +including the RGB/mask tokens that otherwise carry no geometry. It replaces the +earlier content-free, order-arbitrary ``view_emb`` index table (which encoded +only "which slot", not where the camera is). Per the locked design (memory +mv-sam3d-6d-design): + * anchor view -> the full set of 6 modality kwargs SAM-3D's single view uses + (image / rgb_image / mask / rgb_image_mask / pointmap / rgb_pointmap) + * each aux view -> the 3 object-centric cropped kwargs + (image / mask / pointmap) +Aux-view pointmaps are expected to be already transformed into the anchor +camera frame *before* SSI normalization (handled in the dataloader). Pointmaps +carry NaN at invalid/background pixels so the pretrained ``invalid_xyz_token`` +pathway in ``PointPatchEmbed`` fires (see dexycb_dataset._build_view). +""" +from __future__ import annotations + +import math +from typing import Dict, List, Optional, Tuple + +import torch +import torch.nn as nn + +# SAM-3D's real, defined modules (vendored in-repo under sam3d_objects/). +from sam3d_objects.model.backbone.dit.embedder.dino import Dino +from sam3d_objects.model.backbone.dit.embedder.pointmap import PointPatchEmbed +from sam3d_objects.model.backbone.dit.embedder.embedder_fuser import EmbedderFuser + +# Our dataloader view-dict keys -> the fuser's expected kwarg names. These map +# the (modality, crop) view dict onto exactly the kwargs declared in the +# ss_generator.yaml embedder_list (see module docstring). +# image_cropped -> image (Dino #0, pos group "cropped") +# image_full -> rgb_image (Dino #0, pos group "full") +# mask_cropped -> mask (Dino #1, pos group "cropped") +# mask_full -> rgb_image_mask (Dino #1, pos group "full") +# pointmap_cropped -> pointmap (PointPatchEmbed, pos group "cropped") +# pointmap_full -> rgb_pointmap (PointPatchEmbed, pos group "full") +VIEWKEY_TO_KWARG = { + "image_cropped": "image", + "image_full": "rgb_image", + "mask_cropped": "mask", + "mask_full": "rgb_image_mask", + "pointmap_cropped": "pointmap", + "pointmap_full": "rgb_pointmap", +} +# Per-view camera-feature width fed to mlp_cam (rot6d 6 + SSI camera centre 3). +# Must match dexycb_dataset._cam_feature / CAM_FEAT_DIM. +CAM_FEAT_DIM = 9 + +# Anchor passes all six fuser kwargs; aux passes only the three cropped ones. +ANCHOR_VIEWKEYS = [ + "image_cropped", "image_full", + "mask_cropped", "mask_full", + "pointmap_cropped", "pointmap_full", +] +AUX_VIEWKEYS = ["image_cropped", "mask_cropped", "pointmap_cropped"] +# The fuser kwargs that belong to the "cropped" object-centric subset (used by +# every aux view). Order matters: it must follow the embedder_list order so the +# tokens line up with the anchor's cropped tokens. +_AUX_KWARGS = {VIEWKEY_TO_KWARG[k] for k in AUX_VIEWKEYS} + + +def _build_embedder_list( + dino_model: str, + dino_input_size: int, + pointmap_input_size: int, + pointmap_patch_size: int, + pointmap_embed_dim: int, + freeze_dino: bool, +) -> List[Tuple[nn.Module, List[Tuple[str, Optional[str]]]]]: + """Construct the exact ``embedder_list`` from ss_generator.yaml's + ``module.condition_embedder.backbone.embedder_list`` (3 embedders, each with + a cropped kwarg + a full kwarg and matching positional-encoding group).""" + dino_image = Dino( + input_size=dino_input_size, + dino_model=dino_model, + normalize_images=True, + freeze_backbone=freeze_dino, + ) + dino_mask = Dino( + input_size=dino_input_size, + dino_model=dino_model, + normalize_images=True, + freeze_backbone=freeze_dino, + ) + pointmap = PointPatchEmbed( + input_size=pointmap_input_size, + patch_size=pointmap_patch_size, + embed_dim=pointmap_embed_dim, + remap_output="linear", + ) + return [ + (dino_image, [("image", "cropped"), ("rgb_image", "full")]), + (dino_mask, [("mask", "cropped"), ("rgb_image_mask", "full")]), + (pointmap, [("pointmap", "cropped"), ("rgb_pointmap", "full")]), + ] + + +class MultiViewConditionEmbedder(nn.Module): + """Multi-view wrapper around SAM-3D's real ``EmbedderFuser``. + + Produces the ``[B, N, model_channels]`` condition sequence for all views by + running the SAME fuser per camera and tagging each view's tokens with a + learned per-view embedding. + """ + + # Real key prefixes inside ss_generator.ckpt["state_dict"]: + # _base_models.condition_embedder.* (module_list.0/1 = DINOs, + # module_list.2 = PointPatchEmbed, projection_nets.0/1/2, idx_emb) + _COND_PREFIX = "_base_models.condition_embedder." + + def __init__( + self, + model_channels: int = 1024, + max_views: int = 8, + dino_model: str = "dinov2_vitl14_reg", + dino_input_size: int = 518, + pointmap_input_size: int = 256, + pointmap_patch_size: int = 8, + pointmap_embed_dim: int = 512, + projection_hidden_mult: float = 4.0, + dropout_prob: float = 0.1, + cam_hidden_dim: int = 256, + freeze: bool = True, + # DINO backbones are frozen at construction (their hub weights never + # train); ``freeze`` then freezes the WHOLE fuser after warm-start, + # matching SAM-3D's freeze: true (the fuser owns freezing). + freeze_dino: bool = True, + ): + super().__init__() + self.model_channels = model_channels + # Retained for API compatibility only. mlp_cam derives the per-view bias + # from the actual extrinsics, so there is no per-slot table to size and no + # upper bound on the view count (any number of cameras is fine). + self.max_views = max_views + self.freeze = freeze + + embedder_list = _build_embedder_list( + dino_model=dino_model, + dino_input_size=dino_input_size, + pointmap_input_size=pointmap_input_size, + pointmap_patch_size=pointmap_patch_size, + pointmap_embed_dim=pointmap_embed_dim, + freeze_dino=freeze_dino, + ) + + # The defined SAM-3D fuser, instantiated with the SAME kwargs as + # ss_generator.yaml (use_pos_embedding=learned, proj mult 4.0, dropout + # 0.1, pointmap/rgb_pointmap drop weight). ``freeze`` is applied here + # by the fuser itself; we additionally re-assert it after warm-start. + self.fuser = EmbedderFuser( + embedder_list=embedder_list, + use_pos_embedding="learned", + projection_pre_norm=True, + projection_net_hidden_dim_multiplier=projection_hidden_mult, + drop_modalities_weight=[(["pointmap", "rgb_pointmap"], 1.0)], + dropout_prob=dropout_prob, + force_drop_modalities=None, + freeze=False, # freeze AFTER warm-start so loading is unaffected + ) + # The fuser's fused output width (== embed_dims == DINO embed dim). We + # need cond_channels == fuser embed_dims for the per-view embedding to be + # additive on the fuser tokens. + self.fuser_dim = self.fuser.embed_dims + if self.fuser_dim != model_channels: + raise ValueError( + f"fuser embed_dims ({self.fuser_dim}) must equal model_channels " + f"({model_channels}); the backbone expects cond_channels=" + f"{model_channels}." + ) + + # Precompute the aux (cropped-only) embedder_list view: same embedder + # objects + same projection-net indices, but each embedder keeps only its + # cropped kwarg. Swapping this into the fuser for aux calls reuses the + # fuser's forward verbatim (per-embedder index -> projection_nets[i] and + # the "cropped" positional group are all preserved). + self._aux_embedder_list = [ + (emb, [(kw, pg) for (kw, pg) in info if kw in _AUX_KWARGS]) + for emb, info in self.fuser.embedder_list + ] + + # The ONLY from-scratch module: a per-view CAMERA embedding. There is no + # pretrained source (SAM-3D's EmbedderFuser is single-view). mlp_cam maps + # the relative aux->anchor extrinsics (CAM_FEAT_DIM = rot6d + SSI camera + # centre) to model_channels; its output is added to every token of the + # view, binding the camera frame to the (otherwise geometry-free) RGB/mask + # tokens. The output layer is ZERO-initialized so at step 0 the per-view + # bias is exactly 0 -> conditioning equals the no-tag baseline and the + # warm-started fuser/backbone are perturbed only as mlp_cam learns. + self.mlp_cam = nn.Sequential( + nn.Linear(CAM_FEAT_DIM, cam_hidden_dim), + nn.SiLU(), + nn.Linear(cam_hidden_dim, model_channels), + ) + nn.init.zeros_(self.mlp_cam[-1].weight) + nn.init.zeros_(self.mlp_cam[-1].bias) + + # ---- pretrained warm-start ----------------------------------------- + @torch.no_grad() + def load_pretrained(self, ckpt_path: str) -> None: + """Load ``_base_models.condition_embedder.*`` into the EmbedderFuser. + + Strict where shapes match; per-submodule counts reported. The two + DINOs, PointPatchEmbed, all three projection_nets and idx_emb are + expected to load. ``mlp_cam`` has no pretrained source and stays at + init. Honors ``self.freeze`` (re-asserted after loading).""" + ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) + sd = ckpt.get("state_dict", ckpt) if isinstance(ckpt, dict) else ckpt + + own = self.fuser.state_dict() + remap, mismatched, missing_src = {}, 0, 0 + for nk in own: + src = self._COND_PREFIX + nk + if src in sd and hasattr(sd[src], "shape"): + if own[nk].shape == sd[src].shape: + remap[nk] = sd[src] + else: + mismatched += 1 + else: + missing_src += 1 + self.fuser.load_state_dict(remap, strict=False) + + # Per-submodule load report. + def _count(prefix: str) -> Tuple[int, int]: + tot = sum(1 for nk in own if nk.startswith(prefix)) + got = sum(1 for nk in remap if nk.startswith(prefix)) + return got, tot + + for label, prefix in [ + ("dino_image (module_list.0)", "module_list.0."), + ("dino_mask (module_list.1)", "module_list.1."), + ("pointmap (module_list.2)", "module_list.2."), + ("projection_nets.0", "projection_nets.0."), + ("projection_nets.1", "projection_nets.1."), + ("projection_nets.2", "projection_nets.2."), + ("idx_emb", "idx_emb"), + ]: + got, tot = _count(prefix) + print(f"[cond warm-start] {label}: loaded {got}/{tot}") + print( + f"[cond warm-start] fuser TOTAL: loaded {len(remap)}/{len(own)} " + f"({mismatched} shape-mismatched, {missing_src} no pretrained source)." + ) + print("[cond warm-start] mlp_cam: from-scratch (no pretrained source; " + "output layer zero-init).") + + if self.freeze: + self.fuser.requires_grad_(False) + self.fuser.eval() + print("[cond warm-start] EmbedderFuser FROZEN (freeze=True).") + else: + print("[cond warm-start] EmbedderFuser TRAINABLE (freeze=False).") + + def train(self, mode: bool = True): + """Keep the (frozen) fuser in eval mode regardless of the parent's mode, + mirroring SAM-3D, while letting mlp_cam participate normally.""" + super().train(mode) + if self.freeze: + self.fuser.eval() + return self + + # ---- forward --------------------------------------------------------- + def _fuser_kwargs(self, view: Dict[str, torch.Tensor], viewkeys: List[str]) -> Dict[str, torch.Tensor]: + return {VIEWKEY_TO_KWARG[k]: view[k] for k in viewkeys if k in view} + + def forward(self, views: List[Dict[str, torch.Tensor]]) -> torch.Tensor: + """views[0] = anchor (6 modality groups), views[1:] = aux (3 cropped). + + Each group tensor is [B, C, H, W]. Returns the concatenated condition + sequence [B, N, model_channels].""" + tokens: List[torch.Tensor] = [] + full_embedder_list = self.fuser.embedder_list + try: + for v, view in enumerate(views): + if v == 0: + # Anchor: full 6-kwarg call, exactly as SAM-3D's single view. + self.fuser.embedder_list = full_embedder_list + kw = self._fuser_kwargs(view, ANCHOR_VIEWKEYS) + else: + # Aux: cropped-only call via the precomputed embedder_list + # view (reuses the fuser's forward verbatim). + self.fuser.embedder_list = self._aux_embedder_list + kw = self._fuser_kwargs(view, AUX_VIEWKEYS) + tok = self.fuser(**kw) # [B, T, model_channels] + cam_emb = self.mlp_cam(view["cam"]) # [B, model_channels] + tok = tok + cam_emb[:, None, :] # bind this view's camera frame + tokens.append(tok) + finally: + self.fuser.embedder_list = full_embedder_list + return torch.cat(tokens, dim=1) diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/flow_model.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/flow_model.py new file mode 100644 index 0000000000000000000000000000000000000000..295dcbcb1bfbe49bfd0215d22b38eb204c3e26a0 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/flow_model.py @@ -0,0 +1,425 @@ +"""Joint flow-matching model over the voxel-latent + layout (R, t, s) streams. + +Plain rectified flow matching (v-prediction, sigma_min=0), NO shape-K/V, NO +shortcut. A single MoT DiT denoises two streams jointly in fully-symmetric +self-attention (mode="default"): + + * shape : the REAL SAM-3D SS latent [B,8,16,16,16] flattened to [B,4096,8] + via view(B,8,-1).permute(0,2,1) (row-major to match + ShapePositionEmbedder's meshgrid ij ordering). + * pose : {6drotation_normalized [B,1,6], translation [B,1,3], scale [B,1,1]} + merged into one 3-token transformer stream. + +Both streams cross-attend to the multi-view condition tokens. The MoT backbone +class (``SparseStructureFlowTdfyWrapper``) and the per-stream ``Latent`` / +position embedders are reused VERBATIM from the sam3d-for-6d fork (vendored +in-repo under ``sam3d_objects/``); everything else (config, flow step, losses, +warm-start) is new and clean. + +The 6D rotation is handled with the VENDORED column-convention utilities +(``matrix_to_rot6d`` / ``rot6d_to_matrix`` / ``normalize_rot6d`` / +``denormalize_rot6d``) consistently with the dataset — there is NO inline +interleaved encode here (that was a real bug in the prior attempt). +""" +from __future__ import annotations + +from typing import Dict, Optional + +import torch +import torch.nn as nn + +from sam3d_objects.model.backbone.tdfy_dit.models.mm_latent import ( + Latent, + LearntPositionEmbedder, + ShapePositionEmbedder, +) +from sam3d_objects.model.backbone.tdfy_dit.models.mot_sparse_structure_flow import ( + SparseStructureFlowTdfyWrapper, +) +from sam3d_objects.model.backbone.generator.flow_matching.solver import Euler +from sam3d_objects.geometry.geometry import ( + denormalize_rot6d, + rot6d_to_matrix, + ssi_to_metric, +) + +from .condition import MultiViewConditionEmbedder + +# Loss weights for the 6D objective. All 1.0 (matches the pretrained backbone's +# training regime; the YAML config is authoritative — this is only the fallback). +DEFAULT_LOSS_WEIGHTS = { + "shape": 1.0, + "6drotation_normalized": 1.0, + "translation": 1.0, + "scale": 1.0, +} + +# Per-stream latent inner dims and token lengths. +SHAPE_RES = 16 # latent grid is 16^3 -> 4096 tokens +SHAPE_CHANNELS = 8 +TIME_SCALE = 1000.0 + + +def decode_pose_rotation(rot6d_normalized: torch.Tensor) -> torch.Tensor: + """Predicted normalized 6D rotation -> rotation matrix [...,3,3]. + + Inverse of the dataset encode ``normalize_rot6d(matrix_to_rot6d(R))``: this + is ``rot6d_to_matrix(denormalize_rot6d(x))`` with the vendored + column-convention functions.""" + return rot6d_to_matrix(denormalize_rot6d(rot6d_normalized)) + + +def decode_pose_translation( + translation_ssi: torch.Tensor, + pm_scale: torch.Tensor, + pm_shift: torch.Tensor, +) -> torch.Tensor: + """SSI-space translation -> metric translation [...,3]. + + Inverse of the dataset encode ``t_ssi = (t - shift) / pm_scale_safe``; + delegates to the VENDORED ``geometry.ssi_to_metric`` (byte-identical to the + fork) — mirrors fork model_wrapper._decode_latents line 326: + ``trans_metric = ssi_to_metric(trans_ssi, pm_scale, pm_shift)``. + ``pm_scale`` is the SSI scale (scalar, shape [...,1] or [...,3]); ``pm_shift`` + is the SSI shift [...,3]. Both are threaded from the batch.""" + return ssi_to_metric(translation_ssi, pm_scale, pm_shift) + + +def decode_pose_scale( + scale_log: torch.Tensor, + pm_scale: torch.Tensor, +) -> torch.Tensor: + """SSI-space log-ratio scale -> metric scale. + + Inverse of the dataset encode ``s_log = log((sc / pm_scale_safe).clamp_min(1e-6))`` + (fork model_wrapper lines 580-581). The fork's canonical scale decode is + ``metric_scale = exp(scale_latent) * pm_scale`` (fork line 568); there is NO + vendored decode helper for scale in geometry.py, so this is the reference + form. ``scale_log`` is the 3-ch isotropic log-ratio (channels identical); + ``pm_scale`` is the SSI scale threaded from the batch.""" + return torch.exp(scale_log) * pm_scale + + +class MVSam3DFlow(nn.Module): + """Multi-view 6D-pose joint flow model.""" + + def __init__( + self, + model_channels: int = 1024, + cond_channels: int = 1024, + num_blocks: int = 24, + num_heads: int = 16, + max_views: int = 8, + loss_weights: Optional[Dict[str, float]] = None, + sigma_min: float = 0.0, + time_scale: float = TIME_SCALE, + timestep_sampler: str = "logit_normal", + timestep_lognorm_mean: float = -1.0, + timestep_lognorm_std: float = 1.0, + pretrained_ckpt: Optional[str] = None, + condition_kwargs: Optional[Dict] = None, + ): + super().__init__() + self.sigma_min = sigma_min + self.time_scale = time_scale + # Timestep sampling for the flow loss. Default = the ORIGINAL SAM-3D + # generator pretraining schedule: logit-normal(mean=-1, std=1) + # (ss_generator.yaml training_time_sampler_fn=lognorm_sampler mean=-1). + # v2's flow_step convention (t=0 noise -> t=1 data) is identical to the + # FlowMatching module's (_generate_xt_tensor), so the mean maps directly + # with NO sign flip. "uniform" reproduces the stage1_flow_lora harness. + self.timestep_sampler = timestep_sampler + self.timestep_lognorm_mean = timestep_lognorm_mean + self.timestep_lognorm_std = timestep_lognorm_std + self.loss_weights = dict(loss_weights or DEFAULT_LOSS_WEIGHTS) + + latent_mapping = { + "shape": Latent( + in_channels=SHAPE_CHANNELS, + model_channels=model_channels, + pos_embedder=ShapePositionEmbedder( + model_channels=model_channels, resolution=SHAPE_RES, patch_size=1 + ), + ), + "6drotation_normalized": Latent( + in_channels=6, + model_channels=model_channels, + pos_embedder=LearntPositionEmbedder(model_channels, token_len=1), + ), + "translation": Latent( + in_channels=3, + model_channels=model_channels, + pos_embedder=LearntPositionEmbedder(model_channels, token_len=1), + ), + "scale": Latent( + in_channels=3, + model_channels=model_channels, + pos_embedder=LearntPositionEmbedder(model_channels, token_len=1), + ), + } + # Pose streams share one transformer; shape is its own stream. -> 2 MoT streams. + # The shared group MUST be named "6drotation_normalized" to match the + # pretrained ss_generator block-stream naming: the fork names a shared MoT + # group after one of its member latents, so the per-block transformer keys + # become blocks.N.<...>.6drotation_normalized.*. load_state_dict matches + # by key name; naming this group anything else (e.g. "pose") leaves all + # 432 shared-stream block tensors at random init. With this name they + # warm-start 432/432 from the checkpoint (verified, 0 shape mismatches). + latent_share_transformer = { + "6drotation_normalized": ["6drotation_normalized", "translation", "scale"] + } + + self.backbone = SparseStructureFlowTdfyWrapper( + latent_mapping=latent_mapping, + latent_share_transformer=latent_share_transformer, + in_channels=SHAPE_CHANNELS, + out_channels=SHAPE_CHANNELS, + model_channels=model_channels, + cond_channels=cond_channels, + num_blocks=num_blocks, + num_heads=num_heads, + mlp_ratio=4, + pe_mode="ape", + qk_rms_norm=True, + use_fp16=False, + use_checkpoint=False, + is_shortcut_model=False, # plain flow, no shortcut distance + condition_embedder=None, # identity; we pass cond tokens directly + force_zeros_cond=False, + freeze_shared_parameters=False, # unfreeze everything + ) + self.condition = MultiViewConditionEmbedder( + model_channels=cond_channels, max_views=max_views, + **(condition_kwargs or {}) + ) + + # Stream key -> latent token length (from each Latent's positional table). + self._token_len = { + name: lat.pos_emb.shape[0] for name, lat in latent_mapping.items() + } + self._inner_dim = { + "shape": SHAPE_CHANNELS, + "6drotation_normalized": 6, + "translation": 3, + "scale": 3, + } + + if pretrained_ckpt is not None: + self.load_pretrained_backbone(pretrained_ckpt) + self.condition.load_pretrained(pretrained_ckpt) + + # ---- helpers --------------------------------------------------------- + def latent_shape_dict(self, B: int, device) -> Dict[str, torch.Size]: + return { + name: (B, self._token_len[name], self._inner_dim[name]) + for name in self._inner_dim + } + + def gt_streams(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """Assemble GT target latents (z1) for every stream from a batch.""" + B = batch["shape_latent"].shape[0] + # [B,8,16,16,16] -> [B,8,4096] -> [B,4096,8] (row-major; matches + # ShapePositionEmbedder meshgrid(indexing="ij")). + shape = batch["shape_latent"].reshape(B, SHAPE_CHANNELS, -1).permute(0, 2, 1) + return { + "shape": shape.contiguous(), # [B,4096,8] + "6drotation_normalized": batch["rot6d"].reshape(B, 1, 6), + "translation": batch["translation"].reshape(B, 1, 3), # SSI space + "scale": batch["scale"].reshape(B, 1, 3), # log-ratio, isotropic + } + + # ---- forward (DDP entrypoint) --------------------------------------- + def forward(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """Training forward == one flow step. + + DDP only all-reduces gradients for parameters touched inside the wrapped + module's ``forward``; routing training through ``model(batch)`` (not + ``model.flow_step(batch)``) is therefore REQUIRED for gradient sync. + ``flow_step`` remains a public alias for non-DDP / test use.""" + return self.flow_step(batch) + + # ---- timestep sampling ---------------------------------------------- + def _sample_timesteps(self, batch_size: int, device) -> torch.Tensor: + """Per-sample t in (0,1). Mirrors the SAM-3D FlowMatching/stage1 sampler: + logit-normal (t = sigmoid(N(mean,std))) or uniform. Per-sample size=(B,) + matches the pretraining FlowMatching._generate_t (size=(batch_size,)).""" + if self.timestep_sampler == "logit_normal": + logit = (torch.randn(batch_size, device=device) * self.timestep_lognorm_std + + self.timestep_lognorm_mean) + return torch.sigmoid(logit) + return torch.rand(batch_size, device=device) + + # ---- flow step ------------------------------------------------------- + def flow_step(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """One training step: returns {loss, loss/} (all scalars).""" + z1 = self.gt_streams(batch) + device = z1["shape"].device + B = z1["shape"].shape[0] + + t = self._sample_timesteps(B, device) # per-sample t in (0,1) + tb = t.view(B, 1, 1) + + z_t, v_target = {}, {} + for key, z1k in z1.items(): + z0 = torch.randn_like(z1k) # noise + # z_t = (1 - (1-sigma_min) t) z0 + t z1 ; v = z1 - (1-sigma_min) z0. + z_t[key] = (1.0 - (1.0 - self.sigma_min) * tb) * z0 + tb * z1k + v_target[key] = z1k - (1.0 - self.sigma_min) * z0 + + cond = self.condition(batch["views"]) + pred = self.backbone(z_t, t * self.time_scale, cond, mode="default") + + losses, total = {}, z_t["shape"].new_zeros(()) + for key, w in self.loss_weights.items(): + mse = torch.mean((pred[key] - v_target[key]) ** 2) + losses[f"loss/{key}"] = mse.detach() + total = total + w * mse + losses["loss"] = total + return losses + + # ---- sampling / inference ------------------------------------------- + @torch.no_grad() + def sample( + self, + views, + num_steps: int = 100, + cond: Optional[torch.Tensor] = None, + generator: Optional[torch.Generator] = None, + z0: Optional[Dict[str, torch.Tensor]] = None, + ) -> Dict[str, torch.Tensor]: + """Integrate the rectified-flow ODE from noise (t=0) to data (t=1) and + return the per-stream sample z1_hat {shape, 6drotation_normalized, + translation, scale}. + + Uses the VENDORED ``Euler`` solver (fork + flow_matching/solver.py, byte-identical). The dynamics are the exact + inverse of ``flow_step``: the backbone predicts the velocity + v = z1 - (1-sigma_min) z0 at each (z_t, t), integrated with the same + ``time_scale`` and ``mode="default"`` used in training. + + ``z0`` (optional): a dict of per-stream initial-noise tensors keyed by + stream name, each shaped exactly ``latent_shape_dict(B, device)[name]``. + When provided it is used VERBATIM instead of drawing fresh ``randn``; + this is the ONLY behavioural hook and exists solely so a batched-vs-B=1 + equivalence test can feed IDENTICAL noise into both paths (the model + math is otherwise untouched). ``None`` reproduces the original sampler + exactly (fresh ``torch.randn`` per stream).""" + if cond is None: + cond = self.condition(views) + B = cond.shape[0] + device = cond.device + + # z0 ~ N(0, I) for every stream, at the stream's token/inner shape. + if z0 is None: + z0 = { + name: torch.randn(shape, device=device, generator=generator) + for name, shape in self.latent_shape_dict(B, device).items() + } + else: + # Validate caller-supplied noise matches the expected per-stream shapes + # and batch size (no silent broadcasting / truncation). + want = self.latent_shape_dict(B, device) + assert set(z0.keys()) == set(want.keys()), ( + f"z0 keys {sorted(z0.keys())} != expected {sorted(want.keys())}") + z0 = {name: z0[name].to(device) for name in want} + for name, shape in want.items(): + assert tuple(z0[name].shape) == tuple(shape), ( + f"z0['{name}'] shape {tuple(z0[name].shape)} != expected {tuple(shape)}") + + def dynamics_fn(x_t, t): + # ``t`` is a scalar in [0,1] from the solver's time grid; broadcast to + # [B] and rescale exactly as in flow_step before the backbone call. + t_vec = torch.full((B,), float(t) * self.time_scale, device=device) + return self.backbone(x_t, t_vec, cond, mode="default") + + times = torch.linspace(0.0, 1.0, num_steps + 1, device=device) + z1 = Euler().solve(dynamics_fn, z0, times) + return z1 + + @torch.no_grad() + def predict(self, batch: Dict[str, torch.Tensor], num_steps: int = 100, + z0: Optional[Dict[str, torch.Tensor]] = None) -> Dict[str, torch.Tensor]: + """Sample, then decode every stream to its metric-space quantity. + + Returns: + shape_latent [B,8,16,16,16] — inverse of gt_streams' flatten, + rotation [B,3,3] — decode_pose_rotation, + translation [B,3] (metric) — decode_pose_translation(SSI, pm_*), + scale [B] (metric) — decode_pose_scale(log-ratio, pm_scale). + ``pm_shift`` / ``pm_scale`` are threaded from the batch (shared anchor SSI). + ``z0`` (optional) is a passthrough to ``sample`` for the equivalence test; + ``None`` is the production path (fresh sampler noise).""" + z1 = self.sample(batch["views"], num_steps=num_steps, z0=z0) + B = z1["shape"].shape[0] + + # shape: [B,4096,8] -> [B,8,16,16,16] (inverse of gt_streams). + shape_latent = ( + z1["shape"].permute(0, 2, 1).reshape(B, SHAPE_CHANNELS, SHAPE_RES, SHAPE_RES, SHAPE_RES) + ).contiguous() + + rot6d = z1["6drotation_normalized"].reshape(B, 6) + rotation = decode_pose_rotation(rot6d) # [B,3,3] + + pm_shift = batch["pm_shift"] # [B,3] + pm_scale = batch["pm_scale"] # [B,1] + translation = decode_pose_translation( + z1["translation"].reshape(B, 3), pm_scale, pm_shift + ) # [B,3] metric + + # 3-ch isotropic log-ratio -> metric; take channel 0 (channels identical). + scale = decode_pose_scale(z1["scale"].reshape(B, 3), pm_scale)[:, 0] # [B] metric + + return { + "shape_latent": shape_latent, + "rotation": rotation, + "translation": translation, + "scale": scale, + } + + # ---- warm start ------------------------------------------------------ + @torch.no_grad() + def load_pretrained_backbone(self, ckpt_path: str) -> None: + """Warm-start the MoT torso + matching latent streams from the pretrained + ss_generator checkpoint. + + Real torso prefix (verified): + _base_models.generator.reverse_fn.backbone..* + + The pretrained model is a SHORTCUT model with an EXTRA ``translation_scale`` + stream and a shortcut ``d_embedder`` we do not have (their keys find no + home and are skipped cleanly). Our ``scale`` stream is now 3-channel + (SSI log-ratio, isotropic) to MATCH the pretrained 3-channel scale head, + so its input/out layers now load (previously shape-mismatched as 1-ch and + skipped). Nothing is ever forced; loading is purely by shape match.""" + ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) + sd = ckpt.get("state_dict", ckpt) if isinstance(ckpt, dict) else ckpt + own = self.backbone.state_dict() + # Verified torso prefix; anchoring on it avoids colliding with the + # condition embedder's own ``blocks.`` keys (DINO / PointPatchEmbed). + torso_prefix = "_base_models.generator.reverse_fn.backbone." + anchors = ("blocks.", "latent_mapping.", "t_embedder.") + remapped, skipped_shape = {}, 0 + for k, v in sd.items(): + if not hasattr(v, "shape") or not k.startswith(torso_prefix): + continue + nk = k[len(torso_prefix):] + if not any(nk.startswith(a) for a in anchors): + continue + if nk in own: + if own[nk].shape == v.shape: + remapped[nk] = v + else: + skipped_shape += 1 + missing = self.backbone.load_state_dict(remapped, strict=False) + # Per-stream load report (which latent streams got pretrained weights). + streams = ("shape", "6drotation_normalized", "translation", "scale") + per_stream = { + s: sum(1 for nk in remapped if nk.startswith(f"latent_mapping.{s}.")) + for s in streams + } + print( + f"[warm-start] backbone: loaded {len(remapped)}/{len(own)} tensors " + f"from {ckpt_path}; {len(missing.missing_keys)} left at init; " + f"{skipped_shape} shape-mismatched (e.g. shortcut 3-ch scale) skipped." + ) + print(f"[warm-start] backbone latent_mapping per-stream loaded: {per_stream}") diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/ss_vae.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/ss_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..828a4acf4cf171f480550c2b55d59b145a2cc9d7 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/model/ss_vae.py @@ -0,0 +1,103 @@ +"""Sparse-structure (SS) VAE loader for the REAL SAM-3D latent space. + +This is the authoritative SAM-3D sparse-structure VAE from +``facebook/sam-3d-objects`` (revision ``2e73555018d2741ccd486e56c24fac41155a1dc6``), +**NOT** TRELLIS's. The distinction matters: + + * The SAM-3D SS encoder (``checkpoints/hf/ss_encoder.ckpt``, ~119 MB) and + TRELLIS's conv3d_16l8 encoder share the SAME architecture but DIFFERENT + weights (max abs diff ~2.5) → mutually incompatible latent spaces. + * The SAM-3D standalone decoder (``checkpoints/hf/ss_decoder.ckpt``, ~147 MB) + round-trips ONLY with the SAM-3D encoder (IoU≈1.0). It is the authoritative + space and the ONE used to validate GT latents. + * WARNING: the decoder embedded inside ``ss_generator.ckpt`` + (``_base_models.decoder.*``) is TRELLIS-identical and would give IoU≈1.0 on + the WRONG latents — do NOT use it to validate; use the standalone decoder. + +The model classes are the fork's REAL ``SparseStructureEncoder`` / +``SparseStructureDecoder`` (vendored verbatim under +``sam3d_objects/model/backbone/tdfy_dit/models/sparse_structure_vae.py``). The +checkpoints load with 0 missing / 0 unexpected keys. + +Encoder: dense ``[B,1,64,64,64]`` 0/1 occupancy → ``[B,8,16,16,16]`` SS latent +(posterior MEAN, ``sample_posterior=False`` → deterministic GT regression +target). Decoder: latent → ``[B,1,64,64,64]`` occupancy logits. +""" +from __future__ import annotations + +import torch + +from sam3d_objects.model.backbone.tdfy_dit.models.sparse_structure_vae import ( + SparseStructureDecoder, + SparseStructureEncoder, +) + +from mvsam3d import paths as _paths + +# Authoritative SAM-3D SS VAE checkpoints (migrator/checkpoints/hf; portable). +SS_ENCODER_CKPT = _paths.SS_ENCODER_CKPT +SS_DECODER_CKPT = _paths.SS_DECODER_CKPT + +# Architecture from ss_encoder.yaml / ss_decoder.yaml (verified: 0 missing / 0 +# unexpected on load). +_ENC_KW = dict(in_channels=1, latent_channels=8, channels=[32, 128, 512], + num_res_blocks=2, num_res_blocks_middle=2) +_DEC_KW = dict(out_channels=1, latent_channels=8, channels=[512, 128, 32], + num_res_blocks=2, num_res_blocks_middle=2, reshape_input_to_cube=False) + + +def _state_dict(path: str) -> dict: + sd = torch.load(path, map_location="cpu", weights_only=False) + if isinstance(sd, dict) and "state_dict" in sd: + sd = sd["state_dict"] + return sd + + +def load_ss_encoder(device="cuda", ckpt_path: str = SS_ENCODER_CKPT) -> SparseStructureEncoder: + """REAL SAM-3D sparse-structure encoder (defines the GT latent space).""" + enc = SparseStructureEncoder(**_ENC_KW) + missing, unexpected = enc.load_state_dict(_state_dict(ckpt_path), strict=True) + assert not missing and not unexpected, (missing, unexpected) + enc = enc.to(device).eval() + for p in enc.parameters(): + p.requires_grad = False + return enc + + +def load_ss_decoder(device="cuda", ckpt_path: str = SS_DECODER_CKPT) -> SparseStructureDecoder: + """REAL standalone SAM-3D sparse-structure decoder (authoritative for IoU).""" + dec = SparseStructureDecoder(**_DEC_KW) + missing, unexpected = dec.load_state_dict(_state_dict(ckpt_path), strict=True) + assert not missing and not unexpected, (missing, unexpected) + dec = dec.to(device).eval() + for p in dec.parameters(): + p.requires_grad = False + return dec + + +@torch.no_grad() +def encode_occupancy(encoder: SparseStructureEncoder, occ: torch.Tensor) -> torch.Tensor: + """occ ``[B,1,64,64,64]`` (0/1 float) → latent ``[B,8,16,16,16]`` (posterior mean). + + Uses ``sample_posterior=False`` so the cached GT latent is the deterministic + posterior mean (no sampling noise) — the correct fixed regression target. + """ + occ = occ.to(encoder.device, dtype=torch.float32) + z = encoder(occ, sample_posterior=False, return_raw=False) + return z.float() + + +@torch.no_grad() +def decode_latent(decoder: SparseStructureDecoder, latent: torch.Tensor) -> torch.Tensor: + """latent ``[B,8,16,16,16]`` → occupancy logits ``[B,1,64,64,64]``.""" + latent = latent.to(decoder.device, dtype=torch.float32) + return decoder(latent).float() + + +def voxel_iou(logits: torch.Tensor, gt_occ: torch.Tensor, thresh: float = 0.0) -> float: + """IoU between decoded occupancy (logits>thresh) and GT occupancy ``[B,1,64,64,64]``.""" + pred = logits > thresh + gt = gt_occ.to(logits.device) > 0.5 + inter = (pred & gt).sum().item() + union = (pred | gt).sum().item() + return inter / union if union > 0 else 1.0 diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/paths.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/paths.py new file mode 100644 index 0000000000000000000000000000000000000000..c0e6c6acdbdddd8fcd5a142e8df785e5fadd4cf9 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/paths.py @@ -0,0 +1,90 @@ +"""Central, portable path resolution for the migrator package. + +Every absolute ``/NHNHOME/...`` path that used to be hardcoded in the code and +configs is derived here from a single root. Nothing below this module hardcodes +an install location, so the package works unchanged when copied elsewhere. + +Resolution order for the package root (``MIGRATOR_ROOT``): + 1. the ``MIGRATOR_ROOT`` environment variable, if set; + 2. otherwise auto-detected from this file's location. This file lives at + ``/code/mv-sam3d-for-6d-v2/mvsam3d/paths.py`` so the root is + three parents up from the package dir. + +Layout assumed under ``MIGRATOR_ROOT``: + code/mv-sam3d-for-6d-v2/ this repo (on PYTHONPATH) + code/sam3d-for-6d/ the fork; contains ``sam3d_objects`` (the + ObjectCentricSSI source, resolved via __path__) + checkpoints/ model + SAM-3D dependency checkpoints + checkpoints/hf/ ss_generator.ckpt, ss_decoder.ckpt, ss_encoder.ckpt + cache/canonical/ precomputed canonical latents (built by the user) + cache/ckpt/ training output checkpoints + +Datasets are huge and downloaded separately, so their root is independent: + MIGRATOR_DATA_ROOT (default ``/data``) + +Individual locations can each be overridden by their own env var (see below), +but the defaults are all relative to ``MIGRATOR_ROOT`` / ``MIGRATOR_DATA_ROOT``. +""" +from __future__ import annotations + +import os +from pathlib import Path + + +def _auto_root() -> str: + # /code/mv-sam3d-for-6d-v2/mvsam3d/paths.py -> parents[3] == + return str(Path(__file__).resolve().parents[3]) + + +MIGRATOR_ROOT = os.environ.get("MIGRATOR_ROOT", _auto_root()) + +CODE_ROOT = os.path.join(MIGRATOR_ROOT, "code") +MV_REPO = os.path.join(CODE_ROOT, "mv-sam3d-for-6d-v2") +CONFIG_DIR = os.path.join(MV_REPO, "configs") + +# Directory that CONTAINS the fork's ``sam3d_objects`` package. It is NOT put on +# PYTHONPATH (the repo's own partial sam3d_objects must stay primary); the data +# loader appends "/sam3d_objects[/data]" to the package __path__ to +# resolve ObjectCentricSSI. See mvsam3d/data/dexycb_dataset._load_object_centric_ssi. +FORK_ROOT = os.environ.get("MIGRATOR_FORK_ROOT", os.path.join(CODE_ROOT, "sam3d-for-6d")) + +CHECKPOINTS = os.environ.get("MIGRATOR_CHECKPOINTS", os.path.join(MIGRATOR_ROOT, "checkpoints")) +HF_DIR = os.path.join(CHECKPOINTS, "hf") +SS_GENERATOR_CKPT = os.path.join(HF_DIR, "ss_generator.ckpt") +SS_DECODER_CKPT = os.path.join(HF_DIR, "ss_decoder.ckpt") +SS_ENCODER_CKPT = os.path.join(HF_DIR, "ss_encoder.ckpt") + +DATA_ROOT = os.environ.get("MIGRATOR_DATA_ROOT", os.path.join(MIGRATOR_ROOT, "data")) + +CACHE_ROOT = os.environ.get("MIGRATOR_CACHE", os.path.join(MIGRATOR_ROOT, "cache")) +CANONICAL_CACHE = os.environ.get("MIGRATOR_CANONICAL_CACHE", os.path.join(CACHE_ROOT, "canonical")) +CKPT_DIR = os.environ.get("MIGRATOR_CKPT_DIR", os.path.join(CACHE_ROOT, "ckpt")) +BAD_FILES_LOG = os.environ.get("MVSAM3D_BAD_FILES_LOG", os.path.join(CACHE_ROOT, "_bad_files.log")) + +DEFAULT_CONFIG = os.path.join(CONFIG_DIR, "train.yaml") + + +def register_env() -> None: + """Export the derived roots into ``os.environ`` (without clobbering values the + user already set) so YAML configs using ``${MIGRATOR_ROOT}`` / + ``${MIGRATOR_DATA_ROOT}`` / ``${MIGRATOR_CHECKPOINTS}`` expand correctly.""" + os.environ.setdefault("MIGRATOR_ROOT", MIGRATOR_ROOT) + os.environ.setdefault("MIGRATOR_MV_REPO", MV_REPO) + os.environ.setdefault("MIGRATOR_DATA_ROOT", DATA_ROOT) + os.environ.setdefault("MIGRATOR_CHECKPOINTS", CHECKPOINTS) + os.environ.setdefault("MIGRATOR_CACHE", CACHE_ROOT) + + +def expand(obj): + """Recursively ``os.path.expandvars`` every string in a nested dict/list. + + Call on a freshly loaded YAML config so ``${MIGRATOR_*}`` placeholders in + dataset roots / checkpoint paths / cache dirs become absolute paths.""" + register_env() + if isinstance(obj, dict): + return {k: expand(v) for k, v in obj.items()} + if isinstance(obj, list): + return [expand(v) for v in obj] + if isinstance(obj, str): + return os.path.expandvars(obj) + return obj diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f39f13241a50ec8a5c34f7c95d9273526862814 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/__pycache__/train.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/__pycache__/train.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da3011de35da7e814202133e42037303ceac1490 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/__pycache__/train.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/__pycache__/validate.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/__pycache__/validate.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..47edb7392d0901f2255d5e619b2db8c9c721f609 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/__pycache__/validate.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/train.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/train.py new file mode 100644 index 0000000000000000000000000000000000000000..d7dc2ba28c8172cc615ce4294b6e77623f060795 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/train.py @@ -0,0 +1,409 @@ +"""Training loop for the multi-view 6D-pose flow model. + +Plain rectified flow matching; loss = weighted sum of per-stream v-prediction +MSE (all stream weights 1.0; see configs/train.yaml loss_weights). No +shape-K/V, no shortcut. GPU-only; launch via sbatch + torchrun (never sets +CUDA_VISIBLE_DEVICES). + +Distributed: the generic DDP/DeepSpeed primitives are REUSED VERBATIM from the +vendored fork module ``sam3d_objects.training.stage1_flow_lora.distributed`` +(init/setup/barrier/reduce/cleanup). This loop is the only from-scratch glue: +it wires those primitives around v2's own ``MVSam3DFlow`` + ``MultiViewDexYCB``, +with gradient accumulation and periodic sampler-based validation. Single-GPU +(WORLD_SIZE=1) transparently skips the process group and DDP wrap. + +ALL hyperparameters come from a YAML config (default configs/train.yaml). +Nothing training-related is hardcoded here. + + sbatch sbatch/train.sh --config configs/train.yaml + # per-key overrides (repeatable, dotted keys): + sbatch sbatch/train.sh --config configs/train.yaml --set train.steps=3 --set data.random_views=false +""" +from __future__ import annotations + +import argparse +import os +import random +import time +from contextlib import nullcontext +from types import SimpleNamespace + +import numpy as np +import torch +import yaml +from torch.utils.data import DataLoader, DistributedSampler + +from sam3d_objects.training.stage1_flow_lora.distributed import ( + barrier, + cleanup_distributed, + init_distributed, + reduce_scalar_dict, + setup_runtime, +) +from sam3d_objects.training.stage1_flow_lora.train_utils import build_scheduler +from sam3d_objects.training.stage1_flow_lora.checkpointing import ( + load_checkpoint, + save_checkpoint, +) +from mvsam3d.data.dexycb_dataset import MultiViewDexYCB, collate_multiview +from mvsam3d.data.integrated import build_integrated_dataset, make_weighted_sampler +from mvsam3d.model.flow_model import MVSam3DFlow +from mvsam3d.train.validate import run_validation +from mvsam3d import paths as _paths + +DEFAULT_CONFIG = _paths.DEFAULT_CONFIG + + +def move_to(obj, device): + if torch.is_tensor(obj): + return obj.to(device, non_blocking=True) + if isinstance(obj, dict): + return {k: move_to(v, device) for k, v in obj.items()} + if isinstance(obj, list): + return [move_to(v, device) for v in obj] + return obj + + +def _coerce(text: str): + """Parse a --set override value using YAML scalar rules.""" + return yaml.safe_load(text) + + +def _apply_override(cfg: dict, dotted_key: str, value) -> None: + """Set cfg[a][b][c] = value for dotted_key 'a.b.c'; intermediate keys must exist.""" + keys = dotted_key.split(".") + node = cfg + for k in keys[:-1]: + if k not in node or not isinstance(node[k], dict): + raise KeyError(f"--set {dotted_key}: '{k}' is not a config section") + node = node[k] + if keys[-1] not in node: + raise KeyError(f"--set {dotted_key}: unknown key '{keys[-1]}'") + node[keys[-1]] = value + + +def load_config(path: str, overrides: list[str]) -> dict: + with open(path) as f: + cfg = yaml.safe_load(f) + # Expand ${MIGRATOR_ROOT} / ${MIGRATOR_DATA_ROOT} / ${MIGRATOR_CHECKPOINTS} + # placeholders in the config into absolute paths (portable install root). + cfg = _paths.expand(cfg) + for ov in overrides or []: + if "=" not in ov: + raise ValueError(f"--set expects key=value, got {ov!r}") + key, raw = ov.split("=", 1) + _apply_override(cfg, key.strip(), _coerce(raw.strip())) + return cfg + + +def seed_everything(seed: int, rank: int) -> None: + seed = int(seed) + int(rank) + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def build_dataset(data_cfg: dict, split: str, *, random_views: bool, + num_aux_views: int | None = None): + return MultiViewDexYCB( + dexycb_root=data_cfg["dexycb_root"], + canonical_cache=data_cfg["canonical_cache"], + num_aux_views=data_cfg["num_aux_views"] if num_aux_views is None else num_aux_views, + frame_stride=data_cfg["frame_stride"], + anchor_sweep=data_cfg["anchor_sweep"], + subjects=data_cfg["subjects"], + image_size=data_cfg["image_size"], + random_views=random_views, + min_aux_views=data_cfg["min_aux_views"], + max_aux_views=data_cfg["max_aux_views"], + depth_scale=data_cfg["depth_scale"], + depth_near=data_cfg["depth_near"], + depth_far=data_cfg["depth_far"], + split=split, + val_fraction=data_cfg["val_fraction"], + split_seed=data_cfg["split_seed"], + ssi_mode=data_cfg.get("ssi_mode", "scene"), + ) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--config", type=str, default=DEFAULT_CONFIG, + help="Path to the training YAML config.") + ap.add_argument("--set", dest="overrides", action="append", default=[], + metavar="KEY.SUBKEY=VALUE", + help="Override a single config value (repeatable).") + args = ap.parse_args() + + cfg = load_config(args.config, args.overrides) + data_cfg = cfg["data"] + model_cfg = cfg["model"] + cond_cfg = cfg["condition"] + loss_cfg = cfg["loss_weights"] + optim_cfg = cfg["optim"] + train_cfg = cfg["train"] + + # ── distributed init (reused vendored infra) ── + dist_state = init_distributed(train_cfg["backend"]) + seed_everything(train_cfg["seed"], dist_state.rank) + device = dist_state.device + is_main = dist_state.is_main + if is_main: + os.makedirs(train_cfg["ckpt_dir"], exist_ok=True) + barrier() + + # max_views is retained only for API compatibility: the per-view bias now + # comes from mlp_cam (derived from each view's extrinsics), so there is no + # per-slot embedding table to size and no hard upper bound on the view count. + # `max_aux` is still used below to pick a fixed, deterministic validation view + # count (not a correctness constraint anymore). + max_aux = data_cfg["max_aux_views"] if data_cfg["random_views"] else data_cfg["num_aux_views"] + max_views = max_aux + 1 + + # ── data ── + # Integrated multi-dataset path (additive): active iff data.datasets is set. + # Otherwise the original single-DexYCB path runs UNCHANGED. + integrated = bool(data_cfg.get("datasets")) + num_workers = train_cfg["num_workers"] + _int_kw = dict(image_size=data_cfg["image_size"], min_aux=data_cfg["min_aux_views"], + max_aux=data_cfg["max_aux_views"], num_aux=data_cfg["num_aux_views"], + depth_near=data_cfg["depth_near"], depth_far=data_cfg["depth_far"], + ssi_mode=data_cfg.get("ssi_mode", "scene"), + heldout_path=data_cfg.get("heldout_manifest")) + if integrated: + train_ds, train_weights = build_integrated_dataset( + data_cfg, "train", random_views=data_cfg["random_views"], **_int_kw) + train_sampler = make_weighted_sampler( + train_weights, world_size=dist_state.world_size, rank=dist_state.rank, + seed=train_cfg["seed"]) + else: + train_ds = build_dataset(data_cfg, "train", random_views=data_cfg["random_views"]) + train_sampler = ( + DistributedSampler(train_ds, num_replicas=dist_state.world_size, + rank=dist_state.rank, shuffle=True) + if dist_state.distributed else None + ) + if is_main: + print(f"[train] train samples: {len(train_ds)}", flush=True) + train_loader = DataLoader( + train_ds, batch_size=train_cfg["batch_size"], sampler=train_sampler, + shuffle=(train_sampler is None), num_workers=num_workers, + collate_fn=collate_multiview, drop_last=True, + persistent_workers=num_workers > 0, pin_memory=True, + ) + + val_every = train_cfg["val_every"] + val_loader = None + if integrated and val_every > 0: + try: + val_ds, _ = build_integrated_dataset(data_cfg, "val", random_views=False, **_int_kw) + except RuntimeError: + val_ds = None + if val_ds is not None: + val_sampler = (DistributedSampler(val_ds, num_replicas=dist_state.world_size, + rank=dist_state.rank, shuffle=False) + if dist_state.distributed else None) + val_loader = DataLoader( + val_ds, batch_size=train_cfg["batch_size"], sampler=val_sampler, + shuffle=False, num_workers=max(1, num_workers // 2), + collate_fn=collate_multiview, drop_last=False, pin_memory=True) + if is_main: + print(f"[train] val samples: {len(val_ds)}", flush=True) + elif val_every > 0 and data_cfg["val_fraction"] > 0.0: + # Validation views are deterministic (fixed count) so metrics aren't noised + # by random view sampling; the held-out split is disjoint from train. The + # aux count is fixed at `max_aux` purely for a stable, representative view + # count (mlp_cam imposes no upper bound, unlike the old view_emb table). + val_ds = build_dataset(data_cfg, "val", random_views=False, num_aux_views=max_aux) + if is_main: + print(f"[train] val samples: {len(val_ds)}", flush=True) + val_sampler = ( + DistributedSampler(val_ds, num_replicas=dist_state.world_size, + rank=dist_state.rank, shuffle=False) + if dist_state.distributed else None + ) + val_loader = DataLoader( + val_ds, batch_size=train_cfg["batch_size"], sampler=val_sampler, + shuffle=False, num_workers=max(1, num_workers // 2), + collate_fn=collate_multiview, drop_last=False, pin_memory=True, + ) + + # ── model + optimizer ── + model = MVSam3DFlow( + model_channels=model_cfg["model_channels"], + cond_channels=model_cfg["cond_channels"], + num_blocks=model_cfg["num_blocks"], + num_heads=model_cfg["num_heads"], + max_views=max_views, + loss_weights=loss_cfg, + sigma_min=model_cfg["sigma_min"], + time_scale=model_cfg["time_scale"], + timestep_sampler=model_cfg["timestep_sampler"], + timestep_lognorm_mean=model_cfg["timestep_lognorm_mean"], + timestep_lognorm_std=model_cfg["timestep_lognorm_std"], + pretrained_ckpt=model_cfg["pretrained_ckpt"], + condition_kwargs=cond_cfg, + ).to(device) + model.train() + optimizer = torch.optim.AdamW( + [p for p in model.parameters() if p.requires_grad], + lr=optim_cfg["lr"], weight_decay=optim_cfg["weight_decay"], + ) + # LR schedule via the REUSED fork build_scheduler (warmup -> cosine -> min_lr, + # or constant). It reads attributes, so adapt the dict; note v2 keeps `lr` + # under optim and uses `steps` as the schedule total (-> max_steps). + sched_cfg = SimpleNamespace( + max_steps=train_cfg["steps"], warmup_steps=train_cfg["warmup_steps"], + min_lr=train_cfg["min_lr"], lr=optim_cfg["lr"], scheduler=train_cfg["scheduler"], + ) + scheduler = build_scheduler(optimizer, sched_cfg) + # DDP-wrap (or DeepSpeed) via the vendored runtime helper. Optimizer was built + # on the same Parameter objects pre-wrap (DDP does not reparametrize). + runtime = setup_runtime( + model=model, optimizer=optimizer, scheduler=scheduler, state=dist_state, + deepspeed_config_path=train_cfg["deepspeed_config"], + ) + model = runtime.model + optimizer = runtime.optimizer + scheduler = runtime.scheduler + trainable = [p for p in model.parameters() if p.requires_grad] + + steps = train_cfg["steps"] + log_every = train_cfg["log_every"] + ckpt_every = train_cfg["ckpt_every"] + ckpt_dir = train_cfg["ckpt_dir"] + grad_clip = optim_cfg["grad_clip"] + grad_accum = max(1, train_cfg["grad_accum_steps"]) + amp = train_cfg["amp"] + amp_dtype = train_cfg["amp_dtype"] + autocast_dt = torch.bfloat16 if str(amp_dtype).lower() == "bf16" else torch.float16 + use_scaler = amp and str(amp_dtype).lower() in {"fp16", "float16"} + scaler = torch.cuda.amp.GradScaler(enabled=use_scaler) + + # ── resume (REUSED fork load_checkpoint) ── + # Restore AFTER model/opt/scheduler/scaler are built and the model is wrapped: + # frozen condition weights were reconstructed from pretrained_ckpt at __init__, + # so the checkpoint (trainable params only) overwrites just the trained tensors. + # The scheduler's last_epoch is restored, so the LR curve continues unbroken. + start_step = 0 + resume_from = train_cfg["resume_from"] + if resume_from: + start_step, _ = load_checkpoint( + checkpoint_path=resume_from, model=model, optimizer=optimizer, + scheduler=scheduler, scaler=scaler, map_location="cpu", + ) + if is_main: + print(f"[resume] from {resume_from} at step {start_step}", flush=True) + + if is_main: + gbs = train_cfg["batch_size"] * grad_accum * dist_state.world_size + print(f"[setup] backend={dist_state.backend} world={dist_state.world_size} " + f"global_batch={gbs} steps={steps} grad_accum={grad_accum} " + f"amp={amp}/{amp_dtype} val_every={val_every}", flush=True) + + # ── training loop ── + optimizer.zero_grad(set_to_none=True) + step, micro, epoch = start_step, 0, 0 + if train_sampler is not None and hasattr(train_sampler, "set_epoch"): + train_sampler.set_epoch(epoch) + train_iter = iter(train_loader) + accum = {} + t0 = time.time() + ok = False + + try: + while step < steps: + try: + batch = next(train_iter) + except StopIteration: + epoch += 1 + if train_sampler is not None and hasattr(train_sampler, "set_epoch"): + train_sampler.set_epoch(epoch) + train_iter = iter(train_loader) + batch = next(train_iter) + batch = move_to(batch, device) + + autocast = (torch.autocast("cuda", dtype=autocast_dt) if amp else nullcontext()) + with autocast: + losses = model(batch) # DDP forward == flow_step + loss = losses["loss"] / grad_accum + + micro += 1 + is_boundary = (micro % grad_accum == 0) + # Skip DDP all-reduce on non-boundary micro-steps. + sync_ctx = nullcontext() if (is_boundary or not hasattr(model, "no_sync")) else model.no_sync() + with sync_ctx: + (scaler.scale(loss) if use_scaler else loss).backward() + + for k, v in losses.items(): + accum[k] = accum.get(k, 0.0) + float(v.detach()) + + if not is_boundary: + continue + + if grad_clip > 0: + if use_scaler: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(trainable, grad_clip) + if use_scaler: + scaler.step(optimizer) + scaler.update() + else: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + scheduler.step() # advance LR schedule + step += 1 + + if step % log_every == 0: + avg = {k: torch.tensor(v / (log_every * grad_accum)) for k, v in accum.items()} + reduced = reduce_scalar_dict(avg, dist_state) + accum.clear() + if is_main: + rate = log_every / (time.time() - t0) + lr = optimizer.param_groups[0]["lr"] + msg = " ".join(f"{k}={reduced[k]:.4f}" for k in sorted(reduced)) + print(f"[train] step {step} lr={lr:.6e} {msg} ({rate:.2f} it/s)", flush=True) + t0 = time.time() + + if val_loader is not None and step % val_every == 0: + metrics = run_validation( + model, val_loader, dist_state, + num_steps=train_cfg["val_num_steps"], + max_batches=train_cfg["val_max_batches"], + canonical_cache=data_cfg["canonical_cache"], + ss_decoder_ckpt=model_cfg["ss_decoder_ckpt"], + amp_dtype=amp_dtype, + ) + model.train() + if is_main: + msg = " ".join(f"{k}={metrics[k]:.4f}" for k in + ["shape_iou", "rot_deg", "trans_cm", "scale_rel", + "loss", "loss/shape", "loss/6drotation_normalized", + "loss/translation", "loss/scale"]) + print(f"[val] step {step} {msg} (n={int(metrics['val_samples'])})", flush=True) + + if step % ckpt_every == 0: + if is_main: + path = f"{ckpt_dir}/step_{step:07d}.pt" + # REUSED fork save_checkpoint: trainable params + optimizer + + # scheduler + scaler + step (frozen weights come from pretrained + # at construction, so they are intentionally not duplicated here). + save_checkpoint( + checkpoint_path=path, model=model, optimizer=optimizer, + scheduler=scheduler, scaler=scaler, step=step, config_dict=cfg, + ) + print(f"[train] saved {path}", flush=True) + barrier() + + if step >= steps: + break + ok = True + finally: + if is_main: + print(f"[train] {'done' if ok else 'TERMINATED (error)'}.", flush=True) + cleanup_distributed() + + +if __name__ == "__main__": + main() diff --git a/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/validate.py b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/validate.py new file mode 100644 index 0000000000000000000000000000000000000000..dec0b7370b2597d32b15ef93bbd45b441f424afa --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/mvsam3d/train/validate.py @@ -0,0 +1,181 @@ +"""In-training validation for the multi-view 6D-pose flow model. + +Runs the model's own rectified-flow sampler (``MVSam3DFlow.predict`` -> vendored +Euler solver) on a held-out split and reports four metrics, all in the same +metric space the GT was decoded from: + + * ``val/shape_iou`` — IoU(decode(pred shape latent) > 0, GT canonical occupancy). + The shape stream is pose-independent, so this is the + canonical voxel reconstruction quality. + * ``val/rot_deg`` — geodesic rotation error (vendored geodesic_distance_deg). + * ``val/trans_cm`` — L2 translation error in centimetres (metric space). + * ``val/scale_rel`` — |pred_scale - gt_scale| / gt_scale (metric extent). + +GT (R, t, s) is decoded with the SAME helpers as the prediction so the two are +compared in identical spaces; GT occupancy is read from the canonical cache by +object name. Metrics are summed with their sample counts and all-reduced across +DDP ranks, so the reported mean is exact regardless of per-rank batch counts. + +In addition to the sampler metrics, validation also reports the held-out **flow +loss** — the SAME per-stream v-prediction MSE used in training, including +``loss/shape`` (the validation shape loss) — from one ``flow_step`` per batch. +The noise/timestep draws are fixed by a seed (RNG saved/restored) so the loss is +comparable across passes and does not perturb the training RNG stream. + +The SS decoder is the REAL standalone SAM-3D ``ss_decoder.ckpt`` (the authoritative +SS latent space; see ss_vae.py), loaded once and cached. +""" +from __future__ import annotations + +import os +from typing import Dict + +import torch +import torch.distributed as dist + +from sam3d_objects.geometry.geometry import geodesic_distance_deg + +from mvsam3d.model.flow_model import ( + decode_pose_rotation, + decode_pose_scale, + decode_pose_translation, +) +from mvsam3d.model.ss_vae import decode_latent, load_ss_decoder + +# Module-level caches so repeated validation calls don't reload the decoder or +# re-read canonical occupancy grids from disk. +_DECODER = None +_OCC_CACHE: Dict[str, torch.Tensor] = {} + +# Fixed seed for the validation noise/timestep draws. The flow loss uses random +# t and z0, so a fixed seed makes the held-out loss (and the sampler metrics) +# comparable across validation passes; the training RNG is saved/restored around +# it so the training stream is unperturbed. +VAL_SEED = 1234 + +# The per-stream held-out flow losses reported alongside the sampler metrics +# (SAME objective as training: v-prediction MSE). "loss/shape" is the requested +# validation shape loss; the others come for free from one flow_step call. +_LOSS_KEYS = ["loss", "loss/shape", "loss/6drotation_normalized", + "loss/translation", "loss/scale"] + + +def _get_decoder(device, ckpt_path: str): + global _DECODER + if _DECODER is None: + _DECODER = load_ss_decoder(device=device, ckpt_path=ckpt_path) + return _DECODER + + +def _gt_occupancy(object_name: str, canonical_cache: str) -> torch.Tensor: + """GT canonical occupancy [64,64,64] (bool) for an object, cached.""" + if object_name not in _OCC_CACHE: + d = torch.load(os.path.join(canonical_cache, f"{object_name}.pt"), + map_location="cpu", weights_only=False) + _OCC_CACHE[object_name] = (d["occupancy"] > 0.5) + return _OCC_CACHE[object_name] + + +def _move(obj, device): + if torch.is_tensor(obj): + return obj.to(device, non_blocking=True) + if isinstance(obj, dict): + return {k: _move(v, device) for k, v in obj.items()} + if isinstance(obj, list): + return [_move(v, device) for v in obj] + return obj + + +@torch.no_grad() +def run_validation(model, val_loader, dist_state, *, num_steps: int, + max_batches: int, canonical_cache: str, ss_decoder_ckpt: str, + amp_dtype: str = "bf16") -> Dict[str, float]: + """Sample + decode on up to ``max_batches`` val batches; return reduced means.""" + inner = model.module if hasattr(model, "module") else model + was_training = inner.training + inner.eval() + device = dist_state.device + decoder = _get_decoder(device, ss_decoder_ckpt) + + # Save the training RNG, then seed deterministically for this pass. + cpu_rng = torch.get_rng_state() + cuda_rng = torch.cuda.get_rng_state_all() if device.type == "cuda" else None + torch.manual_seed(VAL_SEED) + + use_amp = (device.type == "cuda" and amp_dtype.lower() in {"bf16", "fp16", "float16"}) + amp_dt = torch.bfloat16 if amp_dtype.lower() == "bf16" else torch.float16 + autocast = (torch.autocast(device_type="cuda", dtype=amp_dt) + if use_amp else torch.autocast(device_type="cpu", enabled=False)) + + # Running sums + sample count (for an exact cross-rank mean). Sampler metrics + # and the held-out flow losses are all accumulated B-weighted, then divided by + # the reduced total count, so the mean is exact regardless of per-rank batches. + sums = {"shape_iou": 0.0, "rot_deg": 0.0, "trans_cm": 0.0, "scale_rel": 0.0} + sums.update({k: 0.0 for k in _LOSS_KEYS}) + n = 0 + # The body runs under the fixed VAL_SEED; the finally restores the training + # RNG stream and train() mode on EVERY exit path (including an exception in + # the val loop), so a transient validation failure can never leave the + # training RNG pinned to VAL_SEED or the model stuck in eval(). + try: + for bi, batch in enumerate(val_loader): + if bi >= max_batches: + break + batch = _move(batch, device) + with autocast: + pred = inner.predict(batch, num_steps=num_steps) + + B = pred["rotation"].shape[0] + + # ---- GT decoded with the same helpers (identical spaces) ---- + gt_R = decode_pose_rotation(batch["rot6d"].float()) # [B,3,3] + gt_t = decode_pose_translation( + batch["translation"].float(), batch["pm_scale"].float(), batch["pm_shift"].float() + ) # [B,3] metric + gt_s = decode_pose_scale(batch["scale"].float(), batch["pm_scale"].float())[:, 0] # [B] + + # ---- pose metrics ---- + rot_deg = geodesic_distance_deg(pred["rotation"].float(), gt_R).reshape(B) + trans_cm = (pred["translation"].float() - gt_t).norm(dim=-1) * 100.0 # m -> cm + scale_rel = (pred["scale"].float() - gt_s).abs() / gt_s.abs().clamp_min(1e-6) + + # ---- shape IoU: decode pred latent -> occupancy vs GT canonical occ ---- + logits = decode_latent(decoder, pred["shape_latent"].float()) # [B,1,64,64,64] + pred_occ = (logits[:, 0] > 0.0) # [B,64,64,64] + for b in range(B): + gt_occ = _gt_occupancy(batch["object_name"][b], canonical_cache).to(device) + inter = (pred_occ[b] & gt_occ).sum().item() + union = (pred_occ[b] | gt_occ).sum().item() + sums["shape_iou"] += (inter / union) if union > 0 else 1.0 + + sums["rot_deg"] += float(rot_deg.sum()) + sums["trans_cm"] += float(trans_cm.sum()) + sums["scale_rel"] += float(scale_rel.sum()) + + # ---- held-out flow loss: SAME objective as training (per-stream + # v-prediction MSE, incl. the shape loss). Cheap extra forward; the + # model is in eval() so fuser modality dropout is off (deterministic). + # flow_step returns batch-mean scalars -> multiply by B for an exact + # B-weighted sum. + with autocast: + losses = inner.flow_step(batch) + for lk in _LOSS_KEYS: + sums[lk] += float(losses[lk]) * B + n += B + + # ---- exact cross-rank reduction: all_reduce SUM of (sums..., count) ---- + keys = ["shape_iou", "rot_deg", "trans_cm", "scale_rel"] + _LOSS_KEYS + packed = torch.tensor([sums[k] for k in keys] + [float(n)], device=device) + if dist_state.distributed: + dist.all_reduce(packed, op=dist.ReduceOp.SUM) + total_n = packed[-1].item() + out = {k: (packed[i].item() / total_n if total_n > 0 else float("nan")) + for i, k in enumerate(keys)} + out["val_samples"] = total_n + return out + finally: + if cuda_rng is not None: + torch.cuda.set_rng_state_all(cuda_rng) + torch.set_rng_state(cpu_rng) + if was_training: + inner.train() diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8506f76cdf80cbe81b80953966a3cd1f7f560a67 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/data/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/data/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/data/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e7785e5891ca3829ed5b186d98637bed30a8a06 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/data/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/data/__pycache__/utils.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/data/__pycache__/utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36ed3783b3b81034df21e8f5d3b8a60fa8aab51c Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/data/__pycache__/utils.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/data/utils.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/data/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e990bafbd803b292124f18ddb5c9cb07e06fdbcb --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/data/utils.py @@ -0,0 +1,242 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from typing import Any, Iterable, Tuple, Union, Dict, Sequence, Mapping, Container +import optree +import torch +from collections.abc import Iterable +import inspect +import ast +import astor +from torch.utils import _pytree + +# None = root, Iterable[Any] = path, Any = path of one +ChildPathType = Union[None, Iterable[Any], Any] +ArgsType = Iterable[ChildPathType] +KwargsType = Mapping[str, ChildPathType] +ArgsKwargsType = Tuple[ArgsType, KwargsType] +MappingType = Union[None, ArgsKwargsType, ArgsType, KwargsType] + + +def tree_transpose_level_one( + structure, + check_children=False, + map_fn=None, + is_leaf=None, +): + _, outer_spec = optree.tree_flatten( + structure, + is_leaf=lambda x: x is not structure, + none_is_leaf=True, + ) + + spec = optree.tree_structure(structure, none_is_leaf=True, is_leaf=is_leaf) + children_spec = spec.children() + if len(children_spec) > 0: + inner_spec = children_spec[0] + if check_children: + for child_spec in children_spec[1:]: + assert ( + inner_spec == child_spec + ), f"one child was found having a different tree structure ({inner_spec} != {child_spec})" + + structure = optree.tree_transpose(outer_spec, inner_spec, structure) + + if map_fn is not None: + structure = optree.tree_map( + map_fn, + structure, + is_leaf=lambda x: optree.tree_structure( + x, is_leaf=is_leaf, none_is_leaf=True + ) + == outer_spec, + none_is_leaf=True, + ) + + return structure + + +def tree_tensor_map(fn, tree, *rest): + return optree.tree_map( + fn, + tree, + *rest, + is_leaf=lambda x: isinstance(x, torch.Tensor), + none_is_leaf=False, + ) + + +def to_device(obj, device): + """Recursively moves all tensors in obj to the specified device. + + Args: + obj: Object to move to device - can be a tensor, list, tuple, dict or any nested combination + device: Target device (e.g. 'cuda', 'cpu', torch.device('cuda:0') etc.) + + Returns: + Same object structure with all contained tensors moved to specified device + """ + to_fn = lambda x: x.to(device) + return optree.tree_map(to_fn, obj, is_leaf=torch.is_tensor, none_is_leaf=False) + + +def expand_right(tensor, target_shape): + """ + e.g. Takes tensor of (a, b, c) and returns a tensor of (a, b, c, 1, 1, ...) + """ + current_shape = tensor.shape + dims_to_add = len(target_shape) - len(current_shape) + result = tensor + for _ in range(dims_to_add): + result = result.unsqueeze(-1) + expand_shape = list(current_shape) + [-1] * dims_to_add + for i in range(len(target_shape)): + if i < len(expand_shape) and expand_shape[i] == -1: + expand_shape[i] = target_shape[i] + return result.expand(*expand_shape) + + +def expand_as_right(tensor, target): + return expand_right(tensor, target.shape) + + +def as_keys(path: ChildPathType): + if isinstance(path, Iterable) and (not isinstance(path, str)): + return tuple(path) + elif path is None: + return () + return (path,) + + +def get_child(obj: Any, *keys: Iterable[Any]): + for key in keys: + obj = obj[key] + return obj + + +def set_child(obj: Any, value: Any, *keys: Iterable[Any]): + parent = None + for key in keys: + parent = obj + obj = obj[key] + if parent is None: + obj = value + else: + parent[key] = value + return obj + + +def build_args_batch_extractor(args_mapping: ArgsType): + def extract_fn(batch): + return tuple(get_child(batch, *as_keys(path)) for path in args_mapping) + + return extract_fn + + +def build_kwargs_batch_extractor(kwargs_mapping: KwargsType): + def extract_fn(batch): + return { + name: get_child(batch, *as_keys(path)) + for name, path in kwargs_mapping.items() + } + + return extract_fn + + +empty_mapping = object() +kwargs_identity_mapping = object() + + +def build_batch_extractor(mapping: MappingType): + extract_args_fn = lambda x: () + extract_kwargs_fn = lambda x: {} + + if mapping is None: + + def extract_args_fn(batch): + return (batch,) + + elif mapping is empty_mapping: + pass + elif mapping is kwargs_identity_mapping: + extract_kwargs_fn = lambda x: x + elif isinstance(mapping, Sequence) and (not isinstance(mapping, str)): + if ( + len(mapping) == 2 + and isinstance(mapping[0], Sequence) + and isinstance(mapping[1], Dict) + ): + extract_args_fn = build_args_batch_extractor(mapping[0]) + extract_kwargs_fn = build_kwargs_batch_extractor(mapping[1]) + else: + extract_args_fn = build_args_batch_extractor(mapping) + elif isinstance(mapping, Mapping): + extract_kwargs_fn = build_kwargs_batch_extractor(mapping) + else: + + def extract_args_fn(batch): + return (get_child(batch, *as_keys(mapping)),) + + def extract_fn(batch): + return extract_args_fn(batch), extract_kwargs_fn(batch) + + return extract_fn + + +# > + + +def right_broadcasting(arr, target): + return arr.reshape(arr.shape + (1,) * (target.ndim - arr.ndim)) + + +def get_stats(tensor: torch.Tensor): + float_tensor = tensor.float() + return { + "shape": tuple(tensor.shape), + "min": tensor.min().item(), + "max": tensor.max().item(), + "mean": float_tensor.mean().item(), + "median": tensor.median().item(), + "std": float_tensor.std().item(), + } + + +def _get_caller_arg_name(argnum=0, parent_frame=1): + try: + frame = inspect.currentframe() # current frame + frame = inspect.getouterframes(frame)[1 + parent_frame] # parent frame + code = inspect.getframeinfo(frame[0]).code_context[0].strip() # get code line + + tree = ast.parse(code) + + for node in ast.walk(tree): + if isinstance(node, ast.Call): + args = node.args + break # only get the first parent call + + # get first argument string (do not handle '=') + label = astor.to_source(args[argnum]).strip() + except: + # TODO(Pierre) log exception + label = "{label}" + return label + + +def print_stats(tensor, label=None): + if label is None: + label = _get_caller_arg_name(argnum=0) + stats = get_stats(tensor) + string = f"{label}:\n" + "\n".join(f"- {k}: {v}" for k, v in stats.items()) + print(string) + + +def tree_reduce_unique(fn, tree, ensure_unique=True, **kwargs): + values = _pytree.tree_flatten(tree, **kwargs)[0] + values = tuple(map(fn, values)) + first = values[0] + if ensure_unique: + for value in values[1:]: + if value != first: + raise RuntimeError( + f"different values found, {value} and {first} should be the same" + ) + return first diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5fd7636eac75e90f537f3db4ac1ad624eb9c1cdc Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/__pycache__/geometry.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/__pycache__/geometry.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f6b0cb68479d001620fef02ea5704973933431c Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/__pycache__/geometry.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/__pycache__/rotation_constants.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/__pycache__/rotation_constants.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f5282bd16567315a5ee5d8c382e4bc17bb4aa05b Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/__pycache__/rotation_constants.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/geometry.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..70fc75cfcb37cf8694b23c104518a09ac58d97e1 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/geometry.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import math + +import torch +from pytorch3d.transforms import axis_angle_to_matrix + +from .rotation_constants import ROTATION_6D_MEAN, ROTATION_6D_STD + + +# ── 6D rotation: COLUMN convention ────────────────────────────────── +# The pretrained SAM-3D pose_decoder uses column-based Gram-Schmidt: +# encode: 6D = [R[:,0], R[:,1]] (first two COLUMNS of R) +# decode: torch.stack([b1, b2, b3], dim=-1) → columns +# +# pytorch3d uses ROW convention (dim=-2) — INCOMPATIBLE with pose_decoder. +# Do NOT use pytorch3d matrix_to_rotation_6d / rotation_6d_to_matrix here. + + +def matrix_to_rot6d(R: torch.Tensor) -> torch.Tensor: + """Rotation matrix → 6D (column convention, matching pose_decoder).""" + return torch.cat([R[..., :, 0], R[..., :, 1]], dim=-1) + + +def rot6d_to_matrix(rot6d: torch.Tensor) -> torch.Tensor: + """6D → rotation matrix (column Gram-Schmidt, matching pose_decoder).""" + a1 = rot6d[..., 0:3] + a2 = rot6d[..., 3:6] + b1 = torch.nn.functional.normalize(a1, dim=-1) + b2 = a2 - torch.sum(b1 * a2, dim=-1, keepdim=True) * b1 + b2 = torch.nn.functional.normalize(b2, dim=-1) + b3 = torch.cross(b1, b2, dim=-1) + return torch.stack([b1, b2, b3], dim=-1) + + +def normalize_rot6d(raw_rot6d: torch.Tensor) -> torch.Tensor: + """raw 6D → normalized 6D (matching pretrained backbone convention).""" + mean = ROTATION_6D_MEAN.to(device=raw_rot6d.device, dtype=raw_rot6d.dtype) + std = ROTATION_6D_STD.to(device=raw_rot6d.device, dtype=raw_rot6d.dtype) + return (raw_rot6d - mean) / std + + +def denormalize_rot6d(norm_rot6d: torch.Tensor) -> torch.Tensor: + """normalized 6D → raw 6D (inverse of normalize_rot6d).""" + mean = ROTATION_6D_MEAN.to(device=norm_rot6d.device, dtype=norm_rot6d.dtype) + std = ROTATION_6D_STD.to(device=norm_rot6d.device, dtype=norm_rot6d.dtype) + return norm_rot6d * std + mean + + +def ssi_to_metric( + translation_ssi: torch.Tensor, + pointmap_scale: torch.Tensor, + pointmap_shift: torch.Tensor, +) -> torch.Tensor: + return pointmap_scale * translation_ssi + pointmap_shift + + + +def geodesic_distance_deg(pred_rotation: torch.Tensor, gt_rotation: torch.Tensor) -> torch.Tensor: + rot = pred_rotation @ gt_rotation.transpose(-1, -2) + trace = torch.diagonal(rot, dim1=-2, dim2=-1).sum(dim=-1) + cos = torch.clamp((trace - 1.0) * 0.5, -1.0, 1.0) + return torch.acos(cos) * (180.0 / math.pi) diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/rotation_constants.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/rotation_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..caa1d6bc2a901b27d0cd890b99052f0cf190d85e --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/geometry/rotation_constants.py @@ -0,0 +1,32 @@ +"""Rotation 6D normalization constants. + +These match the ROTATION_6D_MEAN / ROTATION_6D_STD in +sam3d_objects.pipeline.inference_utils (pretrained SAM-3D checkpoint). +Duplicated here to avoid importing the heavy inference module (open3d, trimesh, etc.). +""" + +from __future__ import annotations + +import torch + +ROTATION_6D_MEAN = torch.tensor( + [ + -0.06366084883674913, + 0.008438224692279752, + 0.00017084786438302483, + 0.0007126610473540038, + -0.0030916726538816417, + 0.5166093753457688, + ] +) + +ROTATION_6D_STD = torch.tensor( + [ + 0.6656971967514863, + 0.6787012271867754, + 0.30345010594844524, + 0.4394504420678794, + 0.39817973931717104, + 0.6176286868761914, + ] +) diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a2eb24da3007d44d6484f58078c93ac209859b6a Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b75a61a4e72ab8659ad7bb95061445bc98ed5cfa Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14894de2718730104dd584b46747e44d63f65648 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dca9ed36b190cb08133f3ee6e59c4235cfca9608 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/dino.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/dino.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f6f543e4d7660dd675b9a3d73de405a7a5a12937 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/dino.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/embedder_fuser.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/embedder_fuser.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b058d29e788bc4319b6f7aa5f410cfcbaf68c2c6 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/embedder_fuser.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/point_remapper.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/point_remapper.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..061b879aa3410d16497cdf679aabe1d94b9b033c Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/point_remapper.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/pointmap.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/pointmap.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1609b9a5aa98c57419721bc97ecdf8f4e28b03d0 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/__pycache__/pointmap.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/dino.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/dino.py new file mode 100644 index 0000000000000000000000000000000000000000..909349a339570e1bfe501d6bf4c620a1b3876831 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/dino.py @@ -0,0 +1,149 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import torch +from typing import Optional, Dict, Any +import warnings +from torchvision.transforms import Normalize +import torch.nn.functional as F +from loguru import logger +import os + + +class Dino(torch.nn.Module): + def __init__( + self, + input_size: int = 224, + repo_or_dir: str = "facebookresearch/dinov2", + dino_model: str = "dinov2_vitb14", + source: str = "github", + backbone_kwargs: Optional[Dict[str, Any]] = None, + normalize_images: bool = True, + # for backward compatible + prenorm_features: bool = False, + freeze_backbone: bool = True, + prune_network: bool = False, # False for backward compatible + ): + super().__init__() + if backbone_kwargs is None: + backbone_kwargs = {} + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + + logger.info(f"Loading DINO model: {dino_model} from {repo_or_dir} (source: {source})") + # SLURM server path rewrite: hub cache was originally created under x3135a07's home; + # remap to current user's scratch on this cluster node. + # TODO: replace with a proper env-var or config-driven hub cache path. + current = os.getcwd() + if "x3302a04" in current: + repo_or_dir = repo_or_dir.replace("x3135a07", "x3302a04") + if backbone_kwargs: + logger.info(f"DINO backbone kwargs: {backbone_kwargs}") + + self.backbone = torch.hub.load( + repo_or_dir=repo_or_dir, + model=dino_model, + source=source, + verbose=False, + **backbone_kwargs, + ) + + # Log model properties after loading + logger.info(f"Loaded DINO model - type: {type(self.backbone)}, " + f"embed_dim: {self.backbone.embed_dim}, " + f"patch_size: {getattr(self.backbone.patch_embed, 'patch_size', 'N/A')}") + + + self.resize_input_size = (input_size, input_size) + self.embed_dim = self.backbone.embed_dim + self.input_size = input_size + self.input_channels = 3 + self.normalize_images = normalize_images + self.prenorm_features = prenorm_features + self.register_buffer('mean', torch.as_tensor([[0.485, 0.456, 0.406]]).view(-1, 1, 1), persistent=False) + self.register_buffer('std', torch.as_tensor([[0.229, 0.224, 0.225]]).view(-1, 1, 1), persistent=False) + + # freeze + if freeze_backbone: + self.requires_grad_(False) + self.eval() + elif not prune_network: + logger.warning( + "Unfreeze encoder w/o prune parameter may lead to error in ddp/fp16 training" + ) + + if prune_network: + self._prune_network() + + def _preprocess_input(self, x): + _resized_images = torch.nn.functional.interpolate( + x, + size=self.resize_input_size, + mode="bilinear", + align_corners=False, + ) + + if x.shape[1] == 1: + _resized_images = _resized_images.repeat(1, 3, 1, 1) + + if self.normalize_images: + _resized_images = _resized_images.sub_(self.mean).div_(self.std) + + return _resized_images + + def _forward_intermediate_layers( + self, input_img, intermediate_layers, cls_token=True + ): + return self.backbone.get_intermediate_layers( + input_img, + intermediate_layers, + return_class_token=cls_token, + ) + + def _forward_last_layer(self, input_img): + output = self.backbone.forward_features(input_img) + if self.prenorm_features: + features = output["x_prenorm"] + tokens = F.layer_norm(features, features.shape[-1:]) + else: + tokens = torch.cat( + [ + output["x_norm_clstoken"].unsqueeze(1), + output["x_norm_patchtokens"], + ], + dim=1, + ) + return tokens + + def forward(self, x, **kwargs): + _resized_images = self._preprocess_input(x) + tokens = self._forward_last_layer(_resized_images) + return tokens.to(x.dtype) + + def _prune_network(self): + """ + Ran this script: + out = model(input) + loss = out.sum() + loss.backward() + + for name, p in dino_model.named_parameters(): + if p.grad is None: + print(name) + model.zero_grad() + """ + self.backbone.mask_token = None + if self.prenorm_features: + self.backbone.norm = torch.nn.Identity() + + +class DinoForMasks(torch.nn.Module): + def __init__( + self, + backbone: Dino, + ): + super().__init__() + self.backbone = backbone + self.embed_dim = self.backbone.embed_dim + + def forward(self, image, mask): + return self.backbone.forward(mask) diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/embedder_fuser.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/embedder_fuser.py new file mode 100644 index 0000000000000000000000000000000000000000..a8b6130420f0716ff864e9e78fd9b164ba1e9da4 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/embedder_fuser.py @@ -0,0 +1,238 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import math +import torch +from loguru import logger +from torch import nn +from typing import Optional, Tuple, List, Literal, Dict +from sam3d_objects.model.layers.llama3.ff import FeedForward +from omegaconf import OmegaConf + +class EmbedderFuser(torch.nn.Module): + """ + Fusing individual condition embedder. Require kwargs for the forward! + Args: + embedder_list: List of Tuples. Each Tuple consists of a condition_embedder + and a list of tuple. In the list, each tuple consists of a string, indicating + a kward, and astring, indicating the group of positional encoding to be used. + use_pos_embedding: whether to add positional embedding. If add, follow the index in + embedder_list. Choices of None (no pos emb), random, and learned. + projection_pre_norm: pre-normalize features before feeding into projector layers. + projection_net_hidden_dim_multiplier: hidden dimension for projection layer. If 0, don't use. + """ + + def __init__( + self, + embedder_list: List[Tuple[nn.Module, List[Tuple[str, Optional[str]]]]], + use_pos_embedding: Optional[Literal["random", "learned"]] = "learned", + projection_pre_norm: bool = True, + projection_net_hidden_dim_multiplier: float = 4.0, + compression_projection_multiplier: float = 0, + freeze: bool = False, + drop_modalities_weight: Dict[List[str], float] = None, + dropout_prob: float = 0.0, + force_drop_modalities: List[str] = None, + ): + super().__init__() + # torch.compile does not support OmegaConf.ListConfig, so we convert to a list + if not isinstance(embedder_list, List): + self.embedder_list = OmegaConf.to_container(embedder_list) + else: + self.embedder_list = embedder_list + + self.embed_dims = 0 + self.compression_projection_multiplier = compression_projection_multiplier + self.concate_embed_dims = 0 + # keep moduleList to be compatible with nn module + self.module_list = [] + max_positional_embed_idx = 0 + self.positional_embed_map = {} + for condition_embedder, kwargs_info in self.embedder_list: + self.embed_dims = max(self.embed_dims, condition_embedder.embed_dim) + self.module_list.append(condition_embedder) + for _, pos_group in kwargs_info: + self.concate_embed_dims += condition_embedder.embed_dim + if pos_group is not None: + if pos_group not in self.positional_embed_map: + self.positional_embed_map[pos_group] = max_positional_embed_idx + max_positional_embed_idx += 1 + self.module_list = nn.ModuleList(self.module_list) + self.use_pos_embedding = use_pos_embedding + if self.use_pos_embedding == "random": + idx_emb = torch.randn(max_positional_embed_idx + 1, 1, self.embed_dims) + self.register_buffer("idx_emb", idx_emb) + elif self.use_pos_embedding == "learned": + self.idx_emb = nn.Parameter( + torch.empty(max_positional_embed_idx + 1, self.embed_dims) + ) + nn.init.normal_( + self.idx_emb, mean=0.0, std=1.0 / math.sqrt(self.embed_dims) + ) + else: + raise NotImplementedError(f"Unknown pos embedding {self.use_pos_embedding}") + + self.projection_pre_norm = projection_pre_norm + self.projection_net_hidden_dim_multiplier = projection_net_hidden_dim_multiplier + if projection_net_hidden_dim_multiplier > 0: + self.projection_nets = [] + for condition_embedder, _ in self.embedder_list: + self.projection_nets.append( + self._make_projection_net( + condition_embedder.embed_dim, + self.embed_dims, + self.projection_net_hidden_dim_multiplier, + ) + ) + self.projection_nets = nn.ModuleList(self.projection_nets) + + if compression_projection_multiplier > 0: + self.compression_projector = self._make_projection_net( + self.concate_embed_dims, + self.embed_dims, + self.compression_projection_multiplier, + ) + + self.drop_modalities_weight = drop_modalities_weight if drop_modalities_weight is not None else [] + self.dropout_prob = dropout_prob + self.force_drop_modalities = force_drop_modalities + + if freeze: + self.requires_grad_(False) + self.eval() + + def _make_projection_net( + self, + input_embed_dim, + output_embed_dim: int, + multiplier: int, + ): + if self.projection_pre_norm: + pre_norm = nn.LayerNorm(input_embed_dim) + else: + pre_norm = nn.Identity() + + # Per-token projection + gated activation + ff_net = FeedForward( + dim=input_embed_dim, + hidden_dim=int(multiplier * output_embed_dim), + output_dim=output_embed_dim, + ) + return nn.Sequential(pre_norm, ff_net) + + def _build_dropout_distribution(self, device): + """ + Build the probability distribution for dropout configurations. + + Returns: + dropout_configs: List of sets containing modalities to drop + cumsum_weights: Cumulative sum of weights for sampling + """ + dropout_configs = [] + weights = [] + + # Add no-dropout configuration with remaining probability + dropout_configs.append(set()) + weights.append(1.0 - self.dropout_prob) + + # Add configured dropout patterns + total_dropout_weight = sum(w for _, w in self.drop_modalities_weight) + assert total_dropout_weight > 0, "Total dropout weight must be positive when drop_modalities_weight is provided" + for modality_list, weight in self.drop_modalities_weight: + dropout_configs.append(set(modality_list)) + # Scale weight by dropout_prob to ensure total probability sums to 1 + weights.append(self.dropout_prob * weight / total_dropout_weight) + + # Convert weights to cumulative distribution + weights_tensor = torch.tensor(weights, device=device) + + was_deterministic = torch.are_deterministic_algorithms_enabled() + torch.use_deterministic_algorithms(False) + cumsum_weights = torch.cumsum(weights_tensor, dim=0) + torch.use_deterministic_algorithms(was_deterministic) + + return dropout_configs, cumsum_weights + + def _apply_force_drop(self, kwarg_names: List[str], tokens: List[torch.Tensor]): + if not self.force_drop_modalities: + return tokens + + force_drop_set = set(self.force_drop_modalities) + result_tokens = [] + + for kwarg_name, token_tensor in zip(kwarg_names, tokens): + # Create mask: 0 for forced drop, 1 otherwise + mask = 0.0 if kwarg_name in force_drop_set else 1.0 + result_tokens.append(token_tensor * mask) + + return result_tokens + + def _dropout_modalities(self, kwarg_names: List[str], tokens: List[torch.Tensor]): + # First apply forced drops (deterministic, always applied) + tokens = self._apply_force_drop(kwarg_names, tokens) + + # Then apply probabilistic dropout (only in training) + if not self.training or self.dropout_prob <= 0 or not self.drop_modalities_weight: + return tokens + + batch_size = tokens[0].shape[0] + device = tokens[0].device + + # Build dropout configurations and sample which to use per batch element + dropout_configs, cumsum_weights = self._build_dropout_distribution(device) + rand_vals = torch.rand(batch_size, device=device) + # Clamp indices to valid range (handle edge case where rand_val == 1.0) + config_indices = torch.searchsorted(cumsum_weights, rand_vals).clamp(max=len(dropout_configs) - 1) + + # Apply dropout masks with vectorized operations + result_tokens = [] + for kwarg_name, token_tensor in zip(kwarg_names, tokens): + # Start with all ones (no dropout) + mask = torch.ones(batch_size, dtype=token_tensor.dtype, device=device) + + # Vectorized mask creation: check all configurations at once + for config_idx, modalities_to_drop in enumerate(dropout_configs): + if kwarg_name in modalities_to_drop: + # Set mask to 0 for all batch elements using this configuration + mask[config_indices == config_idx] = 0.0 + + # Reshape mask to match token dimensions + mask = mask.view([batch_size] + [1] * (token_tensor.ndim - 1)) + result_tokens.append(token_tensor * mask) + + return result_tokens + + def forward(self, *args, **kwargs): + tokens = [] + kwarg_names = [] + + for i, (condition_embedder, kwargs_info) in enumerate(self.embedder_list): + # Ideally, we would batch the inputs; but that assumes same-sized inputs + for kwarg_name, pos_group in kwargs_info: + if kwarg_name not in kwargs: + logger.warning(f"{kwarg_name} not in kwargs to condition embedder!") + input_cond = kwargs[kwarg_name] + cond_token = condition_embedder(input_cond) + if self.projection_net_hidden_dim_multiplier > 0: + cond_token = self.projection_nets[i](cond_token) + if pos_group is not None: + pos_idx = self.positional_embed_map[pos_group] + if self.use_pos_embedding == "random": + cond_token += self.idx_emb[pos_idx : pos_idx + 1] + elif self.use_pos_embedding == "learned": + cond_token += self.idx_emb[pos_idx : pos_idx + 1, None] + else: + raise NotImplementedError( + f"Unknown pos embedding {self.use_pos_embedding}" + ) + tokens.append(cond_token) + kwarg_names.append(kwarg_name) + + # Apply dropout modalities with preserved order + tokens = self._dropout_modalities(kwarg_names, tokens) + + if self.compression_projection_multiplier > 0: + tokens = torch.cat(tokens, dim=-1) + tokens = self.compression_projector(tokens) + else: + tokens = torch.cat(tokens, dim=1) + + return tokens diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/point_remapper.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/point_remapper.py new file mode 100644 index 0000000000000000000000000000000000000000..f971de1afc47eb3d726dff870c78a9809c17540c --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/point_remapper.py @@ -0,0 +1,78 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import torch +import torch.nn as nn + + +class PointRemapper(nn.Module): + """Handles remapping of 3D point coordinates and their inverse transformations.""" + + VALID_TYPES = ["linear", "sinh", "exp", "sinh_exp", "exp_disparity"] + + def __init__(self, remap_type: str = "exp"): + super().__init__() + self.remap_type = remap_type + + if remap_type not in self.VALID_TYPES: + raise ValueError( + f"Invalid remap type: {remap_type}. Must be one of {self.VALID_TYPES}" + ) + + def forward(self, points: torch.Tensor) -> torch.Tensor: + """Apply remapping to point coordinates.""" + if self.remap_type == "linear": + return points + + elif self.remap_type == "sinh": + return torch.asinh(points) + + elif self.remap_type == "exp": + xy_scaled, z_exp = points.split([2, 1], dim=-1) + # Use log1p for better numerical stability near zero + z = torch.log1p(z_exp) + xy = xy_scaled / (1 + z_exp) + return torch.cat([xy, z], dim=-1) + + elif self.remap_type == "exp_disparity": + xy_scaled, z_exp = points.split([2, 1], dim=-1) + xy = xy_scaled / z_exp + z = torch.log(z_exp) + return torch.cat([xy, z], dim=-1) + + elif self.remap_type == "sinh_exp": + xy_sinh, z_exp = points.split([2, 1], dim=-1) + xy = torch.asinh(xy_sinh) + z = torch.log(z_exp.clamp(min=1e-8)) + return torch.cat([xy, z], dim=-1) + + else: + raise ValueError(f"Unknown remap type: {self.remap_type}") + + def inverse(self, points: torch.Tensor) -> torch.Tensor: + """Apply inverse remapping to recover original point coordinates.""" + if self.remap_type == "linear": + return points + + elif self.remap_type == "sinh": + return torch.sinh(points) + + elif self.remap_type == "exp": + xy, z = points.split([2, 1], dim=-1) + # Inverse of log1p is expm1(z) = exp(z) - 1 + z_exp = torch.expm1(z) + # Inverse of xy/(1+z_exp) is xy*(1+z_exp) + return torch.cat([xy * (1 + z_exp), z_exp], dim=-1) + + elif self.remap_type == "exp_disparity": + xy, z = points.split([2, 1], dim=-1) + z_exp = torch.exp(z) + return torch.cat([xy * z_exp, z_exp], dim=-1) + + elif self.remap_type == "sinh_exp": + xy, z = points.split([2, 1], dim=-1) + return torch.cat([torch.sinh(xy), torch.exp(z)], dim=-1) + + else: + raise ValueError(f"Unknown remap type: {self.remap_type}") + + def extra_repr(self) -> str: + return f"remap_type='{self.remap_type}'" diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/pointmap.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/pointmap.py new file mode 100644 index 0000000000000000000000000000000000000000..62328b5d9c708ba58d48993eb96eeeac804dc32c --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/dit/embedder/pointmap.py @@ -0,0 +1,238 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from timm.models.vision_transformer import Block +import torch +from torch import nn +import torch.nn.functional as F +from functools import partial +from loguru import logger + +from .point_remapper import PointRemapper + + +class PointPatchEmbed(nn.Module): + """ + Projects (x,y,z) → D + Splits into patches (patch_size x patch_size) + Runs a tiny self-attention block inside each window + Returns one token per window. + """ + + def __init__( + self, + input_size: int = 256, + patch_size: int = 8, + embed_dim: int = 768, + remap_output: str = "exp", # Add remap_output parameter + dropout_prob: float = 0.0, # Dropout probability for pointmap + force_dropout_always: bool = False, # Force dropout during validation/inference + ): + super().__init__() + self.input_size = input_size + self.patch_size = patch_size + self.embed_dim = embed_dim + self.dropout_prob = dropout_prob + self.force_dropout_always = force_dropout_always + + # Point remapper + self.point_remapper = PointRemapper(remap_output) + + # (1) point embedding + self.point_proj = nn.Linear(3, embed_dim) + self.invalid_xyz_token = nn.Parameter(torch.zeros(embed_dim)) + + # Special embedding for dropped patches (used during dropout) + # Alternative dropout strategies to consider: + # 1. Drop all tokens entirely or use a single token only + # 2. Different dropout patterns per window + # 3. Use dropped_xyz_token/invalid_xyz_token per pixel + if dropout_prob > 0: + self.dropped_xyz_token = nn.Parameter(torch.zeros(embed_dim)) + + # (2) positional embedding + num_patches = input_size // patch_size + # For patches + self.pos_embed = nn.Parameter( + torch.zeros(1, embed_dim, num_patches, num_patches) + ) + # For points in a patch + self.pos_embed_window = nn.Parameter( + torch.zeros(1, 1 + patch_size * patch_size, embed_dim) + ) + self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + + # (3) within-patch transformer block(s) + # From MCC: https://github.com/facebookresearch/MCC/blob/b04c97518360e4fdedfb6f090db7e90d0c2f8ae6/mcc_model.py#L97 + self.blocks = nn.ModuleList( + [ + Block( + embed_dim, + num_heads=16, + mlp_ratio=2.0, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + ) + ] + ) + self.initialize_weights() + + def initialize_weights(self): + # Initialize positional embeddings with small std + nn.init.normal_(self.pos_embed, std=0.02) + nn.init.normal_(self.pos_embed_window, std=0.02) + nn.init.normal_(self.cls_token, std=0.02) + nn.init.normal_(self.invalid_xyz_token, std=0.02) + + # Initialize dropped pointmap token if dropout is enabled + if self.dropout_prob > 0: + nn.init.normal_(self.dropped_xyz_token, std=0.02) + + # Initialize point projection with xavier uniform for better gradient flow + # This is crucial since pointmaps can have large value ranges + nn.init.xavier_uniform_(self.point_proj.weight, gain=0.02) + if self.point_proj.bias is not None: + nn.init.constant_(self.point_proj.bias, 0) + + def _get_pos_embed(self, hw): + h, w = hw + pos_embed = F.interpolate( + self.pos_embed, size=(h, w), mode="bilinear", align_corners=False + ) + pos_embed = pos_embed.permute(0, 2, 3, 1) # (B, H, W, C) + return pos_embed + + def resize_input(self, xyz: torch.Tensor) -> torch.Tensor: + resized_xyz = F.interpolate(xyz, size=self.input_size, mode="nearest") + resized_xyz = resized_xyz.permute(0, 2, 3, 1) # (B, H, W, C) + return resized_xyz + + def apply_pointmap_dropout(self, embeddings: torch.Tensor) -> torch.Tensor: + """ + Apply dropout to pointmap embeddings. + Drops entire pointmap for selected samples during training or when forced. + + When force_dropout_always is True, always drops pointmap regardless of training mode. + """ + # Check if we should apply dropout + should_apply_dropout = (self.training or self.force_dropout_always) and self.dropout_prob > 0 + + if not should_apply_dropout: + return embeddings + + # Check if dropout infrastructure exists + if not hasattr(self, 'dropped_xyz_token'): + if self.force_dropout_always: + raise RuntimeError( + "Cannot force dropout: model was initialized with dropout_prob=0. " + "Re-initialize with dropout_prob > 0 to enable forced dropout." + ) + return embeddings + + batch_size, n_windows, embed_dim = embeddings.shape + + # Decide dropout behavior + if self.force_dropout_always and not self.training: + # When forced during inference, always drop (100% dropout) + drop_mask = torch.ones(batch_size, device=embeddings.device, dtype=torch.bool) + else: + # Normal training dropout - use configured probability + drop_mask = torch.rand(batch_size, device=embeddings.device) < self.dropout_prob + + # Create dropped embedding for all windows - use same token for all patches + # Shape: (batch_size, n_windows, embed_dim) + dropped_embedding = self.dropped_xyz_token.view(1, 1, embed_dim).expand(batch_size, n_windows, embed_dim) + + # Add positional embeddings to dropped tokens (same as regular embeddings get) + n_windows_h = n_windows_w = int(n_windows ** 0.5) + pos_embed_patch = self._get_pos_embed((n_windows_h, n_windows_w)).reshape( + 1, n_windows, embed_dim + ) + dropped_embedding = dropped_embedding + pos_embed_patch + drop_mask_expanded = drop_mask.view(batch_size, 1, 1).expand_as(embeddings) + embeddings = torch.where(drop_mask_expanded, dropped_embedding, embeddings) + + return embeddings + + @torch._dynamo.disable() + def embed_pointmap_windows( + self, xyz: torch.Tensor, valid_mask: torch.Tensor = None + ) -> torch.Tensor: + """Process pointmap into window embeddings without positional encoding""" + with torch.no_grad(): + xyz = self.resize_input(xyz) + if valid_mask is None: + valid_mask = xyz.isfinite().all(dim=-1) + + B, H, W, _ = xyz.shape + assert ( + H % self.patch_size == 0 and W % self.patch_size == 0 + ), "image must be divisible by patch_size" + + # (1) Handle NaN values before remapping to prevent propagation + xyz_safe = xyz.clone() + xyz_safe[~valid_mask] = 0.0 # Set invalid points to 0 before remapping + + # (1b) remap points to normalize their range + xyz_remapped = self.point_remapper(xyz_safe) + + # (2) project + invalid token + x = self.point_proj(xyz_remapped) # (B,H,W,D) + + x[~valid_mask] = 0.0 # Stop gradient for invalid points + x[~valid_mask] += self.invalid_xyz_token + + return x, B, H, W + + def inner_forward( + self, x: torch.Tensor, B: int, H: int, W: int + ) -> torch.Tensor: + x = x.view( + B, + H // self.patch_size, + self.patch_size, + W // self.patch_size, + self.patch_size, + self.embed_dim, + ) # (B, hW, wW, ws, ws, D) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous() # (B, hW, wW, ws, ws, D) + x = x.view(-1, self.patch_size * self.patch_size, self.embed_dim) + + # (4) CLS token that contains the patch information + cls_tok = self.cls_token.expand(x.shape[0], -1, -1) + toks = torch.cat([cls_tok, x], dim=1) + + # (5) add positional embedding for window + toks = toks + self.pos_embed_window + + # (6) intra-window attention + for blk in self.blocks: + toks = blk(toks) + + # (7) Extract CLS tokens and reshape to (B, n_windows, embed_dim) + n_windows_h = H // self.patch_size + n_windows_w = W // self.patch_size + window_embeddings = toks[:, 0].view(B, n_windows_h * n_windows_w, self.embed_dim) + + # Add positional embeddings + pos_embed_patch = self._get_pos_embed((n_windows_h, n_windows_w)).reshape( + 1, n_windows_h * n_windows_w, self.embed_dim + ) + out = window_embeddings + pos_embed_patch + + # Apply dropout if enabled (during training OR when forced) + if (self.training or self.force_dropout_always) and self.dropout_prob > 0: + out = self.apply_pointmap_dropout(out) + + return out + + def forward( + self, xyz: torch.Tensor, valid_mask: torch.Tensor = None + ) -> torch.Tensor: + """ + xyz : (B, 3, H, W) map of (x,y,z) coordinates + valid_mask : (B, H, W) boolean - True for valid points (optional) + + returns: (B, num_windows, D) + """ + # Get window embeddings + x, B, H, W = self.embed_pointmap_windows(xyz, valid_mask) + return self.inner_forward(x, B, H, W) diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..043b3297cb11ac67d0a0a884ca5f3b8fbbf89c77 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/flow_matching/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/flow_matching/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/flow_matching/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/flow_matching/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da17be99cc8f294d8947019b4484f224900436f3 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/flow_matching/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/flow_matching/__pycache__/solver.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/flow_matching/__pycache__/solver.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b6f3f1277a4b068f0f57c5473aaf4d898e2c377 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/flow_matching/__pycache__/solver.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/flow_matching/solver.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/flow_matching/solver.py new file mode 100644 index 0000000000000000000000000000000000000000..ead1b58c921faebe0451227200bc93d16bda52aa --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/generator/flow_matching/solver.py @@ -0,0 +1,126 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import optree +import torch +from functools import partial + +from sam3d_objects.data.utils import tree_tensor_map + + +def linear_approximation_step(x_t, dt, velocity): + # x_tp1 = x_t + velocity * dt + x_tp1 = tree_tensor_map(lambda x, v: x + v * dt, x_t, velocity) + return x_tp1 + + +def gradient(output, x, create_graph: bool = False): + tensors, pyspec = optree.tree_flatten( + x, is_leaf=lambda x: isinstance(x, torch.Tensor) + ) + grad_outputs = [torch.ones_like(output).detach() for _ in tensors] + grads = torch.autograd.grad( + output, + tensors, + grad_outputs=grad_outputs, + create_graph=create_graph, + ) + return optree.tree_unflatten(pyspec, grads) + + +class ODESolver: + def step(self, dynamics_fn, x_t, t, dt, *args, **kwargs): + raise NotImplementedError + + def solve_iter(self, dynamics_fn, x_init, times, *args, **kwargs): + x_t = x_init + for t0, t1 in zip(times[:-1], times[1:]): + dt = t1 - t0 + x_t = self.step(dynamics_fn, x_t, t0, dt, *args, **kwargs) + yield x_t, t0 + + def solve(self, dynamics_fn, x_init, times, *args, **kwargs): + for x_t, _ in self.solve_iter(dynamics_fn, x_init, times, *args, **kwargs): + pass + return x_t + + +# https://en.wikipedia.org/wiki/Euler_method +class Euler(ODESolver): + def step(self, dynamics_fn, x_t, t, dt, *args, **kwargs): + velocity = dynamics_fn(x_t, t, *args, **kwargs) + x_tp1 = linear_approximation_step(x_t, dt, velocity) + return x_tp1 + + +# https://arxiv.org/abs/2505.05470 +class SDE(ODESolver): + def __init__(self, **kwargs): + super().__init__() + self.sde_strength = kwargs.get("sde_strength", 0.1) + + def step(self, dynamics_fn, x_t, t, dt, *args, **kwargs): + velocity = dynamics_fn(x_t, t, *args, **kwargs) + sigma = 1 - t + var_t = sigma / (1 - torch.tensor(sigma).clamp(min=dt)) + std_dev_t = ( + torch.sqrt(variance) * self.sde_strength + ) # self.sde_strength = alpha + + def compute_mean(x, v): + drift_term = x * (std_dev_t**2 / (2 * sigma) * dt) + velocity_term = v * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt + return x + drift_term + velocity_term + + prev_sample_mean = tree_tensor_map(compute_mean, x_t, velocity) + + # Generate noise and compute final sample using tree_tensor_map + def add_noise(mean_val): + variance_noise = torch.randn_like(mean_val) + return mean_val + std_dev_t * torch.sqrt(torch.tensor(dt)) * variance_noise + + prev_sample = tree_tensor_map(add_noise, prev_sample_mean) + + return prev_sample + + +# https://en.wikipedia.org/wiki/Midpoint_method +class Midpoint(ODESolver): + def step(self, dynamics_fn, x_t, t, dt, *args, **kwargs): + half_dt = 0.5 * dt + + x_mid = Euler.step(self, dynamics_fn, x_t, t, half_dt, *args, **kwargs) + + velocity_mid = dynamics_fn(x_mid, t + half_dt, *args, **kwargs) + x_tp1 = linear_approximation_step(x_t, dt, velocity_mid) + return x_tp1 + + +# https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods +class RungeKutta4(ODESolver): + + def k1(self, dynamics_fn, x_t, t, dt, *args, **kwargs): + return dynamics_fn(x_t, t, *args, **kwargs) + + def k2(self, dynamics_fn, x_t, t, dt, k1, *args, **kwargs): + x_k1 = linear_approximation_step(x_t, dt * 0.5, k1) + return dynamics_fn(x_k1, t + dt * 0.5, *args, **kwargs) + + def k3(self, dynamics_fn, x_t, t, dt, k2, *args, **kwargs): + x_k2 = linear_approximation_step(x_t, dt * 0.5, k2) + return dynamics_fn(x_k2, t + dt * 0.5, *args, **kwargs) + + def k4(self, dynamics_fn, x_t, t, dt, k3, *args, **kwargs): + x_k3 = linear_approximation_step(x_t, dt, k3) + return dynamics_fn(x_k3, t + dt, *args, **kwargs) + + def step(self, dynamics_fn, x_t, t, dt, *args, **kwargs): + k1 = self.k1(dynamics_fn, x_t, t, dt, *args, **kwargs) + k2 = self.k2(dynamics_fn, x_t, t, dt, k1, *args, **kwargs) + k3 = self.k3(dynamics_fn, x_t, t, dt, k2, *args, **kwargs) + k4 = self.k4(dynamics_fn, x_t, t, dt, k3, *args, **kwargs) + + def compute_velocity(k1, k2, k3, k4): + return (k1 + 2 * k2 + 2 * k3 + k4) / 6 + + velocity_k = tree_tensor_map(compute_velocity, k1, k2, k3, k4) + x_tp1 = linear_approximation_step(x_t, dt, velocity_k) + return x_tp1 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51fe9695f70d4c8d6e67ef1b56cec18bb2b67a4a Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8fdaca46214e7231491cdbe666e72265d4671f49 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/mm_latent.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/mm_latent.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f49641917274cb964474875ff3ed7d503f56a73d Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/mm_latent.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/mot_sparse_structure_flow.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/mot_sparse_structure_flow.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1223f0efc77fdf4fc3cff9228d3bff238a851935 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/mot_sparse_structure_flow.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/sparse_structure_vae.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/sparse_structure_vae.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e07e4ca1a5e642d7f0f0a2a28814314f701136c Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/sparse_structure_vae.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/timestep_embedder.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/timestep_embedder.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..186ed7dac47bfec2c0bfb2eec69119862538205a Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/__pycache__/timestep_embedder.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/mm_latent.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/mm_latent.py new file mode 100644 index 0000000000000000000000000000000000000000..4abacfea9298097430c3a51846f383dd998324c8 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/mm_latent.py @@ -0,0 +1,90 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import torch +from torch import nn +from ..modules.transformer import ( + AbsolutePositionEmbedder, +) +from ..modules import spatial +import torch.nn.functional as F +from typing import Callable + + +class Latent(nn.Module): + def __init__( + self, in_channels, model_channels: int, pos_embedder: Callable[[], torch.Tensor] + ): + super().__init__() + self.input_layer = nn.Linear(in_channels, model_channels) + self.out_layer = nn.Linear(model_channels, in_channels) + + pos_emb = pos_embedder() + if isinstance(pos_emb, torch.nn.Parameter): + # learnt position embedding + self.register_parameter("pos_emb", pos_emb) + elif isinstance(pos_emb, torch.Tensor): + # fixed position embedding + self.register_buffer("pos_emb", pos_emb) + else: + raise NotImplementedError + + self.initialize_weights() + + def initialize_weights(self) -> None: + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + self.apply(_basic_init) + + # Zero-out output layers: + # nn.init.constant_(self.out_layer.weight, 0) + nn.init.xavier_uniform_(self.out_layer.weight) + nn.init.constant_(self.out_layer.bias, 0) + + def to_input(self, x): + x = self.input_layer(x) + x = x + self.pos_emb[None] + + return x + + def to_output(self, h): + h = F.layer_norm(h, h.shape[-1:]) + h = self.out_layer(h) + + return h + + +def ShapePositionEmbedder(model_channels, resolution, patch_size): + def embedder(): + pos_embedder = AbsolutePositionEmbedder(model_channels, 3) + coords = torch.meshgrid( + *[torch.arange(res) for res in [resolution // patch_size] * 3], + indexing="ij", + ) + coords = torch.stack(coords, dim=-1).reshape(-1, 3) + pos_emb = pos_embedder(coords) + + return pos_emb + + return embedder + + +def RandomPositionEmbedder(model_channels, token_len): + def embedder(): + pos_emb = torch.randn(token_len, model_channels) + + return pos_emb + + return embedder + + +def LearntPositionEmbedder(model_channels, token_len): + def embedder(): + pos_emb = torch.nn.Parameter(torch.randn(token_len, model_channels)) + + return pos_emb + + return embedder diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/mot_sparse_structure_flow.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/mot_sparse_structure_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..8cd2b4f1f980f1bdea6d92c0e08e6ed1b9e88636 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/mot_sparse_structure_flow.py @@ -0,0 +1,303 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from functools import partial +from typing import * +from torch.utils import _pytree +import torch +import torch.nn as nn +from ..modules.utils import convert_module_to_f16, convert_module_to_f32 +from collections import namedtuple +from ..modules.utils import FP16_TYPE +from ..modules.transformer import ( + MOTModulatedTransformerCrossBlock, +) +from sam3d_objects.data.utils import ( + tree_reduce_unique, +) +from .timestep_embedder import TimestepEmbedder +from omegaconf import OmegaConf + +class SparseStructureFlowModel(nn.Module): + def __init__( + self, + in_channels: int, + model_channels: int, + cond_channels: int, + out_channels: int, + num_blocks: int, + num_heads: Optional[int] = None, + num_head_channels: Optional[int] = 64, + mlp_ratio: float = 4, + pe_mode: Literal["ape", "rope"] = "ape", + use_fp16: bool = False, + use_checkpoint: bool = False, + share_mod: bool = False, + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + freeze_shared_parameters: bool = False, + is_shortcut_model: bool = False, + *args, + **kwargs, + ): + super().__init__() + self.in_channels = in_channels + self.model_channels = model_channels + self.cond_channels = cond_channels + self.out_channels = out_channels + self.num_blocks = num_blocks + self.num_heads = num_heads or model_channels // num_head_channels + self.mlp_ratio = mlp_ratio + self.pe_mode = pe_mode + self.use_fp16 = use_fp16 + self.use_checkpoint = use_checkpoint + self.share_mod = share_mod + self.qk_rms_norm = qk_rms_norm + self.qk_rms_norm_cross = qk_rms_norm_cross + self.dtype = FP16_TYPE if use_fp16 else torch.float32 + self.is_shortcut_model = is_shortcut_model + if is_shortcut_model: + self.d_embedder = TimestepEmbedder(model_channels) # for shortcut model + + self.t_embedder = TimestepEmbedder(model_channels, freeze=freeze_shared_parameters) + if share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(model_channels, 6 * model_channels, bias=True) + ) + + self.blocks = nn.ModuleList( + [ + MOTModulatedTransformerCrossBlock( + model_channels, + cond_channels, + num_heads=self.num_heads, + mlp_ratio=self.mlp_ratio, + attn_mode="full", + use_checkpoint=self.use_checkpoint, + use_rope=(pe_mode == "rope"), + share_mod=share_mod, + qk_rms_norm=self.qk_rms_norm, + qk_rms_norm_cross=self.qk_rms_norm_cross, + latent_names=self.latent_names, + freeze_shared_parameters=freeze_shared_parameters + ) + for _ in range(num_blocks) + ] + ) + + self.initialize_weights() + if use_fp16: + self.convert_to_fp16() + + @property + def device(self) -> torch.device: + """ + Return the device of the model. + """ + return next(self.parameters()).device + + def convert_to_fp16(self) -> None: + """ + Convert the torso of the model to float16. + """ + self.blocks.apply(convert_module_to_f16) + + def convert_to_fp32(self) -> None: + """ + Convert the torso of the model to float32. + """ + self.blocks.apply(convert_module_to_f32) + + + + def initialize_weights(self) -> None: + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + self.apply(_basic_init) + + # Initialize timestep embedding MLP: + nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) + + # zero init like controlnet, for MLP should only zero + # the weight of the last layer only + if self.is_shortcut_model: + nn.init.constant_(self.d_embedder.mlp[2].weight, 0) + + # Zero-out adaLN modulation layers in DiT blocks: + if self.share_mod: + nn.init.constant_(self.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.adaLN_modulation[-1].bias, 0) + else: + for block in self.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + def _cast_type(self, x, dtype): + return x.type(dtype) + + def forward( + self, + h: Dict, + t: torch.Tensor, + cond: torch.Tensor, + d: torch.Tensor = None, + mode="default", + ) -> torch.Tensor: + t_emb = self.t_embedder(t) + if d is not None: + d_emb = self.d_embedder(d) + t_emb = t_emb + d_emb + if self.share_mod: + t_emb = self.adaLN_modulation(t_emb) + + input_dtype = tree_reduce_unique(lambda tensor: tensor.dtype, h) + t_emb = t_emb.type(self.dtype) + h = _pytree.tree_map( + partial(self._cast_type, dtype=self.dtype), + h, + ) + cond = cond.type(self.dtype) + + for block in self.blocks: + h = block(h, t_emb, cond, mode=mode) + + h = _pytree.tree_map( + partial(self._cast_type, dtype=input_dtype), + h, + ) + + return h + + +class SparseStructureFlowTdfyWrapper(SparseStructureFlowModel): + def __init__( + self, + latent_mapping: dict, + latent_share_transformer: dict = {}, + *args, + **kwargs, + ): + condition_embedder = kwargs.pop("condition_embedder", None) + # if enabled, model will record the condition_shape in one run and uses zeros for all that afterwards + force_zeros_cond = kwargs.pop("force_zeros_cond", False) + # backward compatible to models trained before PR #87 + kwargs.pop("shape_attend_pose", None) + merge_latent_names = [i for _, v in latent_share_transformer.items() for i in v] + self.latent_names = [ + latent_name + for latent_name in list(latent_mapping.keys()) + if latent_name not in merge_latent_names + ] + list(latent_share_transformer.keys()) + super().__init__(*args, **kwargs) + if condition_embedder is not None: + self.condition_embedder = condition_embedder + else: + self.condition_embedder = lambda x: x + self.force_zeros_cond = force_zeros_cond + self.latent_mapping = nn.ModuleDict(latent_mapping) + if not isinstance(latent_share_transformer, dict): + self.latent_share_transformer = OmegaConf.to_container(latent_share_transformer) + else: + self.latent_share_transformer = latent_share_transformer + self.input_latent_mappings = list(self.latent_mapping.keys()) + + def forward( + self, + latents_dict: dict, + t: torch.Tensor, + *condition_args, + mode="default", + **condition_kwargs, + ) -> dict: + d = condition_kwargs.pop("d", None) + + cfg_activate = condition_kwargs.pop("cfg", False) + if self.force_zeros_cond and cfg_activate: + cond = self.condition_embedder(*condition_args, **condition_kwargs) + cond = cond * 0 + else: + cond = self.condition_embedder(*condition_args, **condition_kwargs) + + # concatenate input + latent_dict = self.project_input(latents_dict) + output = super().forward(latent_dict, t, cond, d, mode=mode) + + # Save pre-project 6D embedding [1, 4, 1024] as attribute (not in output dict, + # which would break the solver's key validation). Retrieved via + # hasattr(inner_model, "_last_6d_embedding") in inference_pipeline.py. + pre_project_6d = output.get("6drotation_normalized", None) + if pre_project_6d is not None: + self._last_6d_embedding = pre_project_6d + + output_latents = self.project_output(output) + + return output_latents + + def project_input( + self, + latents_dict: Dict, + ) -> Dict: + # concatenate input from multiple modalities + latent_dict = {} + for latent_name in self.input_latent_mappings: + assert ( + latent_name in latents_dict + ), f"'{latent_name}' not found in latents_dict" + latent_input = latents_dict[latent_name] + x = self.latent_mapping[latent_name].to_input(latent_input) + latent_dict[latent_name] = x + + latent_dict = self.merge_latent_share_transformer(latent_dict) + return latent_dict + + def project_output(self, output: Dict) -> Dict: + ### slicing latents --> pass through final MLP + output = self.split_latent_share_transformer(output) + output_latents = {} + for latent_name, latent_tensor in output.items(): + if latent_name in self.latent_mapping: + output_latents[latent_name] = self.latent_mapping[latent_name].to_output(latent_tensor) + else: + output_latents[latent_name] = latent_tensor + + return output_latents + + def merge_latent_share_transformer(self, latent_dict): + visited_latent_names = set() + return_dict = {} + for merged_name, latent_names in self.latent_share_transformer.items(): + tensors = [] + for latent_name in latent_names: + visited_latent_names.add(latent_name) + tensors.append(latent_dict[latent_name]) + tensors = torch.cat(tensors, dim=1) + return_dict[merged_name] = tensors + + for latent_name in latent_dict: + if latent_name not in visited_latent_names: + return_dict[latent_name] = latent_dict[latent_name] + + return return_dict + + def split_latent_share_transformer(self, output_latents): + return_dict = {} + visited_latent_names = set() + for merged_name, latent_names in self.latent_share_transformer.items(): + start = 0 + visited_latent_names.add(merged_name) + tensors = output_latents[merged_name] + for latent_name in latent_names: + token_len = self.latent_mapping[latent_name].pos_emb.shape[0] + latent = tensors[:, start : start + token_len] + return_dict[latent_name] = latent + start += token_len + + for latent_name in output_latents: + if latent_name not in visited_latent_names: + return_dict[latent_name] = output_latents[latent_name] + + return return_dict diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/sparse_structure_vae.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/sparse_structure_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..0db7f5790f6698199cc9873ab032e23a0d85a98c --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/sparse_structure_vae.py @@ -0,0 +1,402 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import os +import math +from typing import * +import torch +import torch.nn as nn +import torch.nn.functional as F +from ..modules.norm import GroupNorm32, ChannelLayerNorm32 +from ..modules.spatial import pixel_shuffle_3d +from ..modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32 +from safetensors.torch import load_file +from loguru import logger + + +def norm_layer(norm_type: str, *args, **kwargs) -> nn.Module: + """ + Return a normalization layer. + """ + if norm_type == "group": + return GroupNorm32(32, *args, **kwargs) + elif norm_type == "layer": + return ChannelLayerNorm32(*args, **kwargs) + else: + raise ValueError(f"Invalid norm type {norm_type}") + + +class ResBlock3d(nn.Module): + def __init__( + self, + channels: int, + out_channels: Optional[int] = None, + norm_type: Literal["group", "layer"] = "layer", + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + + self.norm1 = norm_layer(norm_type, channels) + self.norm2 = norm_layer(norm_type, self.out_channels) + self.conv1 = nn.Conv3d(channels, self.out_channels, 3, padding=1) + self.conv2 = zero_module( + nn.Conv3d(self.out_channels, self.out_channels, 3, padding=1) + ) + self.skip_connection = ( + nn.Conv3d(channels, self.out_channels, 1) + if channels != self.out_channels + else nn.Identity() + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.norm1(x) + h = F.silu(h) + h = self.conv1(h) + h = self.norm2(h) + h = F.silu(h) + h = self.conv2(h) + h = h + self.skip_connection(x) + return h + + +class DownsampleBlock3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + mode: Literal["conv", "avgpool"] = "conv", + ): + assert mode in ["conv", "avgpool"], f"Invalid mode {mode}" + + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + if mode == "conv": + self.conv = nn.Conv3d(in_channels, out_channels, 2, stride=2) + elif mode == "avgpool": + assert ( + in_channels == out_channels + ), "Pooling mode requires in_channels to be equal to out_channels" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if hasattr(self, "conv"): + return self.conv(x) + else: + return F.avg_pool3d(x, 2) + + +class UpsampleBlock3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + mode: Literal["conv", "nearest"] = "conv", + ): + assert mode in ["conv", "nearest"], f"Invalid mode {mode}" + + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + if mode == "conv": + self.conv = nn.Conv3d(in_channels, out_channels * 8, 3, padding=1) + elif mode == "nearest": + assert ( + in_channels == out_channels + ), "Nearest mode requires in_channels to be equal to out_channels" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if hasattr(self, "conv"): + x = self.conv(x) + return pixel_shuffle_3d(x, 2) + else: + return F.interpolate(x, scale_factor=2, mode="nearest") + + +class SparseStructureEncoder(nn.Module): + """ + Encoder for Sparse Structure (\mathcal{E}_S in the paper Sec. 3.3). + + Args: + in_channels (int): Channels of the input. + latent_channels (int): Channels of the latent representation. + num_res_blocks (int): Number of residual blocks at each resolution. + channels (List[int]): Channels of the encoder blocks. + num_res_blocks_middle (int): Number of residual blocks in the middle. + norm_type (Literal["group", "layer"]): Type of normalization layer. + use_fp16 (bool): Whether to use FP16. + """ + + def __init__( + self, + in_channels: int, + latent_channels: int, + num_res_blocks: int, + channels: List[int], + num_res_blocks_middle: int = 2, + norm_type: Literal["group", "layer"] = "layer", + use_fp16: bool = False, + ): + super().__init__() + self.in_channels = in_channels + self.latent_channels = latent_channels + self.num_res_blocks = num_res_blocks + self.channels = channels + self.num_res_blocks_middle = num_res_blocks_middle + self.norm_type = norm_type + self.use_fp16 = use_fp16 + self.dtype = torch.float16 if use_fp16 else torch.float32 + + self.input_layer = nn.Conv3d(in_channels, channels[0], 3, padding=1) + + self.blocks = nn.ModuleList([]) + for i, ch in enumerate(channels): + self.blocks.extend([ResBlock3d(ch, ch) for _ in range(num_res_blocks)]) + if i < len(channels) - 1: + self.blocks.append(DownsampleBlock3d(ch, channels[i + 1])) + + self.middle_block = nn.Sequential( + *[ + ResBlock3d(channels[-1], channels[-1]) + for _ in range(num_res_blocks_middle) + ] + ) + + self.out_layer = nn.Sequential( + norm_layer(norm_type, channels[-1]), + nn.SiLU(), + nn.Conv3d(channels[-1], latent_channels * 2, 3, padding=1), + ) + + if use_fp16: + self.convert_to_fp16() + + @property + def device(self) -> torch.device: + """ + Return the device of the model. + """ + return next(self.parameters()).device + + def convert_to_fp16(self) -> None: + """ + Convert the torso of the model to float16. + """ + self.use_fp16 = True + self.dtype = torch.float16 + self.blocks.apply(convert_module_to_f16) + self.middle_block.apply(convert_module_to_f16) + + def convert_to_fp32(self) -> None: + """ + Convert the torso of the model to float32. + """ + self.use_fp16 = False + self.dtype = torch.float32 + self.blocks.apply(convert_module_to_f32) + self.middle_block.apply(convert_module_to_f32) + + def forward( + self, x: torch.Tensor, sample_posterior: bool = False, return_raw: bool = False + ) -> torch.Tensor: + x = x.float() + h = self.input_layer(x) + h = h.type(self.dtype) + + for block in self.blocks: + h = block(h) + h = self.middle_block(h) + + h = h.type(x.dtype) + h = self.out_layer(h) + + mean, logvar = h.chunk(2, dim=1) + + if sample_posterior: + std = torch.exp(0.5 * logvar) + z = mean + std * torch.randn_like(std) + else: + z = mean + + if return_raw: + return z, mean, logvar + return z + + +class SparseStructureDecoder(nn.Module): + """ + Decoder for Sparse Structure (\mathcal{D}_S in the paper Sec. 3.3). + + Args: + out_channels (int): Channels of the output. + latent_channels (int): Channels of the latent representation. + num_res_blocks (int): Number of residual blocks at each resolution. + channels (List[int]): Channels of the decoder blocks. + num_res_blocks_middle (int): Number of residual blocks in the middle. + norm_type (Literal["group", "layer"]): Type of normalization layer. + use_fp16 (bool): Whether to use FP16. + """ + + def __init__( + self, + out_channels: int, + latent_channels: int, + num_res_blocks: int, + channels: List[int], + num_res_blocks_middle: int = 2, + norm_type: Literal["group", "layer"] = "layer", + reshape_input_to_cube: bool = False, + use_fp16: bool = False, + ): + super().__init__() + self.out_channels = out_channels + self.latent_channels = latent_channels + self.num_res_blocks = num_res_blocks + self.channels = channels + self.num_res_blocks_middle = num_res_blocks_middle + self.norm_type = norm_type + self.use_fp16 = use_fp16 + self.dtype = torch.float16 if use_fp16 else torch.float32 + self.reshape_input_to_cube = reshape_input_to_cube + self.input_layer = nn.Conv3d(latent_channels, channels[0], 3, padding=1) + + self.middle_block = nn.Sequential( + *[ + ResBlock3d(channels[0], channels[0]) + for _ in range(num_res_blocks_middle) + ] + ) + + self.blocks = nn.ModuleList([]) + for i, ch in enumerate(channels): + self.blocks.extend([ResBlock3d(ch, ch) for _ in range(num_res_blocks)]) + if i < len(channels) - 1: + self.blocks.append(UpsampleBlock3d(ch, channels[i + 1])) + + self.out_layer = nn.Sequential( + norm_layer(norm_type, channels[-1]), + nn.SiLU(), + nn.Conv3d(channels[-1], out_channels, 3, padding=1), + ) + + if use_fp16: + self.convert_to_fp16() + + @property + def device(self) -> torch.device: + """ + Return the device of the model. + """ + return next(self.parameters()).device + + def convert_to_fp16(self) -> None: + """ + Convert the torso of the model to float16. + """ + self.use_fp16 = True + self.dtype = torch.float16 + self.blocks.apply(convert_module_to_f16) + self.middle_block.apply(convert_module_to_f16) + + def convert_to_fp32(self) -> None: + """ + Convert the torso of the model to float32. + """ + self.use_fp16 = False + self.dtype = torch.float32 + self.blocks.apply(convert_module_to_f32) + self.middle_block.apply(convert_module_to_f32) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.reshape_input_to_cube: + x = self.flat_to_cube(x) + + h = self.input_layer(x) + + h = h.type(self.dtype) + + h = self.middle_block(h) + for block in self.blocks: + h = block(h) + + h = h.type(x.dtype) + h = self.out_layer(h) + return h + + @staticmethod + def flat_to_cube(flat_latent: torch.Tensor) -> torch.Tensor: + """ + For converting latent tokens from generator to cube + + Args: + flat_latent: (B, T, C) + Returns: + cube: (B, C, D, H, W) + """ + k = round(math.pow(flat_latent.shape[1], 1 / 3)) + assert ( + k**3 == flat_latent.shape[1] + ), f"Flat latent must be a cube {k**3} != {flat_latent.shape[1]}" + latent = flat_latent.view( + flat_latent.shape[0], k, k, k, flat_latent.shape[2] + ).permute(0, 4, 1, 2, 3) + return latent + + +class SparseStructureDecoderTdfyWrapper(SparseStructureDecoder): + def __init__(self, *args, **kwargs): + pretrained_ckpt_path = kwargs.pop("pretrained_ckpt_path", None) + super().__init__(*args, **kwargs) + if pretrained_ckpt_path is not None: + if os.path.exists(pretrained_ckpt_path): + logger.info( + f"Loading pretrained ss decoder from {pretrained_ckpt_path}" + ) + file_type = os.path.splitext(pretrained_ckpt_path)[1] + if file_type == ".safetensors": + self.load_state_dict(load_file(pretrained_ckpt_path)) + else: + self.load_state_dict( + torch.load(pretrained_ckpt_path, weights_only=True) + ) + else: + raise FileNotFoundError( + f"The path for the SS decoder does not exist: {pretrained_ckpt_path}" + ) + + +class SparseStructureEncoderTdfyWrapper(SparseStructureEncoder): + def __init__(self, sample_posterior=True, return_raw=True, *args, **kwargs): + pretrained_ckpt_path = kwargs.pop("pretrained_ckpt_path", None) + super().__init__(*args, **kwargs) + if pretrained_ckpt_path is not None: + if os.path.exists(pretrained_ckpt_path): + logger.info( + f"Loading pretrained ss encoder from {pretrained_ckpt_path}" + ) + file_type = os.path.splitext(pretrained_ckpt_path)[1] + if file_type == ".safetensors": + self.load_state_dict(load_file(pretrained_ckpt_path)) + else: + self.load_state_dict( + torch.load(pretrained_ckpt_path, weights_only=True) + ) + else: + raise FileNotFoundError( + f"The path for the SS encoder does not exist: {pretrained_ckpt_path}" + ) + self.sample_posterior = sample_posterior + self.return_raw = return_raw + + def forward(self, x: torch.Tensor) -> torch.Tensor: + z, mean, logvar = super().forward( + x, sample_posterior=self.sample_posterior, return_raw=True + ) + if self.return_raw: + return { + "z": z, + "mean": mean, + "logvar": logvar, + } + else: + return {"z": z} diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/timestep_embedder.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/timestep_embedder.py new file mode 100644 index 0000000000000000000000000000000000000000..fa8b5293028c76d7c8470e933df9ba98410353f5 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/models/timestep_embedder.py @@ -0,0 +1,59 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import torch.nn as nn +import torch +import numpy as np + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__(self, hidden_size, frequency_embedding_size=256, freeze=False): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + if freeze: + self.requires_grad_(False) + self.eval() + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + dim: the dimension of the output. + max_period: controls the minimum frequency of the embeddings. + + Returns: + an (N, D) Tensor of positional embeddings. + """ + # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + half = dim // 2 + freqs = torch.exp( + -np.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ).to(device=t.device) + t = t[:, None].float() + args = t * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 + ) + return embedding + + def forward(self, t): + if t.ndim == 0: + t = t.unsqueeze(0) + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_emb = self.mlp(t_freq) + return t_emb diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b038b9c96ba9df517c83e37dd7e90181178c0360 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/norm.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/norm.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4ee86843dc8827f9c816904d373f0b622edf7ce8 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/norm.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/spatial.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/spatial.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f593b2ac4494b5b7868f602bd5b7d9f0efff421b Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/spatial.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/utils.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c8b5f0045a928ab4ed0c408640808ac349b7fcbd Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/__pycache__/utils.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..14053c17d579ddee4b6224c0b8a1be87cd7d23f4 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__init__.py @@ -0,0 +1,47 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from typing import * +from loguru import logger + +BACKEND = "sdpa" +DEBUG = False + + +def __from_env(): + import os + + global BACKEND + global DEBUG + + env_attn_backend = os.environ.get("ATTN_BACKEND") + env_sttn_debug = os.environ.get("ATTN_DEBUG") + + if env_attn_backend is not None and env_attn_backend in [ + "xformers", + "flash_attn", + "torch_flash_attn", + "sdpa", + "naive", + ]: + BACKEND = env_attn_backend + # BACKEND = "sdpa" + if env_sttn_debug is not None: + DEBUG = env_sttn_debug == "1" + + logger.info(f"[ATTENTION] Using backend: {BACKEND}") + + +__from_env() + + +def set_backend(backend: Literal["xformers", "flash_attn", "torch_flash_attn"]): + global BACKEND + BACKEND = backend + + +def set_debug(debug: bool): + global DEBUG + DEBUG = debug + + +from .full_attn import * +from .modules import * diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..038c8b51149cfedaebf574d7157d3318404f4ecd Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__pycache__/full_attn.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__pycache__/full_attn.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..113d25660c3c110f3fd532c84a9fcd2ed4d74a39 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__pycache__/full_attn.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__pycache__/modules.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__pycache__/modules.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a59f6712ce85577f6070c35507f3151da856a3fd Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/__pycache__/modules.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/full_attn.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/full_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..53e1ec8a8c399fb2d86cf3231adc6d3b0cadce29 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/full_attn.py @@ -0,0 +1,181 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from typing import * +import torch +import math +from . import DEBUG, BACKEND +from torch.nn.functional import scaled_dot_product_attention as torch_sdpa + +if BACKEND == "xformers": + import xformers.ops as xops +elif BACKEND == "flash_attn": + import flash_attn +elif BACKEND == "torch_flash_attn": + pass +elif BACKEND == "sdpa": + sdpa = torch_sdpa +elif BACKEND == "naive": + pass +else: + raise ValueError(f"Unknown attention backend: {BACKEND}") + + +__all__ = [ + "scaled_dot_product_attention", +] + + +def _naive_sdpa(q, k, v): + """ + Naive implementation of scaled dot product attention. + """ + q = q.permute(0, 2, 1, 3) # [N, H, L, C] + k = k.permute(0, 2, 1, 3) # [N, H, L, C] + v = v.permute(0, 2, 1, 3) # [N, H, L, C] + scale_factor = 1 / math.sqrt(q.size(-1)) + attn_weight = q @ k.transpose(-2, -1) * scale_factor + attn_weight = torch.softmax(attn_weight, dim=-1) + out = attn_weight @ v + out = out.permute(0, 2, 1, 3) # [N, L, H, C] + return out + + +@overload +def scaled_dot_product_attention(qkv: torch.Tensor) -> torch.Tensor: + """ + Apply scaled dot product attention. + + Args: + qkv (torch.Tensor): A [N, L, 3, H, C] tensor containing Qs, Ks, and Vs. + """ + ... + + +@overload +def scaled_dot_product_attention(q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor: + """ + Apply scaled dot product attention. + + Args: + q (torch.Tensor): A [N, L, H, C] tensor containing Qs. + kv (torch.Tensor): A [N, L, 2, H, C] tensor containing Ks and Vs. + """ + ... + + +@overload +def scaled_dot_product_attention( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor +) -> torch.Tensor: + """ + Apply scaled dot product attention. + + Args: + q (torch.Tensor): A [N, L, H, Ci] tensor containing Qs. + k (torch.Tensor): A [N, L, H, Ci] tensor containing Ks. + v (torch.Tensor): A [N, L, H, Co] tensor containing Vs. + + Note: + k and v are assumed to have the same coordinate map. + """ + ... + + +def scaled_dot_product_attention(*args, **kwargs): + arg_names_dict = {1: ["qkv"], 2: ["q", "kv"], 3: ["q", "k", "v"]} + num_all_args = len(args) + len(kwargs) + assert ( + num_all_args in arg_names_dict + ), f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3" + for key in arg_names_dict[num_all_args][len(args) :]: + assert key in kwargs, f"Missing argument {key}" + + if num_all_args == 1: + qkv = args[0] if len(args) > 0 else kwargs["qkv"] + assert ( + len(qkv.shape) == 5 and qkv.shape[2] == 3 + ), f"Invalid shape for qkv, got {qkv.shape}, expected [N, L, 3, H, C]" + device = qkv.device + + elif num_all_args == 2: + q = args[0] if len(args) > 0 else kwargs["q"] + kv = args[1] if len(args) > 1 else kwargs["kv"] + assert ( + q.shape[0] == kv.shape[0] + ), f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}" + assert ( + len(q.shape) == 4 + ), f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]" + assert ( + len(kv.shape) == 5 + ), f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]" + device = q.device + + elif num_all_args == 3: + q = args[0] if len(args) > 0 else kwargs["q"] + k = args[1] if len(args) > 1 else kwargs["k"] + v = args[2] if len(args) > 2 else kwargs["v"] + assert ( + q.shape[0] == k.shape[0] == v.shape[0] + ), f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}" + assert ( + len(q.shape) == 4 + ), f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]" + assert ( + len(k.shape) == 4 + ), f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]" + assert ( + len(v.shape) == 4 + ), f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]" + device = q.device + + effective_backend = BACKEND + if device.type != "cuda" and BACKEND in {"flash_attn", "torch_flash_attn", "xformers"}: + effective_backend = "sdpa" + + if effective_backend == "xformers": + if num_all_args == 1: + q, k, v = qkv.unbind(dim=2) + elif num_all_args == 2: + k, v = kv.unbind(dim=2) + out = xops.memory_efficient_attention(q, k, v) + elif effective_backend == "flash_attn": + if num_all_args == 1: + out = flash_attn.flash_attn_qkvpacked_func(qkv) + elif num_all_args == 2: + out = flash_attn.flash_attn_kvpacked_func(q, kv) + elif num_all_args == 3: + out = flash_attn.flash_attn_func(q, k, v) + elif effective_backend == "sdpa": + if num_all_args == 1: + q, k, v = qkv.unbind(dim=2) + elif num_all_args == 2: + k, v = kv.unbind(dim=2) + q = q.permute(0, 2, 1, 3) # [N, H, L, C] + k = k.permute(0, 2, 1, 3) # [N, H, L, C] + v = v.permute(0, 2, 1, 3) # [N, H, L, C] + out = torch_sdpa(q, k, v) # [N, H, L, C] + out = out.permute(0, 2, 1, 3) # [N, L, H, C] + elif effective_backend == "torch_flash_attn": + if num_all_args == 1: + q, k, v = qkv.unbind(dim=2) + elif num_all_args == 2: + k, v = kv.unbind(dim=2) + q = q.permute(0, 2, 1, 3) # [N, H, L, C] + k = k.permute(0, 2, 1, 3) # [N, H, L, C] + v = v.permute(0, 2, 1, 3) # [N, H, L, C] + original_dtype = q.dtype + with torch.nn.attention.sdpa_kernel(backends=[torch.nn.attention.SDPBackend.FLASH_ATTENTION]): + with torch.autocast(device_type=device.type, dtype=torch.bfloat16): + out = torch.nn.functional.scaled_dot_product_attention(q, k, v) # [N, H, L, C] + out = out.permute(0, 2, 1, 3) # [N, L, H, C] + out = out.to(original_dtype) + elif effective_backend == "naive": + if num_all_args == 1: + q, k, v = qkv.unbind(dim=2) + elif num_all_args == 2: + k, v = kv.unbind(dim=2) + out = _naive_sdpa(q, k, v) + else: + raise ValueError(f"Unknown attention module: {effective_backend}") + + return out diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/modules.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..e17577329da4ab8b8abccc1ff4e78e0a85d2f906 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/attention/modules.py @@ -0,0 +1,451 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from functools import partial +from typing import * +from torch.utils import _pytree +import torch +import torch.nn as nn +import torch.nn.functional as F +from .full_attn import scaled_dot_product_attention +from sam3d_objects.data.utils import ( + tree_reduce_unique, +) + + +class MultiHeadRMSNorm(nn.Module): + def __init__(self, dim: int, heads: int): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(heads, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (F.normalize(x, dim=-1) * self.gamma * self.scale).to(x.dtype) + + +class RotaryPositionEmbedder(nn.Module): + def __init__(self, hidden_size: int, in_channels: int = 3): + super().__init__() + assert hidden_size % 2 == 0, "Hidden size must be divisible by 2" + self.hidden_size = hidden_size + self.in_channels = in_channels + self.freq_dim = hidden_size // in_channels // 2 + self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim + self.freqs = 1.0 / (10000**self.freqs) + + def _get_phases(self, indices: torch.Tensor) -> torch.Tensor: + self.freqs = self.freqs.to(indices.device) + phases = torch.outer(indices, self.freqs) + phases = torch.polar(torch.ones_like(phases), phases) + return phases + + def _rotary_embedding(self, x: torch.Tensor, phases: torch.Tensor) -> torch.Tensor: + x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + x_rotated = x_complex * phases + x_embed = ( + torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1).to(x.dtype) + ) + return x_embed + + def forward( + self, q: torch.Tensor, k: torch.Tensor, indices: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + q (sp.SparseTensor): [..., N, D] tensor of queries + k (sp.SparseTensor): [..., N, D] tensor of keys + indices (torch.Tensor): [..., N, C] tensor of spatial positions + """ + if indices is None: + indices = torch.arange(q.shape[-2], device=q.device) + if len(q.shape) > 2: + indices = indices.unsqueeze(0).expand(q.shape[:-2] + (-1,)) + + phases = self._get_phases(indices.reshape(-1)).reshape(*indices.shape[:-1], -1) + if phases.shape[1] < self.hidden_size // 2: + phases = torch.cat( + [ + phases, + torch.polar( + torch.ones( + *phases.shape[:-1], + self.hidden_size // 2 - phases.shape[1], + device=phases.device, + ), + torch.zeros( + *phases.shape[:-1], + self.hidden_size // 2 - phases.shape[1], + device=phases.device, + ), + ), + ], + dim=-1, + ) + q_embed = self._rotary_embedding(q, phases) + k_embed = self._rotary_embedding(k, phases) + return q_embed, k_embed + + +class MultiHeadAttention(nn.Module): + def __init__( + self, + channels: int, + num_heads: int, + ctx_channels: Optional[int] = None, + type: Literal["self", "cross"] = "self", + attn_mode: Literal["full", "windowed"] = "full", + window_size: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + qkv_bias: bool = True, + use_rope: bool = False, + qk_rms_norm: bool = False, + ): + super().__init__() + assert channels % num_heads == 0 + assert type in ["self", "cross"], f"Invalid attention type: {type}" + assert attn_mode in ["full", "windowed"], f"Invalid attention mode: {attn_mode}" + assert ( + type == "self" or attn_mode == "full" + ), "Cross-attention only supports full attention" + + if attn_mode == "windowed": + raise NotImplementedError("Windowed attention is not yet implemented") + + self.channels = channels + self.head_dim = channels // num_heads + self.ctx_channels = ctx_channels if ctx_channels is not None else channels + self.num_heads = num_heads + self._type = type + self.attn_mode = attn_mode + self.window_size = window_size + self.shift_window = shift_window + self.use_rope = use_rope + self.qk_rms_norm = qk_rms_norm + + if self._type == "self": + self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias) + else: + self.to_q = nn.Linear(channels, channels, bias=qkv_bias) + self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias) + + if self.qk_rms_norm: + self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads) + self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads) + + self.to_out = nn.Linear(channels, channels) + + if use_rope: + self.rope = RotaryPositionEmbedder(channels) + + def forward( + self, + x: torch.Tensor, + context: Optional[torch.Tensor] = None, + indices: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + B, L, C = x.shape + if self._type == "self": + qkv = self.to_qkv(x) + qkv = qkv.reshape(B, L, 3, self.num_heads, -1) + if self.use_rope: + q, k, v = qkv.unbind(dim=2) + q, k = self.rope(q, k, indices) + qkv = torch.stack([q, k, v], dim=2) + if self.attn_mode == "full": + if self.qk_rms_norm: + q, k, v = qkv.unbind(dim=2) + q = self.q_rms_norm(q) + k = self.k_rms_norm(k) + h = scaled_dot_product_attention(q, k, v) + else: + h = scaled_dot_product_attention(qkv) + elif self.attn_mode == "windowed": + raise NotImplementedError("Windowed attention is not yet implemented") + else: + Lkv = context.shape[1] + q = self.to_q(x) + kv = self.to_kv(context) + q = q.reshape(B, L, self.num_heads, -1) + kv = kv.reshape(B, Lkv, 2, self.num_heads, -1) + if self.qk_rms_norm: + q = self.q_rms_norm(q) + k, v = kv.unbind(dim=2) + k = self.k_rms_norm(k) + h = scaled_dot_product_attention(q, k, v) + else: + h = scaled_dot_product_attention(q, kv) + h = h.reshape(B, L, -1) + h = self.to_out(h) + return h + + +class MOTMultiHeadSelfAttention(nn.Module): + def __init__( + self, + channels: int, + num_heads: int, + ctx_channels: Optional[int] = None, + type: Literal["self", "cross"] = "self", + attn_mode: Literal["full", "windowed"] = "full", ### this is not injection stuff + window_size: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + qkv_bias: bool = True, + use_rope: bool = False, + qk_rms_norm: bool = False, + latent_names: List = None, + protect_modality_list: List = ["shape"], + ): + super().__init__() + assert channels % num_heads == 0 + assert type in ["self", "cross"], f"Invalid attention type: {type}" + assert attn_mode in ["full", "windowed"], f"Invalid attention mode: {attn_mode}" + assert ( + type == "self" or attn_mode == "full" + ), "Cross-attention only supports full attention" + + if attn_mode == "windowed": + raise NotImplementedError("Windowed attention is not yet implemented") + + self.channels = channels + self.head_dim = channels // num_heads + self.ctx_channels = ctx_channels if ctx_channels is not None else channels + self.num_heads = num_heads + self._type = type + self.attn_mode = attn_mode + self.window_size = window_size + self.shift_window = shift_window + self.use_rope = use_rope + self.qk_rms_norm = qk_rms_norm + self.protect_modality_list = protect_modality_list + + self._cached_shape_k = None + self._cached_shape_v = None + # Per-timestep caches for inject mode (cond and uncond slots) + self._cached_shape_k_steps = [] # conditional branch cache + self._cached_shape_v_steps = [] + self._cached_shape_k_steps_uncond = [] # unconditional branch cache + self._cached_shape_v_steps_uncond = [] + self._cache_locked = False + self._cache_locked_uncond = False + self._cfg_branch = "cond" # set by _set_cfg_branch before each backbone call + self._inject_step_idx = 0 + + if self._type == "self": + self.to_qkv = torch.nn.ModuleDict( + { + latent_name: nn.Linear(channels, channels * 3, bias=qkv_bias) + for latent_name in latent_names + } + ) + else: + self.to_q = torch.nn.ModuleDict( + { + latent_name: nn.Linear(channels, channels, bias=qkv_bias) + for latent_name in latent_names + } + ) + self.to_kv = torch.nn.ModuleDict( + { + latent_name: nn.Linear( + self.ctx_channels, channels * 2, bias=qkv_bias + ) + for latent_name in latent_names + } + ) + + if self.qk_rms_norm: + self.q_rms_norm = torch.nn.ModuleDict( + { + latent_name: MultiHeadRMSNorm(self.head_dim, num_heads) + for latent_name in latent_names + } + ) + self.k_rms_norm = torch.nn.ModuleDict( + { + latent_name: MultiHeadRMSNorm(self.head_dim, num_heads) + for latent_name in latent_names + } + ) + + self.to_out = torch.nn.ModuleDict( + {latent_name: nn.Linear(channels, channels) for latent_name in latent_names} + ) + + if use_rope: + self.rope = RotaryPositionEmbedder(channels) + + def _reshape(self, qkv, tensor_shape, num_heads): + B, L, _ = tensor_shape + return qkv.reshape(B, L, 3, num_heads, -1) + + def _reshape_back(self, qkv, tensor_shape): + B, L, _ = tensor_shape + return qkv.reshape(B, L, -1) + + def _apply_module(self, x, module): + return module(x) + + # This is stupid, _pytree does not support ModuleDict + def _moduledict_to_dict(self, module): + return {key: module for key, module in module.items()} + + def unbind_qkv(self, qkv): + q, k, v = {}, {}, {} + for latent_name, _qkv in qkv.items(): + _q, _k, _v = _qkv.unbind(dim=2) + q[latent_name] = _q + k[latent_name] = _k + v[latent_name] = _v + + return q, k, v + + def _get_shape(self, x): + return x.shape + + def concatenate_tensor(self, tensor_dict, latent_names): + merged = [] + indicies_mapping = {} + total_tokens = 0 + for latent_name in latent_names: + merged.append(tensor_dict[latent_name]) + cur_token_len = tensor_dict[latent_name].shape[1] + indicies_mapping[latent_name] = [total_tokens, cur_token_len] + total_tokens += cur_token_len + # merge along token dimension + return torch.cat(merged, dim=1), indicies_mapping + + def unpack_tensors(self, h_others, indicies_mapping): + h = {} + for latent_name, (start, cur_token_len) in indicies_mapping.items(): + h[latent_name] = h_others[:, start : start + cur_token_len] + + return h + + def mm_scale_dot_product_attention(self, q, k, v, mode="default"): + h = {} + latent_names = list(q.keys()) + + if mode in ("default", "first_frame_inject"): + # Path 1: Shape self-attends to its own K/V (normal) + for protect_modality in self.protect_modality_list: + _q = q[protect_modality] + _k = k[protect_modality] + _v = v[protect_modality] + h[protect_modality] = scaled_dot_product_attention(_q, _k, _v) + + # Path 2: Layout queries attend to [layout K/V + shape K/V (detached)] + other_modalities = [ + n for n in latent_names if n not in self.protect_modality_list + ] + _q, indicies_mapping = self.concatenate_tensor(q, other_modalities) + o_k, _ = self.concatenate_tensor(k, other_modalities) + o_v, _ = self.concatenate_tensor(v, other_modalities) + _k, _ = self.concatenate_tensor(k, self.protect_modality_list) + _v, _ = self.concatenate_tensor(v, self.protect_modality_list) + # Cache per-timestep K/V during first_frame_inject, per branch. + # NOTE: shortcut model does not use CFG, so only the cond slot is used. + # The uncond slot code is kept for potential future CFG use. + if mode == "first_frame_inject": + branch = getattr(self, "_cfg_branch", "cond") + if branch == "cond" and not self._cache_locked: + self._cached_shape_k_steps.append(_k.detach().clone()) + self._cached_shape_v_steps.append(_v.detach().clone()) + self._cache_locked = True + elif branch == "uncond" and not self._cache_locked_uncond: + self._cached_shape_k_steps_uncond.append(_k.detach().clone()) + self._cached_shape_v_steps_uncond.append(_v.detach().clone()) + self._cache_locked_uncond = True + # NOTE: alternative that fills both slots from a single cond pass + # (for cases where CFG is not active during first_frame_inject): + # if branch == "cond" and not self._cache_locked: + # k_clone = _k.detach().clone() + # v_clone = _v.detach().clone() + # self._cached_shape_k_steps.append(k_clone) + # self._cached_shape_v_steps.append(v_clone) + # self._cached_shape_k_steps_uncond.append(k_clone) + # self._cached_shape_v_steps_uncond.append(v_clone) + # self._cache_locked = True + # self._cache_locked_uncond = True + _k = _k.detach() + _v = _v.detach() + _k = torch.cat([o_k, _k], dim=1) + _v = torch.cat([o_v, _v], dim=1) + h_others = scaled_dot_product_attention(_q, _k, _v) + h.update(self.unpack_tensors(h_others, indicies_mapping)) + + elif mode == "inject": + assert len(self._cached_shape_k_steps) > 0, \ + "Cache is empty! Run 'first_frame_inject' mode first." + idx = self._inject_step_idx + branch = getattr(self, "_cfg_branch", "cond") + if branch == "uncond" and len(self._cached_shape_k_steps_uncond) > 0: + cached_k = self._cached_shape_k_steps_uncond[idx] + cached_v = self._cached_shape_v_steps_uncond[idx] + else: + # Fall back to cond cache (backward compat or when uncond cache not populated) + cached_k = self._cached_shape_k_steps[idx] + cached_v = self._cached_shape_v_steps[idx] + + # Batched shape-KV: the cache may be primed at batch 1 (one first frame) + # while the query is a batch of B frames that all share that shape. Broadcast + # the batch-1 cache to B with expand (stride-0 view, no copy) so the persistent + # cache stays 1x; only the per-layer cat below is transiently materialized at B. + q_bs = q[self.protect_modality_list[0]].shape[0] + if cached_k.shape[0] == 1 and q_bs > 1: + cached_k = cached_k.expand(q_bs, *cached_k.shape[1:]) + cached_v = cached_v.expand(q_bs, *cached_v.shape[1:]) + + # Path 1: Shape self-attention uses matching-timestep cached K/V + for protect_modality in self.protect_modality_list: + _q = q[protect_modality] + h[protect_modality] = scaled_dot_product_attention(_q, cached_k, cached_v) + + # Path 2: Layout queries attend to [layout K/V + matching-timestep cached shape K/V] + other_modalities = [ + n for n in latent_names if n not in self.protect_modality_list + ] + _q, indicies_mapping = self.concatenate_tensor(q, other_modalities) + o_k, _ = self.concatenate_tensor(k, other_modalities) + o_v, _ = self.concatenate_tensor(v, other_modalities) + _k = torch.cat([o_k, cached_k], dim=1) + _v = torch.cat([o_v, cached_v], dim=1) + h_others = scaled_dot_product_attention(_q, _k, _v) + h.update(self.unpack_tensors(h_others, indicies_mapping)) + + return h + + def forward( + self, + x: Dict, + mode="default", + ) -> torch.Tensor: + shapes = _pytree.tree_map(self._get_shape, x) + if self._type == "self": + qkv = _pytree.tree_map( + self._apply_module, x, self._moduledict_to_dict(self.to_qkv) + ) + qkv = _pytree.tree_map( + partial(self._reshape, num_heads=self.num_heads), qkv, shapes + ) + if self.use_rope: + raise NotImplementedError + if self.attn_mode == "full": + if self.qk_rms_norm: + q, k, v = self.unbind_qkv(qkv) + q = _pytree.tree_map( + self._apply_module, q, self._moduledict_to_dict(self.q_rms_norm) + ) + k = _pytree.tree_map( + self._apply_module, k, self._moduledict_to_dict(self.k_rms_norm) + ) + h = self.mm_scale_dot_product_attention(q, k, v, mode=mode) + else: + raise NotImplementedError + elif self.attn_mode == "windowed": + raise NotImplementedError("Windowed attention is not yet implemented") + else: + raise NotImplementedError + + h = _pytree.tree_map(self._reshape_back, h, shapes) + h = _pytree.tree_map( + self._apply_module, h, self._moduledict_to_dict(self.to_out) + ) + return h diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/norm.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/norm.py new file mode 100644 index 0000000000000000000000000000000000000000..5e9c505fd7b2dc33dac3f9b9ce6df60ec241d385 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/norm.py @@ -0,0 +1,26 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import torch +import torch.nn as nn + + +class LayerNorm32(nn.LayerNorm): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return super().forward(x.float()).type(x.dtype) + + +class GroupNorm32(nn.GroupNorm): + """ + A GroupNorm layer that converts to float32 before the forward pass. + """ + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return super().forward(x.float()).type(x.dtype) + + +class ChannelLayerNorm32(LayerNorm32): + def forward(self, x: torch.Tensor) -> torch.Tensor: + DIM = x.dim() + x = x.permute(0, *range(2, DIM), 1).contiguous() + x = super().forward(x) + x = x.permute(0, DIM - 1, *range(1, DIM - 1)).contiguous() + return x diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/spatial.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/spatial.py new file mode 100644 index 0000000000000000000000000000000000000000..2a042b12a85cb077c4cc649d9671c1699bab9451 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/spatial.py @@ -0,0 +1,65 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import torch + + +def pixel_shuffle_3d(x: torch.Tensor, scale_factor: int) -> torch.Tensor: + """ + 3D pixel shuffle. + """ + B, C, H, W, D = x.shape + C_ = C // scale_factor**3 + x = x.reshape(B, C_, scale_factor, scale_factor, scale_factor, H, W, D) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4) + x = x.reshape(B, C_, H * scale_factor, W * scale_factor, D * scale_factor) + return x + + +def patchify(x: torch.Tensor, patch_size: int): + """ + Patchify a tensor. + + Args: + x (torch.Tensor): (N, C, *spatial) tensor + patch_size (int): Patch size + """ + DIM = x.dim() - 2 + for d in range(2, DIM + 2): + assert ( + x.shape[d] % patch_size == 0 + ), f"Dimension {d} of input tensor must be divisible by patch size, got {x.shape[d]} and {patch_size}" + + x = x.reshape( + *x.shape[:2], + *sum([[x.shape[d] // patch_size, patch_size] for d in range(2, DIM + 2)], []), + ) + x = x.permute( + 0, 1, *([2 * i + 3 for i in range(DIM)] + [2 * i + 2 for i in range(DIM)]) + ) + x = x.reshape(x.shape[0], x.shape[1] * (patch_size**DIM), *(x.shape[-DIM:])) + return x + + +def unpatchify(x: torch.Tensor, patch_size: int): + """ + Unpatchify a tensor. + + Args: + x (torch.Tensor): (N, C, *spatial) tensor + patch_size (int): Patch size + """ + DIM = x.dim() - 2 + assert ( + x.shape[1] % (patch_size**DIM) == 0 + ), f"Second dimension of input tensor must be divisible by patch size to unpatchify, got {x.shape[1]} and {patch_size ** DIM}" + + x = x.reshape( + x.shape[0], + x.shape[1] // (patch_size**DIM), + *([patch_size] * DIM), + *(x.shape[-DIM:]), + ) + x = x.permute(0, 1, *(sum([[2 + DIM + i, 2 + i] for i in range(DIM)], []))) + x = x.reshape( + x.shape[0], x.shape[1], *[x.shape[2 + 2 * i] * patch_size for i in range(DIM)] + ) + return x diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e94bfeea8b366b91cc24c284dd3e75b4aa7da4f8 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from .blocks import * +from .modulated import * diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c76387f6f2f6345159acd74a74677b79b6612236 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__pycache__/blocks.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__pycache__/blocks.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..881d073c61109c0afb3d26c060965037e369adcc Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__pycache__/blocks.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__pycache__/modulated.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__pycache__/modulated.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32154c105a3deac7a761aa02c5f191e7a45fd4df Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/__pycache__/modulated.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/blocks.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..bb8566c065f59fc454b87118cec1f98221ed1609 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/blocks.py @@ -0,0 +1,197 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from typing import * +import torch +import torch.nn as nn +from ..attention import MultiHeadAttention +from ..norm import LayerNorm32 + + +class AbsolutePositionEmbedder(nn.Module): + """ + Embeds spatial positions into vector representations. + """ + + def __init__(self, channels: int, in_channels: int = 3): + super().__init__() + self.channels = channels + self.in_channels = in_channels + self.freq_dim = channels // in_channels // 2 + self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim + self.freqs = 1.0 / (10000**self.freqs) + + def _sin_cos_embedding(self, x: torch.Tensor) -> torch.Tensor: + """ + Create sinusoidal position embeddings. + + Args: + x: a 1-D Tensor of N indices + + Returns: + an (N, D) Tensor of positional embeddings. + """ + self.freqs = self.freqs.to(x.device) + out = torch.outer(x, self.freqs) + out = torch.cat([torch.sin(out), torch.cos(out)], dim=-1) + return out + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x (torch.Tensor): (N, D) tensor of spatial positions + """ + N, D = x.shape + assert ( + D == self.in_channels + ), "Input dimension must match number of input channels" + embed = self._sin_cos_embedding(x.reshape(-1)) + embed = embed.reshape(N, -1) + if embed.shape[1] < self.channels: + embed = torch.cat( + [ + embed, + torch.zeros(N, self.channels - embed.shape[1], device=embed.device), + ], + dim=-1, + ) + return embed + + +class FeedForwardNet(nn.Module): + def __init__(self, channels: int, mlp_ratio: float = 4.0): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(channels, int(channels * mlp_ratio)), + nn.GELU(approximate="tanh"), + nn.Linear(int(channels * mlp_ratio), channels), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.mlp(x) + + +class TransformerBlock(nn.Module): + """ + Transformer block (MSA + FFN). + """ + + def __init__( + self, + channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "windowed"] = "full", + window_size: Optional[int] = None, + shift_window: Optional[int] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qkv_bias: bool = True, + ln_affine: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.attn = MultiHeadAttention( + channels, + num_heads=num_heads, + attn_mode=attn_mode, + window_size=window_size, + shift_window=shift_window, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.mlp = FeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + + def _forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.norm1(x) + h = self.attn(h) + x = x + h + h = self.norm2(x) + h = self.mlp(h) + x = x + h + return x + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint( + self._forward, x, use_reentrant=False + ) + else: + return self._forward(x) + + +class TransformerCrossBlock(nn.Module): + """ + Transformer cross-attention block (MSA + MCA + FFN). + """ + + def __init__( + self, + channels: int, + ctx_channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "windowed"] = "full", + window_size: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + qkv_bias: bool = True, + ln_affine: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.norm3 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.self_attn = MultiHeadAttention( + channels, + num_heads=num_heads, + type="self", + attn_mode=attn_mode, + window_size=window_size, + shift_window=shift_window, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.cross_attn = MultiHeadAttention( + channels, + ctx_channels=ctx_channels, + num_heads=num_heads, + type="cross", + attn_mode="full", + qkv_bias=qkv_bias, + qk_rms_norm=qk_rms_norm_cross, + ) + self.mlp = FeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + + def _forward(self, x: torch.Tensor, context: torch.Tensor): + h = self.norm1(x) + h = self.self_attn(h) + x = x + h + h = self.norm2(x) + h = self.cross_attn(h, context) + x = x + h + h = self.norm3(x) + h = self.mlp(h) + x = x + h + return x + + def forward(self, x: torch.Tensor, context: torch.Tensor): + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint( + self._forward, x, context, use_reentrant=False + ) + else: + return self._forward(x, context) diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/modulated.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/modulated.py new file mode 100644 index 0000000000000000000000000000000000000000..9552e40d0fe7f1e4c79e7397cdd5df0356c133c1 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/transformer/modulated.py @@ -0,0 +1,341 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from functools import partial +from typing import * +from torch.utils import _pytree +import torch +import torch.nn as nn +from ..attention import MultiHeadAttention, MOTMultiHeadSelfAttention +from ..norm import LayerNorm32 +from .blocks import FeedForwardNet + + +class ModulatedTransformerBlock(nn.Module): + """ + Transformer block (MSA + FFN) with adaptive layer norm conditioning. + """ + + def __init__( + self, + channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "windowed"] = "full", + window_size: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qkv_bias: bool = True, + share_mod: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.share_mod = share_mod + self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.attn = MultiHeadAttention( + channels, + num_heads=num_heads, + attn_mode=attn_mode, + window_size=window_size, + shift_window=shift_window, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.mlp = FeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + if not share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(channels, 6 * channels, bias=True) + ) + + def _forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor: + if self.share_mod: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk( + 6, dim=1 + ) + else: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.adaLN_modulation(mod).chunk(6, dim=1) + ) + h = self.norm1(x) + h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) + h = self.attn(h) + h = h * gate_msa.unsqueeze(1) + x = x + h + h = self.norm2(x) + h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) + h = self.mlp(h) + h = h * gate_mlp.unsqueeze(1) + x = x + h + return x + + def forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor: + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint( + self._forward, x, mod, use_reentrant=False + ) + else: + return self._forward(x, mod) + + +class ModulatedTransformerCrossBlock(nn.Module): + """ + Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning. + """ + + def __init__( + self, + channels: int, + ctx_channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "windowed"] = "full", + window_size: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + qkv_bias: bool = True, + share_mod: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.share_mod = share_mod + self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6) + self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.self_attn = MultiHeadAttention( + channels, + num_heads=num_heads, + type="self", + attn_mode=attn_mode, + window_size=window_size, + shift_window=shift_window, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.cross_attn = MultiHeadAttention( + channels, + ctx_channels=ctx_channels, + num_heads=num_heads, + type="cross", + attn_mode="full", + qkv_bias=qkv_bias, + qk_rms_norm=qk_rms_norm_cross, + ) + self.mlp = FeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + if not share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(channels, 6 * channels, bias=True) + ) + + def _forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor): + if self.share_mod: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk( + 6, dim=1 + ) + else: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.adaLN_modulation(mod).chunk(6, dim=1) + ) + h = self.norm1(x) + h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) + h = self.self_attn(h) + h = h * gate_msa.unsqueeze(1) + x = x + h + h = self.norm2(x) + h = self.cross_attn(h, context) + x = x + h + h = self.norm3(x) + h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) + h = self.mlp(h) + h = h * gate_mlp.unsqueeze(1) + x = x + h + return x + + def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor): + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint( + self._forward, x, mod, context, use_reentrant=False + ) + else: + return self._forward(x, mod, context) + + +class MOTModulatedTransformerCrossBlock(nn.Module): + """ + Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning. + """ + + def __init__( + self, + channels: int, + ctx_channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "windowed"] = "full", + window_size: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + qkv_bias: bool = True, + share_mod: bool = False, + latent_names: List = None, + freeze_shared_parameters: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.share_mod = share_mod + self.norm1 = torch.nn.ModuleDict( + { + latent_name: LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + for latent_name in latent_names + } + ) + self.norm2 = torch.nn.ModuleDict( + { + latent_name: LayerNorm32(channels, elementwise_affine=True, eps=1e-6) + for latent_name in latent_names + } + ) + self.norm3 = torch.nn.ModuleDict( + { + latent_name: LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + for latent_name in latent_names + } + ) + self.self_attn = MOTMultiHeadSelfAttention( + channels, + num_heads=num_heads, + type="self", + attn_mode=attn_mode, + window_size=window_size, + shift_window=shift_window, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + latent_names=latent_names, + ) + self.cross_attn = torch.nn.ModuleDict( + { + latent_name: MultiHeadAttention( + channels, + ctx_channels=ctx_channels, + num_heads=num_heads, + type="cross", + attn_mode="full", + qkv_bias=qkv_bias, + qk_rms_norm=qk_rms_norm_cross, + ) + for latent_name in latent_names + } + ) + self.mlp = torch.nn.ModuleDict( + { + latent_name: FeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + for latent_name in latent_names + } + ) + if not share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(channels, 6 * channels, bias=True) + ) + if freeze_shared_parameters: + self.adaLN_modulation.eval() + self.adaLN_modulation.requires_grad_(False) + + def _apply_module(self, h, module): + return module(h) + + def _apply_cross_attn(self, h, cross_attn, context): + return cross_attn(h, context) + + def _apply_msa(self, h, scale_msa, shift_msa): + return h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) + + def _apply_mlp(self, h, scale_mlp, shift_mlp): + return h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) + + def _apply_add(self, x, h): + return x + h + + def _apply_multiplication(self, h, multiplier): + return h * multiplier.unsqueeze(1) + + # This is stupid, _pytree does not support ModuleDict + def _moduledict_to_dict(self, module): + return {key: module for key, module in module.items()} + + def _forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor, mode="default"): + if self.share_mod: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk( + 6, dim=1 + ) + else: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.adaLN_modulation(mod).chunk(6, dim=1) + ) + h = _pytree.tree_map(self._apply_module, x, self._moduledict_to_dict(self.norm1)) + h = _pytree.tree_map( + partial(self._apply_msa, scale_msa=scale_msa, shift_msa=shift_msa), + h + ) + h = self.self_attn(h, mode=mode) + h = _pytree.tree_map( + partial(self._apply_multiplication, multiplier=gate_msa), + h + ) + x = _pytree.tree_map( + self._apply_add, + x, + h + ) + h = _pytree.tree_map(self._apply_module, x, self._moduledict_to_dict(self.norm2)) + h = _pytree.tree_map( + partial(self._apply_cross_attn, context=context), + h, + self._moduledict_to_dict(self.cross_attn), + ) + x = _pytree.tree_map( + self._apply_add, + x, + h + ) + h = _pytree.tree_map(self._apply_module, x, self._moduledict_to_dict(self.norm3)) + h = _pytree.tree_map( + partial(self._apply_mlp, scale_mlp=scale_mlp, shift_mlp=shift_mlp), + h + ) + h = _pytree.tree_map(self._apply_module, h, self._moduledict_to_dict(self.mlp)) + h = _pytree.tree_map( + partial(self._apply_multiplication, multiplier=gate_mlp), + h + ) + x = _pytree.tree_map( + self._apply_add, + x, + h + ) + return x + + def forward(self, x: Dict, mod: torch.Tensor, context: torch.Tensor, mode="default"): + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint( + self._forward, x, mod, context, mode, use_reentrant=False + ) + else: + return self._forward(x, mod, context, mode=mode) diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/utils.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..35f436a32e8259fde9a08d6b89a77527f2575617 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/backbone/tdfy_dit/modules/utils.py @@ -0,0 +1,72 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import torch.nn as nn +import torch + +FP16_TYPE = torch.float16 + +FP16_MODULES = [ + nn.Conv1d, + nn.Conv2d, + nn.Conv3d, + nn.ConvTranspose1d, + nn.ConvTranspose2d, + nn.ConvTranspose3d, + nn.Linear, +] + +# If we add sparse modules back in, they are compatible with FP16. +# But for now we don't have them and avoid the dependency on FlashAttention +# Instead using the torch implementation of FlashAttention in SDPA. +try: + from ..modules import sparse as sp + + FP16_MODULES += [ + sp.SparseConv3d, + sp.SparseInverseConv3d, + sp.SparseLinear, + ] +except ImportError: + pass + +FP16_MODULES = tuple(FP16_MODULES) + + +def convert_module_to_f16(l): + """ + Convert primitive modules to float16. + """ + if isinstance(l, FP16_MODULES): + for p in l.parameters(): + # p.data = p.data.half() + p.data = p.data.to(FP16_TYPE) + + +def convert_module_to_f32(l): + """ + Convert primitive modules to float32, undoing convert_module_to_f16(). + """ + if isinstance(l, FP16_MODULES): + for p in l.parameters(): + p.data = p.data.float() + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module + + +def scale_module(module, scale): + """ + Scale the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().mul_(scale) + return module + + +def modulate(x, shift, scale): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..930ccc32d2b80319b8e2fc785c3cee7d3eff14ce Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/llama3/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/llama3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/llama3/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/llama3/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a9c2694fe10368a142ed31ab72fee38b9f07b33a Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/llama3/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/llama3/__pycache__/ff.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/llama3/__pycache__/ff.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c8457f7c75da726a6529ec9573a4b5308c053e2 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/llama3/__pycache__/ff.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/llama3/ff.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/llama3/ff.py new file mode 100644 index 0000000000000000000000000000000000000000..ddde75d4ca01fe1aef5750a6ba5ef0ad88e8d5e5 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/model/layers/llama3/ff.py @@ -0,0 +1,41 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from torch import nn +import torch.nn.functional as F +from typing import Optional + + +class FeedForward(nn.Module): + def __init__( + self, + dim: int, + hidden_dim: int, + multiple_of: int = 256, + ffn_dim_multiplier: Optional[float] = None, + output_dim: Optional[int] = None, + skip_w2: bool = False, + ): + """ + Llama3 FeedForward layer + https://github.com/meta-llama/llama3/blob/a0940f9cf7065d45bb6675660f80d305c041a754/llama/model.py#L193 + """ + super().__init__() + hidden_dim = int(2 * hidden_dim / 3) + # custom dim factor multiplier + if ffn_dim_multiplier is not None: + hidden_dim = int(ffn_dim_multiplier * hidden_dim) + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + + if output_dim is None: + output_dim = dim + + self.skip_w2 = skip_w2 + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + if not self.skip_w2: + self.w2 = nn.Linear(hidden_dim, output_dim, bias=False) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + + def forward(self, x): + x = F.silu(self.w1(x)) * self.w3(x) + if self.skip_w2: + return x + return self.w2(x) diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/pipeline/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/pipeline/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/pipeline/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/pipeline/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65f781c6f3ab16074b76f0d1299953e70fe0d0f0 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/pipeline/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/pipeline/__pycache__/inference_utils.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/pipeline/__pycache__/inference_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bd229b50d2f126e25cf6ab2b24c637820aa54e53 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/pipeline/__pycache__/inference_utils.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/pipeline/inference_utils.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/pipeline/inference_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..684b98e38441b3fc0c3af4ed84a70a7da39de29d --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/pipeline/inference_utils.py @@ -0,0 +1,63 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +"""Standalone copy of the three voxelization helpers from the SAM-3D fork's +``sam3d_objects/pipeline/inference_utils.py``. + +These three functions (``normalize_mesh_verts``, ``voxelize_mesh``, +``trimesh2o3d_mesh``) are reproduced VERBATIM from the fork (only the surrounding +heavy imports — pytorch3d / layout-optimization — are dropped, since these +functions touch none of them). Keeping them byte-identical guarantees the SS +encoder sees exactly the occupancy distribution it was trained on. +""" +import numpy as np +import open3d as o3d +import torch + + +def normalize_mesh_verts(verts): + vmin = verts.min(axis=0) + vmax = verts.max(axis=0) + center = (vmax + vmin) / 2.0 + extent = vmax - vmin # largest side length + max_extent = np.max(extent) + if max_extent == 0: + vertices = verts - center + scale = 1 + else: + scale = 1.0 / max_extent + vertices = (verts - center) * scale + return vertices, scale, center + + +def voxelize_mesh(mesh, resolution=64): + verts = np.asarray(mesh.vertices) + # rotate mesh (from z-up to y-up) + verts = verts @ np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]]).T + # normalize vertices + # skip vertices to avoid losing points, likely already normalized + if np.abs(verts.min() + 0.5) < 1e-3 and np.abs(verts.max() - 0.5) < 1e-3: + vertices, scale, center = verts, None, None + else: + vertices, scale, center = normalize_mesh_verts(verts) + + vertices = np.clip(vertices, -0.5 + 1e-6, 0.5 - 1e-6) + mesh.vertices = o3d.utility.Vector3dVector(vertices) + voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds( + mesh, + voxel_size=1 / 64, + min_bound=(-0.5, -0.5, -0.5), + max_bound=(0.5, 0.5, 0.5), + ) + vertices = np.array([voxel.grid_index for voxel in voxel_grid.get_voxels()]) + vertices = (vertices + 0.5) / 64 - 0.5 + coords = ((torch.tensor(vertices) + 0.5) * resolution).int().contiguous() + ss = torch.zeros(1, resolution, resolution, resolution, dtype=torch.long) + ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1 + return ss, scale, center + + +def trimesh2o3d_mesh(trimesh_mesh): + verts = np.asarray(trimesh_mesh.vertices) + faces = np.asarray(trimesh_mesh.faces) + return o3d.geometry.TriangleMesh( + o3d.utility.Vector3dVector(verts), o3d.utility.Vector3iVector(faces) + ) diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ae1c71ce7c2324e7fbf2f441bf7bba1de521a7b7 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__init__.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/__init__.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b2428d3d55ffd87fe4cdef1b6790780a8bd91aa8 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/__init__.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/checkpointing.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/checkpointing.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ece815ceb8aae2ae6e6027093ddbbb90ba1105d5 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/checkpointing.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/distributed.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/distributed.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9df3304949f0b4f2bfd1998c74dc6c63bc34db49 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/distributed.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/train_utils.cpython-311.pyc b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/train_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4f7ac4a40f8ee55d2ef13569f88d20003455fb4 Binary files /dev/null and b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/__pycache__/train_utils.cpython-311.pyc differ diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/checkpointing.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/checkpointing.py new file mode 100644 index 0000000000000000000000000000000000000000..543bd0ff2d611d41d1f06e1c593294f7d095eeae --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/checkpointing.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch + + + +def unwrap_model(model: torch.nn.Module) -> torch.nn.Module: + return model.module if hasattr(model, "module") else model + + + +def trainable_state_dict(model: torch.nn.Module) -> dict[str, torch.Tensor]: + m = unwrap_model(model) + required = {name for name, p in m.named_parameters() if p.requires_grad} + full = m.state_dict() + return {k: v for k, v in full.items() if k in required} + + + +def optimizer_sidecar_path(checkpoint_path: str) -> str: + """Path of the optimizer sidecar for a given model checkpoint. + + ``.../step_0002000.pt`` -> ``.../step_0002000_optimizer.pt``. + """ + p = Path(checkpoint_path) + return str(p.with_name(p.stem + "_optimizer.pt")) + + +def save_checkpoint( + checkpoint_path: str, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + scheduler, + scaler, + step: int, + config_dict: dict[str, Any], + history: dict[str, Any] | None = None, +) -> None: + """[custom edit to reused fork] Write model and optimizer to SEPARATE files. + + ``step_XXXX.pt`` = model (trainable) + config + step + history + ``step_XXXX_optimizer.pt`` = optimizer + scheduler + scaler + + Splitting them lets optimizer states be pruned with a plain + ``rm *_optimizer.pt`` without rewriting model weights, and keeps the + checkpoints that inference/eval load small. Backward compatible: the model + file is written first (always self-contained for inference). + """ + checkpoint_path = str(Path(checkpoint_path).resolve()) + Path(checkpoint_path).parent.mkdir(parents=True, exist_ok=True) + model_payload = { + "step": int(step), + "trainable_state_dict": trainable_state_dict(model), + "config": config_dict, + "history": history if history is not None else {}, + } + torch.save(model_payload, checkpoint_path) + opt_payload = { + "step": int(step), + "optimizer": optimizer.state_dict(), + "scheduler": scheduler.state_dict() if scheduler is not None else None, + "scaler": scaler.state_dict() if scaler is not None else None, + } + torch.save(opt_payload, optimizer_sidecar_path(checkpoint_path)) + + + +def load_checkpoint( + checkpoint_path: str, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer | None = None, + scheduler=None, + scaler=None, + map_location: str = "cpu", +) -> tuple[int, dict[str, Any]]: + """[custom edit to reused fork] Restore model + (optionally) optimizer state. + + Optimizer/scheduler/scaler are read from whichever source exists: + (a) the new sidecar ``step_XXXX_optimizer.pt`` (split format), or + (b) the legacy bundled keys inside the main file (old single-file ckpts). + Missing optimizer state is skipped silently, so model-only checkpoints load + fine for inference/eval. + """ + checkpoint_path = str(Path(checkpoint_path).resolve()) + ckpt = torch.load(checkpoint_path, map_location=map_location, weights_only=False) + m = unwrap_model(model) + m.load_state_dict(ckpt["trainable_state_dict"], strict=False) + + # Resolve optimizer-state source: legacy bundled keys, else the sidecar file. + opt_state = ckpt.get("optimizer") + sched_state = ckpt.get("scheduler") + scaler_state = ckpt.get("scaler") + if opt_state is None: + sidecar = optimizer_sidecar_path(checkpoint_path) + if Path(sidecar).is_file(): + side = torch.load(sidecar, map_location=map_location, weights_only=False) + opt_state = side.get("optimizer") + sched_state = side.get("scheduler") + scaler_state = side.get("scaler") + + if optimizer is not None and opt_state is not None: + optimizer.load_state_dict(opt_state) + if scheduler is not None and sched_state is not None: + scheduler.load_state_dict(sched_state) + if scaler is not None and scaler_state is not None: + scaler.load_state_dict(scaler_state) + return int(ckpt.get("step", 0)), ckpt.get("history", {}) diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/distributed.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/distributed.py new file mode 100644 index 0000000000000000000000000000000000000000..21e26ebf3b90e26e3c0cdf08edc4f1ef614c3bc9 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/distributed.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +from dataclasses import dataclass +import json +import os + +import torch +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP + + +@dataclass +class DistributedState: + rank: int + local_rank: int + world_size: int + distributed: bool + device: torch.device + backend: str + + @property + def is_main(self) -> bool: + return self.rank == 0 + + +@dataclass +class RuntimeHandles: + model: torch.nn.Module + optimizer: torch.optim.Optimizer + scheduler: object + use_deepspeed: bool + + + +def init_distributed(backend: str = "ddp") -> DistributedState: + rank = int(os.environ.get("RANK", "0")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + else: + device = torch.device("cpu") + + distributed = world_size > 1 + if distributed and not dist.is_initialized(): + dist_backend = "nccl" if device.type == "cuda" else "gloo" + dist.init_process_group(backend=dist_backend, init_method="env://") + + return DistributedState( + rank=rank, + local_rank=local_rank, + world_size=world_size, + distributed=distributed, + device=device, + backend=backend.lower().strip(), + ) + + + +def setup_runtime( + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + scheduler, + state: DistributedState, + deepspeed_config_path: str, +) -> RuntimeHandles: + backend = state.backend + + if backend == "deepspeed": + try: + import deepspeed + except Exception as exc: + if state.is_main: + print( + f"[distributed] DeepSpeed unavailable ({exc}); falling back to DDP.", + flush=True, + ) + backend = "ddp" + else: + cfg_path = os.path.abspath(deepspeed_config_path) + if not os.path.isfile(cfg_path): + raise FileNotFoundError(f"DeepSpeed config not found: {cfg_path}") + with open(cfg_path, "r", encoding="utf-8") as f: + _ = json.load(f) + + engine, ds_optimizer, _, ds_scheduler = deepspeed.initialize( + model=model, + model_parameters=[p for p in model.parameters() if p.requires_grad], + optimizer=optimizer, + lr_scheduler=scheduler, + config=cfg_path, + ) + if state.is_main: + print(f"[distributed] DeepSpeed initialized with {cfg_path}", flush=True) + return RuntimeHandles( + model=engine, + optimizer=ds_optimizer, + scheduler=ds_scheduler, + use_deepspeed=True, + ) + + wrapped_model = model + if state.distributed: + wrapped_model = DDP( + model, + device_ids=[state.local_rank] if state.device.type == "cuda" else None, + output_device=state.local_rank if state.device.type == "cuda" else None, + find_unused_parameters=False, + ) + return RuntimeHandles( + model=wrapped_model, + optimizer=optimizer, + scheduler=scheduler, + use_deepspeed=False, + ) + + + +def barrier() -> None: + if dist.is_available() and dist.is_initialized(): + if torch.cuda.is_available(): + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + dist.barrier(device_ids=[local_rank]) + else: + dist.barrier() + + + +def cleanup_distributed() -> None: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + + +def reduce_scalar_dict(metrics: dict[str, torch.Tensor], state: DistributedState) -> dict[str, float]: + out: dict[str, float] = {} + for key, value in metrics.items(): + tensor = value.detach().to(state.device, non_blocking=True).float() + if state.distributed: + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + tensor = tensor / float(state.world_size) + out[key] = float(tensor.item()) + return out diff --git a/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/train_utils.py b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/train_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f3b889b79df14434eb48e1dba4a9a71ae154243e --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sam3d_objects/training/stage1_flow_lora/train_utils.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import math +from contextlib import nullcontext + +import torch + + + +def get_autocast_dtype(amp_dtype: str) -> torch.dtype: + amp_dtype = amp_dtype.lower().strip() + if amp_dtype in {"none", "off", "fp32", "float32"}: + raise ValueError("No autocast dtype for disabled autocast mode") + if amp_dtype == "bf16": + return torch.bfloat16 + if amp_dtype in {"fp16", "float16"}: + return torch.float16 + raise ValueError(f"Unsupported amp dtype: {amp_dtype}") + + + +def get_autocast_context(device: torch.device, amp_dtype: str): + if device.type != "cuda": + return nullcontext() + if amp_dtype.lower().strip() in {"none", "off", "fp32", "float32"}: + return nullcontext() + dtype = get_autocast_dtype(amp_dtype) + # cache_enabled=False is REQUIRED for the 2-pass forward: Pass 1 (first_frame + # shape K/V) runs under torch.no_grad(); with the autocast weight cache on, the + # trainable layout weights get cast to bf16 and cached DETACHED during Pass 1, + # then Pass 2 (grad on) reuses the detached cache -> v_pred loses grad_fn and + # loss.backward() fails. (LoRA dodged this via disable_adapter in Pass 1.) + return torch.autocast(device_type="cuda", dtype=dtype, cache_enabled=False) + + + +def build_scheduler(optimizer: torch.optim.Optimizer, train_cfg): + total = int(train_cfg.max_steps) + warmup = int(train_cfg.warmup_steps) + min_lr = float(train_cfg.min_lr) + base_lr = float(train_cfg.lr) + + if train_cfg.scheduler.lower() == "constant": + return torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0) + + floor_ratio = min_lr / max(base_lr, 1e-12) + + def lr_lambda(step: int) -> float: + if step < warmup: + return max(1e-8, float(step + 1) / max(1, warmup)) + progress = (step - warmup) / max(1, total - warmup) + progress = min(1.0, max(0.0, progress)) + cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) + return floor_ratio + (1.0 - floor_ratio) * cosine + + return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) + + + +def format_metrics(step: int, lr: float, metrics: dict[str, float]) -> str: + ordered = [ + f"step={step}", + f"lr={lr:.6e}", + f"loss={metrics.get('loss', float('nan')):.6f}", + f"rot={metrics.get('loss_rot', float('nan')):.6f}", + f"trans={metrics.get('loss_trans', float('nan')):.6f}", + f"scale={metrics.get('loss_scale', float('nan')):.6f}", + f"rot_deg={metrics.get('metric_rot_deg', float('nan')):.4f}", + f"trans_l2_m={metrics.get('metric_trans_l2_m', float('nan')):.6f}", + ] + if "translation_valid_samples" in metrics: + ordered.append(f"trans_valid={int(round(metrics['translation_valid_samples']))}") + if "translation_skipped_samples" in metrics: + ordered.append(f"trans_skip={int(round(metrics['translation_skipped_samples']))}") + return " | ".join(ordered) diff --git a/migrator/code/mv-sam3d-for-6d-v2/sbatch/infer_smoke.sh b/migrator/code/mv-sam3d-for-6d-v2/sbatch/infer_smoke.sh new file mode 100644 index 0000000000000000000000000000000000000000..fd8f4a61159c989ab359c34948687f412ba6a13d --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sbatch/infer_smoke.sh @@ -0,0 +1,71 @@ +#!/bin/sh +#SBATCH --job-name=mv2_infer_smoke +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=12 +#SBATCH --time=00:30:00 +#SBATCH --output=%x_%j.log +#SBATCH --error=%x_%j.log + +# Smoke test for the v2 inference scripts. GPU only; never sets +# CUDA_VISIBLE_DEVICES. Runs SCRIPT 1 (tools/infer_demo.py) and SCRIPT 2 +# (mvsam3d/infer/predict.py CLI) against one checkpoint on a few subject-01 +# DexYCB samples (large frame_stride for speed). +set -eu +if [ -z "${MIGRATOR_ROOT:-}" ]; then HERE="$(cd "$(dirname "$0")" && pwd)"; MIGRATOR_ROOT="$(cd "$HERE/../../.." && pwd)"; fi +export MIGRATOR_ROOT +REPO="$MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2" +PY="${MIGRATOR_PY:-python}" +export MIGRATOR_DATA_ROOT="${MIGRATOR_DATA_ROOT:-$MIGRATOR_ROOT/data}" + +REPO="$MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2" +CKPT="${CKPT:-$REPO/cache/ckpt/step_0014000.pt}" +OUTDIR="${OUTDIR:-$REPO/cache/infer_smoke}" + +cd "$REPO" +# Standard env exports (same as sbatch/train.sh). +export PYTHONUNBUFFERED=1 +export PYTHONPATH="$REPO" +export OMP_NUM_THREADS=8 +export ATTN_BACKEND=sdpa +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +mkdir -p "$OUTDIR" +echo "[smoke] host=$(hostname) date=$(date) CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-unset}" +echo "[smoke] checkpoint=$CKPT" +nvidia-smi --query-gpu=index,name,memory.free --format=csv + +# Fast smoke knobs: one subject, large frame stride, few aux views, few ODE steps. +SUBJECTS="subject-01" +FRAME_STRIDE=300 +NUM_AUX=2 +NUM_STEPS=20 + +echo "==================================================================" +echo "[smoke] SCRIPT 1 — tools/infer_demo.py (idx 0)" +echo "==================================================================" +$PY -m tools.infer_demo \ + --checkpoint "$CKPT" \ + --idx 0 \ + --num-steps "$NUM_STEPS" \ + --subjects "$SUBJECTS" \ + --frame-stride "$FRAME_STRIDE" \ + --num-aux-views "$NUM_AUX" + +echo "==================================================================" +echo "[smoke] SCRIPT 2 — mvsam3d/infer/predict.py CLI (4 samples -> .pt/.json)" +echo "==================================================================" +$PY -m mvsam3d.infer.predict \ + --checkpoint "$CKPT" \ + --out "$OUTDIR/preds.pt" \ + --num-samples 4 \ + --num-steps "$NUM_STEPS" \ + --subjects "$SUBJECTS" \ + --frame-stride "$FRAME_STRIDE" \ + --num-aux-views "$NUM_AUX" + +echo "[smoke] output files:" +ls -la "$OUTDIR" +echo "[smoke] done" diff --git a/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_dexh2r.sh b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_dexh2r.sh new file mode 100644 index 0000000000000000000000000000000000000000..190c9ed38fca4eae7ae65a041f16be1e8771860b --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_dexh2r.sh @@ -0,0 +1,24 @@ +#!/bin/sh +#SBATCH --job-name=lat_dexh2r +#SBATCH --partition=background +#SBATCH --qos=preempt +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --time=02:00:00 +#SBATCH --output=%x_%j.log +#SBATCH --error=%x_%j.log +set -e +if [ -z "${MIGRATOR_ROOT:-}" ]; then HERE="$(cd "$(dirname "$0")" && pwd)"; MIGRATOR_ROOT="$(cd "$HERE/../../.." && pwd)"; fi +export MIGRATOR_ROOT +REPO="$MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2" +PY="${MIGRATOR_PY:-python}" +export MIGRATOR_DATA_ROOT="${MIGRATOR_DATA_ROOT:-$MIGRATOR_ROOT/data}" +cd "$REPO" +export PYTHONPATH="$REPO" +# Root points at whatever DexH2R tree exposes all 56 object dirs under extracted/; +# latents are mesh-only (object_model/.obj) and namespaced dexh2r/. +ROOT="${DEXH2R_ROOT:-$MIGRATOR_DATA_ROOT/dexh2r}" +$PY tools/build_canonical_latents_multi.py --dataset dexh2r --root "$ROOT" +echo "[lat_dexh2r DONE]" diff --git a/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_graspnet.sh b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_graspnet.sh new file mode 100644 index 0000000000000000000000000000000000000000..c4a77792f4ef2c63fd1b40fdb9562117e1b31f39 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_graspnet.sh @@ -0,0 +1,21 @@ +#!/bin/sh +#SBATCH --job-name=lat_gn +#SBATCH --partition=background +#SBATCH --qos=preempt +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --time=02:00:00 +#SBATCH --output=%x_%j.log +#SBATCH --error=%x_%j.log +set -e +if [ -z "${MIGRATOR_ROOT:-}" ]; then HERE="$(cd "$(dirname "$0")" && pwd)"; MIGRATOR_ROOT="$(cd "$HERE/../../.." && pwd)"; fi +export MIGRATOR_ROOT +REPO="$MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2" +PY="${MIGRATOR_PY:-python}" +export MIGRATOR_DATA_ROOT="${MIGRATOR_DATA_ROOT:-$MIGRATOR_ROOT/data}" +cd "$REPO" +export PYTHONPATH="$REPO" +$PY tools/build_canonical_latents_multi.py --dataset graspnet --root "$MIGRATOR_DATA_ROOT/graspnet" +echo "[lat_gn DONE]" diff --git a/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_h2o.sh b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_h2o.sh new file mode 100644 index 0000000000000000000000000000000000000000..f7085d9a1272891aa7d1dfd016418ebe0779e353 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_h2o.sh @@ -0,0 +1,21 @@ +#!/bin/sh +#SBATCH --job-name=lat_h2o +#SBATCH --partition=background +#SBATCH --qos=preempt +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --time=02:00:00 +#SBATCH --output=%x_%j.log +#SBATCH --error=%x_%j.log +set -e +if [ -z "${MIGRATOR_ROOT:-}" ]; then HERE="$(cd "$(dirname "$0")" && pwd)"; MIGRATOR_ROOT="$(cd "$HERE/../../.." && pwd)"; fi +export MIGRATOR_ROOT +REPO="$MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2" +PY="${MIGRATOR_PY:-python}" +export MIGRATOR_DATA_ROOT="${MIGRATOR_DATA_ROOT:-$MIGRATOR_ROOT/data}" +cd "$REPO" +export PYTHONPATH="$REPO" +$PY tools/build_canonical_latents_multi.py --dataset h2o --root "$MIGRATOR_DATA_ROOT/h2o" +echo "[lat_h2o DONE]" diff --git a/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_hocap_cp.sh b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_hocap_cp.sh new file mode 100644 index 0000000000000000000000000000000000000000..741765ac0dca038a17a6c12245306ee8800c70bb --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_hocap_cp.sh @@ -0,0 +1,23 @@ +#!/bin/sh +#SBATCH --job-name=lat_hc_cp +#SBATCH --partition=background +#SBATCH --qos=preempt +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --time=02:00:00 +#SBATCH --output=%x_%j.log +#SBATCH --error=%x_%j.log +set -e +if [ -z "${MIGRATOR_ROOT:-}" ]; then HERE="$(cd "$(dirname "$0")" && pwd)"; MIGRATOR_ROOT="$(cd "$HERE/../../.." && pwd)"; fi +export MIGRATOR_ROOT +REPO="$MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2" +PY="${MIGRATOR_PY:-python}" +export MIGRATOR_DATA_ROOT="${MIGRATOR_DATA_ROOT:-$MIGRATOR_ROOT/data}" +cd "$REPO" +export PYTHONPATH="$REPO" +DS="$MIGRATOR_DATA_ROOT" +$PY tools/build_canonical_latents_multi.py --dataset contactpose --root $DS/contactpose +$PY tools/build_canonical_latents_multi.py --dataset hocap --root $DS/hocap +echo "[lat_hc_cp DONE]" diff --git a/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_hograspnet.sh b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_hograspnet.sh new file mode 100644 index 0000000000000000000000000000000000000000..a924d53fdd771d956082174896d6291db355ff5f --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_hograspnet.sh @@ -0,0 +1,21 @@ +#!/bin/sh +#SBATCH --job-name=lat_hog +#SBATCH --partition=background +#SBATCH --qos=preempt +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --time=02:00:00 +#SBATCH --output=%x_%j.log +#SBATCH --error=%x_%j.log +set -e +if [ -z "${MIGRATOR_ROOT:-}" ]; then HERE="$(cd "$(dirname "$0")" && pwd)"; MIGRATOR_ROOT="$(cd "$HERE/../../.." && pwd)"; fi +export MIGRATOR_ROOT +REPO="$MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2" +PY="${MIGRATOR_PY:-python}" +export MIGRATOR_DATA_ROOT="${MIGRATOR_DATA_ROOT:-$MIGRATOR_ROOT/data}" +cd "$REPO" +export PYTHONPATH="$REPO" +$PY tools/build_canonical_latents_multi.py --dataset hograspnet --root "$MIGRATOR_DATA_ROOT/hograspnet" +echo "[lat_hog DONE]" diff --git a/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_hograspnet_all30.sh b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_hograspnet_all30.sh new file mode 100644 index 0000000000000000000000000000000000000000..f04a81a3553d31b0b272d27515da554fcb6fb11f --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sbatch/lat_hograspnet_all30.sh @@ -0,0 +1,57 @@ +#!/bin/sh +#SBATCH --job-name=lat_hog30 +#SBATCH --partition=background +#SBATCH --qos=preempt +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --time=02:00:00 +#SBATCH --output=%x_%j.log +#SBATCH --error=%x_%j.log +set -e +if [ -z "${MIGRATOR_ROOT:-}" ]; then HERE="$(cd "$(dirname "$0")" && pwd)"; MIGRATOR_ROOT="$(cd "$HERE/../../.." && pwd)"; fi +export MIGRATOR_ROOT +REPO="$MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2" +PY="${MIGRATOR_PY:-python}" +export MIGRATOR_DATA_ROOT="${MIGRATOR_DATA_ROOT:-$MIGRATOR_ROOT/data}" +cd "$REPO" +export PYTHONPATH="$REPO" +$PY - <<'PYEOF' +# Build canonical latents for ALL 30 HOGraspNet scanned meshes (idempotent). +# Reuses the SAME path as tools/build_canonical_latents_multi.py: adapter.mesh_for +# (applies the mandatory per-object scale-fix), voxelize_canonical, encode. +# Independent of the frame download since latents are purely mesh-derived. +import os, numpy as np, torch +from mvsam3d.data.canonical import voxelize_canonical +from mvsam3d.data.adapters import HOGraspNetAdapter +from mvsam3d.data.base_dataset import ViewSample +from mvsam3d.model.ss_vae import encode_occupancy, load_ss_encoder + +ROOT = os.path.expandvars("$MIGRATOR_DATA_ROOT/hograspnet") +CACHE = os.path.expandvars("$MIGRATOR_CACHE")+"/canonical" +ad = HOGraspNetAdapter(ROOT, CACHE) +objmap = ad._objdir() +enc = load_ss_encoder("cuda") +done = skip = 0 +for oi in sorted(objmap): + objdir = objmap[oi]; key = f"hograspnet/{objdir}" + out = os.path.join(CACHE, f"{key}.pt") + if os.path.exists(out): + skip += 1; continue + s = ViewSample(anchor=None, aux_pool=[], obj_key=key, meta={"obj_idx": oi, "objdir": objdir}) + mesh = ad.mesh_for(s) + occ, ext, cen = voxelize_canonical(mesh, resolution=64) + occ_f = occ.float().unsqueeze(0) if occ.ndim == 4 else occ.float().unsqueeze(0).unsqueeze(0) + latent = encode_occupancy(enc, occ_f) + os.makedirs(os.path.dirname(out), exist_ok=True) + torch.save({"latent": latent.squeeze(0).cpu(), + "occupancy": (occ.squeeze(0) if occ.ndim == 4 else occ).to(torch.uint8).cpu(), + "max_extent": float(ext), + "center_cad": np.asarray(cen, np.float64)}, out) + done += 1 + print(f" built {key} extent={float(ext):.3f}", flush=True) +print(f"[all30] built {done}, skipped {skip}; total now: " + f"{len(os.listdir(os.path.join(CACHE,'hograspnet')))}", flush=True) +PYEOF +echo "[lat_hog30 DONE]" diff --git a/migrator/code/mv-sam3d-for-6d-v2/sbatch/train.sh b/migrator/code/mv-sam3d-for-6d-v2/sbatch/train.sh new file mode 100644 index 0000000000000000000000000000000000000000..fbebf06dab68cb930a523ff5c78ccda0d2de52af --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sbatch/train.sh @@ -0,0 +1,57 @@ +#!/bin/sh +#SBATCH --job-name=mv2_train +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=12 +#SBATCH --time=24:00:00 +#SBATCH --output=mv2_train_%j.out +#SBATCH --error=mv2_train_%j.err + +# Train the multi-view 6D-pose flow model. GPU only; never sets +# CUDA_VISIBLE_DEVICES. All hyperparameters come from the YAML config. +# +# PORTABLE: set MIGRATOR_ROOT (package root that contains code/ checkpoints/ ...) +# and, if the conda env is not already active, MIGRATOR_PY (python interpreter). +# Datasets live under MIGRATOR_DATA_ROOT (default $MIGRATOR_ROOT/data). +# +# export MIGRATOR_ROOT=/path/to/migrator +# export MIGRATOR_DATA_ROOT=/path/to/datasets # optional +# sbatch sbatch/train.sh --config configs/train.yaml +# sbatch sbatch/train.sh --config configs/train.yaml --set train.steps=60 +set -eu + +# --- resolve package root --------------------------------------------------- +if [ -z "${MIGRATOR_ROOT:-}" ]; then + # this script lives at $MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2/sbatch/train.sh + HERE="$(cd "$(dirname "$0")" && pwd)" + MIGRATOR_ROOT="$(cd "$HERE/../../.." && pwd)" +fi +export MIGRATOR_ROOT +REPO="$MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2" +PY="${MIGRATOR_PY:-python}" + +cd "$REPO" +export PYTHONUNBUFFERED=1 +# The sam3d_objects backbone/embedder/SS-VAE closure is VENDORED in-repo under +# sam3d_objects/. ONLY the repo is on PYTHONPATH (the repo's partial +# sam3d_objects stays primary); the fork's full sam3d_objects is resolved at +# runtime via package __path__ extension (see mvsam3d.paths.FORK_ROOT). +export PYTHONPATH="$REPO" +export OMP_NUM_THREADS=8 +export ATTN_BACKEND=sdpa +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +echo "[train] host=$(hostname) date=$(date) MIGRATOR_ROOT=$MIGRATOR_ROOT CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-unset}" +nvidia-smi --query-gpu=index,name,memory.free --format=csv + +# One process per allocated GPU. DDP device selection is by LOCAL_RANK +# (torch.cuda.set_device) inside distributed.init_distributed — we never set +# CUDA_VISIBLE_DEVICES. Override the process count with NPROC=N if needed. +NPROC="${NPROC:-${SLURM_GPUS_ON_NODE:-1}}" +echo "[train] launching torchrun --nproc_per_node=${NPROC}" + +$PY -m torch.distributed.run --standalone --nproc_per_node="${NPROC}" -m mvsam3d.train.train "$@" + +echo "[train] done" diff --git a/migrator/code/mv-sam3d-for-6d-v2/sbatch/watchdog_resume.sh b/migrator/code/mv-sam3d-for-6d-v2/sbatch/watchdog_resume.sh new file mode 100644 index 0000000000000000000000000000000000000000..6f5fa5bb45b974267d234f6b74bddcaa37e5c197 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/sbatch/watchdog_resume.sh @@ -0,0 +1,67 @@ +#!/bin/bash +#SBATCH --job-name=mv2_watchdog +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=1 +#SBATCH --time=00:10:00 +#SBATCH --output=%x_%j.log +#SBATCH --error=%x_%j.log + +# Watchdog: runs after the training job ends. +# - If step_0100000.pt exists, training is done -> exit. +# - Otherwise, find the latest checkpoint and submit a new 8-GPU resume job. +# Chaining: the new resume job gets its own watchdog submitted with afterany:. + +set -eu +if [ -z "${MIGRATOR_ROOT:-}" ]; then HERE="$(cd "$(dirname "$0")" && pwd)"; MIGRATOR_ROOT="$(cd "$HERE/../../.." && pwd)"; fi +export MIGRATOR_ROOT +REPO="$MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2" +PY="${MIGRATOR_PY:-python}" +export MIGRATOR_DATA_ROOT="${MIGRATOR_DATA_ROOT:-$MIGRATOR_ROOT/data}" + +REPO="$MIGRATOR_ROOT/code/mv-sam3d-for-6d-v2" +CKPT_DIR="$REPO/cache/ckpt" +DONE_CKPT="$CKPT_DIR/step_0100000.pt" + +echo "[watchdog] host=$(hostname) date=$(date)" +echo "[watchdog] checking for done checkpoint: $DONE_CKPT" + +if [ -f "$DONE_CKPT" ]; then + echo "[watchdog] step_0100000.pt EXISTS — training is complete. Nothing to do." + exit 0 +fi + +# Find the latest checkpoint +LATEST=$(ls -1 "$CKPT_DIR"/step_*.pt 2>/dev/null | sort | tail -1) +if [ -z "$LATEST" ]; then + echo "[watchdog] ERROR: no checkpoints found in $CKPT_DIR — cannot resume." + exit 1 +fi + +echo "[watchdog] latest checkpoint: $LATEST" +STEP=$(basename "$LATEST" .pt | sed 's/step_0*//') +echo "[watchdog] resuming from step $STEP" + +# Submit a new training job with 8x B200, 48h time limit +NEW_JOB=$(sbatch \ + --gres=gpu:b200:8 \ + --time=48:00:00 \ + --parsable \ + "$REPO/sbatch/train.sh" \ + --config "$REPO/configs/train.yaml" \ + --set train.batch_size=12 \ + --set train.steps=100000 \ + --set "train.resume_from=$LATEST") + +echo "[watchdog] submitted resume job: $NEW_JOB" + +# Submit another watchdog after the new resume job ends +WD_JOB=$(sbatch \ + --dependency=afterany:"$NEW_JOB" \ + --parsable \ + "$REPO/sbatch/watchdog_resume.sh") + +echo "[watchdog] submitted successor watchdog: $WD_JOB (afterany:$NEW_JOB)" +echo "[watchdog] chain: $NEW_JOB -> watchdog $WD_JOB" diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/_badfile_gate.py b/migrator/code/mv-sam3d-for-6d-v2/tools/_badfile_gate.py new file mode 100644 index 0000000000000000000000000000000000000000..335cd6a8fcd4e4b85cc1829c1036c1c383da9e65 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/_badfile_gate.py @@ -0,0 +1,27 @@ +import os +import sys, os, yaml, torch +REPO=os.path.abspath(os.path.join(os.path.dirname(__file__), "..")); sys.path.insert(0,REPO) +from mvsam3d.data.integrated import build_integrated_dataset, make_weighted_sampler +from mvsam3d.data.dexycb_dataset import collate_multiview +cfg=sys.argv[1]; N=int(sys.argv[2]) if len(sys.argv)>2 else 4300 +d=yaml.safe_load(open(cfg))["data"]; mode=d.get("ssi_mode","scene") +kw=dict(image_size=d["image_size"],random_views=d["random_views"],min_aux=d["min_aux_views"], + max_aux=d["max_aux_views"],num_aux=d["num_aux_views"],depth_near=d["depth_near"], + depth_far=d["depth_far"],ssi_mode=mode,heldout_path=d.get("heldout_manifest")) +ds,w=build_integrated_dataset(d,"train",**kw) +print(f"[gate] ssi={mode} total={len(ds)} iterating {N} batches (bs=1, workers=8, seed=0) past step-3440",flush=True) +sampler=make_weighted_sampler(w,world_size=1,rank=0,seed=0) +dl=torch.utils.data.DataLoader(ds,batch_size=1,sampler=sampler,num_workers=8, + collate_fn=collate_multiview,drop_last=True,persistent_workers=False) +it=iter(dl); bad=0 +for b in range(N): + try: + batch=next(it) + for k in ("shape_latent","rot6d","translation","scale"): + assert torch.isfinite(batch[k]).all(), f"nonfinite {k} @ {b}" + except Exception as e: + bad+=1; print(f" !! UNCAUGHT @batch {b}: {type(e).__name__}: {e}",flush=True) + import traceback; traceback.print_exc() + if bad>3: break + if (b+1)%500==0: print(f" ...{b+1}/{N} ok (past 3440 = {'YES' if b+1>3440 else 'no'})",flush=True) +print(f"[gate {mode}] {'PASS' if bad==0 else 'FAIL'} uncaught={bad}",flush=True) diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/_hard_smoke_ssi.py b/migrator/code/mv-sam3d-for-6d-v2/tools/_hard_smoke_ssi.py new file mode 100644 index 0000000000000000000000000000000000000000..ddccdb3b518e6ee03ce67a848ff8e98d6708cf1d --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/_hard_smoke_ssi.py @@ -0,0 +1,86 @@ +import os +"""Hard smoke for the integrated dataloader under a given ssi_mode. +Usage: python tools/_hard_smoke_ssi.py [n_batches] +Checks: build + per-dataset counts; per-adapter file-existence audit (random +samples, anchor+aux, must not FileNotFoundError); DataLoader pass (num_workers>0) +with finite targets and zero missing-file; object_centric mask-present. +""" +import sys, os, random, traceback +import numpy as np, torch, yaml +REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +sys.path.insert(0, REPO) +from mvsam3d.data.integrated import build_integrated_dataset, make_weighted_sampler +from mvsam3d.data.dexycb_dataset import collate_multiview + +cfg_path = sys.argv[1] +n_batches = int(sys.argv[2]) if len(sys.argv) > 2 else 300 +data = yaml.safe_load(open(cfg_path))["data"] +mode = data.get("ssi_mode", "scene") +print(f"\n===== HARD SMOKE {cfg_path} ssi_mode={mode} =====", flush=True) + +kw = dict(image_size=data["image_size"], random_views=data["random_views"], + min_aux=data["min_aux_views"], max_aux=data["max_aux_views"], + num_aux=data["num_aux_views"], depth_near=data["depth_near"], + depth_far=data["depth_far"], ssi_mode=mode, + heldout_path=data.get("heldout_manifest")) +ds, weights = build_integrated_dataset(data, "train", **kw) +print(f"[build] total train samples = {len(ds)}", flush=True) + +# ---- per-adapter file-existence audit (anchor + every aux ref) ---- +HEAVY = {"h2o", "dexh2r", "contactpose"} # load_view renders -> cap sample count +print("[audit] per-adapter existence (random samples, anchor+aux load_view):", flush=True) +audit_fail = 0 +for part in ds.datasets: + ad = part.adapter + name = getattr(ad, "name", type(ad).__name__) + n = min(120 if name in HEAVY else 500, len(part.samples)) + if n == 0: + print(f" {name}: 0 samples"); continue + idxs = random.sample(range(len(part.samples)), n) + miss = 0; checked = 0 + for i in idxs: + s = part.samples[i] + refs = [s.anchor] + list(s.aux_pool) + for r in refs: + checked += 1 + try: + ad.load_view(r) + except FileNotFoundError as e: + miss += 1 + if miss <= 2: + print(f" !! {name} MISSING: {e}") + except Exception: + pass # non-file errors are not this audit's concern + audit_fail += miss + print(f" {name}: {n} samples, {checked} view-loads, missing-file={miss}", flush=True) +print(f"[audit] total missing-file across adapters = {audit_fail}", flush=True) + +# ---- DataLoader pass ---- +sampler = make_weighted_sampler(weights, world_size=1, rank=0, seed=0) +dl = torch.utils.data.DataLoader(ds, batch_size=2, sampler=sampler, num_workers=4, + collate_fn=collate_multiview, drop_last=True, + persistent_workers=True, pin_memory=False) +print(f"[loader] iterating {n_batches} batches (bs=2, workers=4)...", flush=True) +bad = 0; nonfinite = 0; seen = 0 +try: + it = iter(dl) + for b in range(n_batches): + batch = next(it) + seen += 1 + for k in ("shape_latent", "rot6d", "translation", "scale"): + if not torch.isfinite(batch[k]).all(): + nonfinite += 1 + print(f" !! non-finite {k} at batch {b}") + break + if (b + 1) % 100 == 0: + print(f" ...{b+1} ok (last views/sample={len(batch['views'])})", flush=True) +except FileNotFoundError as e: + bad += 1; print(f" !! FileNotFoundError during loader: {e}") +except Exception as e: + bad += 1; print(f" !! loader error: {type(e).__name__}: {e}"); traceback.print_exc() + +print(f"\n[RESULT {mode}] batches={seen}/{n_batches} loader_missing_file_or_err={bad} " + f"nonfinite_batches={nonfinite} audit_missing={audit_fail}", flush=True) +ok = (seen == n_batches and bad == 0 and nonfinite == 0 and audit_fail == 0) +print(f"[RESULT {mode}] {'PASS' if ok else 'FAIL'}", flush=True) +sys.exit(0 if ok else 1) diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/_smoke_integrated.py b/migrator/code/mv-sam3d-for-6d-v2/tools/_smoke_integrated.py new file mode 100644 index 0000000000000000000000000000000000000000..4fac0d66a920d31c668ab953edd6876ec647058c --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/_smoke_integrated.py @@ -0,0 +1,47 @@ +import os +"""Full-integrated DataLoader smoke test (no training). Mirrors train.py's build.""" +import sys, yaml, numpy as np, torch +from torch.utils.data import DataLoader +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +from mvsam3d.data.integrated import build_integrated_dataset, make_weighted_sampler, load_heldout, reserved_object_keys +from mvsam3d.data.dexycb_dataset import collate_multiview +from mvsam3d import paths as _paths + +cfg = _paths.expand(yaml.safe_load(open(os.path.join(_paths.CONFIG_DIR, "train_integrated.yaml")))) +d = cfg["data"] +kw = dict(image_size=d["image_size"], min_aux=d["min_aux_views"], max_aux=d["max_aux_views"], + num_aux=d["num_aux_views"], depth_near=d["depth_near"], depth_far=d["depth_far"], + heldout_path=d.get("heldout_manifest")) + +print("=== building integrated TRAIN split ===", flush=True) +ds, w = build_integrated_dataset(d, "train", random_views=d["random_views"], **kw) + +# zero-leak check +ho = load_heldout(d.get("heldout_manifest")) +reserved = reserved_object_keys(ho, d["canonical_cache"]) +leak = 0 +for part in ds.datasets: + for s in part.samples: + if s.obj_key in reserved: + leak += 1 +print(f"[zero-leak] reserved={len(reserved)} keys; leaked train samples={leak}", flush=True) + +# DataLoader across datasets via weighted sampler +sampler = make_weighted_sampler(w, world_size=1, rank=0, seed=0) +dl = DataLoader(ds, batch_size=2, sampler=sampler, num_workers=0, + collate_fn=collate_multiview, drop_last=True) +print("=== 10 batches ===", flush=True) +it = iter(dl); nv = [] +for i in range(10): + b = next(it) + for k in ("shape_latent", "rot6d", "translation", "scale"): + t = b[k] + assert torch.isfinite(t).all(), f"NON-FINITE {k} at batch {i}" + assert len(b["views"]) >= 2, f"batch {i} <2 views" + for v in b["views"]: + assert torch.isfinite(v["image_cropped"]).all(), f"non-finite image_cropped b{i}" + nv.append(len(b["views"])) + print(f" batch {i}: views={len(b['views'])} shape_latent={tuple(b['shape_latent'].shape)} " + f"rot6d={tuple(b['rot6d'].shape)} objs={b.get('obj_key', b.get('object_name'))}", flush=True) +print(f"OK: 10 batches, finite targets, view-counts={nv}", flush=True) +print("SMOKE_PASS", flush=True) diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/amodal_reproject.py b/migrator/code/mv-sam3d-for-6d-v2/tools/amodal_reproject.py new file mode 100644 index 0000000000000000000000000000000000000000..c13ddd09839e5b8b55679e7a8bfed2237d1c4723 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/amodal_reproject.py @@ -0,0 +1,168 @@ +"""Amodal (unoccluded) reproject check — isolates coordinate convention. + +seg==id in DexYCB is the MODAL mask (occluded by other objects), which caps IoU +even when geometry is perfect. Here we compare two AMODAL silhouettes of the +same object at the same GT pose: + + REFERENCE (independent of our math): densely sample the raw CAD surface and + project with the raw pose_y in pure OpenCV: + X_cv = R_cv @ p_cad + t_cv (no M, no D, no canonical) + OURS : project the canonical voxel through the full gt_layout chain: + R = R_dex @ M^T, t = R_dex c_cad + t_dex, s = max_extent, + X_p3d = R (s v) + t, X_cv = X_p3d @ D + +Both are rasterized identically (project points -> morphological close -> +fill holes) so the comparison is apples-to-apples. IoU(REFERENCE, OURS) should +be ~1.0 for EVERY object if the M / D / center_cad / max_extent math is correct. +A low value for some object/camera = a real convention bug there. CPU-only. +""" +from __future__ import annotations + +import argparse +import glob +import os +from collections import defaultdict + +import numpy as np +import torch +import trimesh +from scipy import ndimage + +from mvsam3d.data import canonical, dexycb_calib +from mvsam3d.data.dexycb_dataset import YCB_ID_TO_NAME + +_D = np.diag([-1.0, -1.0, 1.0]) +ROOT = os.path.expandvars("${MIGRATOR_DATA_ROOT}")+"/dexycb" +MODELS = os.path.join(ROOT, "models") +CACHE = os.path.expandvars("${MIGRATOR_CACHE}")+"/canonical" + +_MESH_PTS = {} # name -> (M,3) dense CAD surface samples +_CANON = {} # name -> (voxel centers, center_cad, max_extent) + + +def mesh_points(name, n=120000): + if name not in _MESH_PTS: + m = trimesh.load(os.path.join(MODELS, name, "textured_simple.obj"), + process=False, force="mesh") + pts, _ = trimesh.sample.sample_surface(m, n) + _MESH_PTS[name] = np.asarray(pts, np.float64) + return _MESH_PTS[name] + + +def canon(name): + if name not in _CANON: + c = torch.load(os.path.join(CACHE, f"{name}.pt"), map_location="cpu", + weights_only=False) + occ = c["occupancy"].numpy() + idx = np.argwhere(occ > 0) + v = (idx + 0.5) / occ.shape[0] - 0.5 + _CANON[name] = (v, np.asarray(c["center_cad"]), float(c["max_extent"])) + return _CANON[name] + + +def rasterize(X_cv, intr, H, W, close_rad): + """Project OpenCV points, close gaps, fill interior -> solid silhouette.""" + Z = X_cv[:, 2] + f = Z > 1e-3 + u = (intr["fx"] * X_cv[f, 0] / Z[f] + intr["ppx"]).round().astype(int) + v = (intr["fy"] * X_cv[f, 1] / Z[f] + intr["ppy"]).round().astype(int) + img = np.zeros((H, W), bool) + ok = (u >= 0) & (u < W) & (v >= 0) & (v < H) + img[v[ok], u[ok]] = True + if close_rad > 0: + img = ndimage.binary_closing(img, iterations=close_rad) + img = ndimage.binary_fill_holes(img) + return img + + +def sample_iou(seq_dir, serial, frame, obj_id, obj_index, intr): + name = YCB_ID_TO_NAME[obj_id] + lab = np.load(os.path.join(seq_dir, serial, f"labels_{frame:06d}.npz")) + py = lab["pose_y"][obj_index] + R_cv, t_cv = py[:, :3].astype(np.float64), py[:, 3].astype(np.float64) + if not np.isfinite(R_cv).all() or np.allclose(py, 0): + return None + H, W = lab["seg"].shape + seg = (lab["seg"] == obj_id) + + # REFERENCE amodal: raw mesh, raw pose_y, pure OpenCV + p = mesh_points(name) + ref_cv = (R_cv @ p.T).T + t_cv + medZ = float(np.median(ref_cv[ref_cv[:, 2] > 1e-3, 2])) + ref = rasterize(ref_cv, intr, H, W, close_rad=2) + + # OURS amodal: canonical voxel through full gt_layout chain + v, center_cad, max_extent = canon(name) + R, t, s = canonical.gt_layout_anchor_frame(_D @ R_cv, _D @ t_cv, center_cad, max_extent) + ours_p3d = (R @ (s * v).T).T + t + ours_cv = ours_p3d @ _D + foot = intr["fx"] * (s / 64.0) / max(medZ, 1e-3) + ours = rasterize(ours_cv, intr, H, W, close_rad=int(max(2, round(foot)))) + + inter = (ref & ours).sum(); union = (ref | ours).sum() + amodal_iou = inter / union if union > 0 else 0.0 + # how much each amodal mask covers the modal seg (sanity: both >= seg) + seg_px = int(seg.sum()) + rec_ref = (ref & seg).sum() / seg_px if seg_px else 0.0 + rec_ours = (ours & seg).sum() / seg_px if seg_px else 0.0 + return dict(amodal_iou=float(amodal_iou), seg_px=seg_px, + ref_px=int(ref.sum()), ours_px=int(ours.sum()), + rec_ref=float(rec_ref), rec_ours=float(rec_ours)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n", type=int, default=150) + ap.add_argument("--min_seg", type=int, default=400) + ap.add_argument("--thr", type=float, default=0.85) + args = ap.parse_args() + rng = np.random.RandomState(0) + seqs = sorted(glob.glob(os.path.join(ROOT, "2020*-subject-*", "2020*"))) + + rows = [] + tried = 0 + while len(rows) < args.n and tried < args.n * 40: + tried += 1 + seq = seqs[rng.randint(len(seqs))] + meta = dexycb_calib._load_yaml(os.path.join(seq, "meta.yml")) + serials = [str(x) for x in meta["serials"]] + ycb_ids = list(meta["ycb_ids"]) + frame = rng.randint(int(meta["num_frames"])) + serial = serials[rng.randint(len(serials))] + oi = rng.randint(len(ycb_ids)) + oid = int(ycb_ids[oi]) + if oid not in YCB_ID_TO_NAME: + continue + intr = dexycb_calib.load_intrinsics(os.path.join(ROOT, "calibration"), serial) + r = sample_iou(seq, serial, frame, oid, oi, intr) + if r is None or r["seg_px"] < args.min_seg: + continue + r.update(obj=YCB_ID_TO_NAME[oid], serial=serial, + seq=os.path.basename(seq), frame=frame) + rows.append(r) + + a = np.array([r["amodal_iou"] for r in rows]) + print(f"\n[amodal] N={len(rows)}") + print(f"[amodal] AMODAL IoU (ours-vs-raw-mesh): mean={a.mean():.3f} " + f"median={np.median(a):.3f} min={a.min():.3f} >0.85: {(a>0.85).mean()*100:.0f}% " + f">0.90: {(a>0.90).mean()*100:.0f}%") + + by = defaultdict(list) + for r in rows: + by[r["obj"]].append(r) + print(f"\n{'object':<26}{'n':>3} {'amodalIoU':>10} {'min':>7}") + print("-" * 48) + for obj in sorted(by): + rs = [x["amodal_iou"] for x in by[obj]] + print(f"{obj:<26}{len(rs):>3} {np.mean(rs):>10.3f} {np.min(rs):>7.3f}") + + bad = sorted([r for r in rows if r["amodal_iou"] < args.thr], + key=lambda r: r["amodal_iou"]) + print(f"\n[amodal] suspects (amodal IoU < {args.thr}): {len(bad)}") + for r in bad[:40]: + print(f" iou={r['amodal_iou']:.3f} {r['obj']:<22} cam {r['serial']} " + f"{r['seq']} f{r['frame']:03d} ref_px={r['ref_px']} ours_px={r['ours_px']}") + + +if __name__ == "__main__": + main() diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/build_canonical_latents.py b/migrator/code/mv-sam3d-for-6d-v2/tools/build_canonical_latents.py new file mode 100644 index 0000000000000000000000000000000000000000..cb42c07793ca6b30cf704488d6df85c421eca9f9 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/build_canonical_latents.py @@ -0,0 +1,130 @@ +"""Build the canonical SS-latent cache for the 21 DexYCB CAD objects. + +For each object: + CAD .obj -> canonical occupancy [1,64,64,64] (vendored voxelize_mesh recipe) + -> SS latent [8,16,16,16] (REAL SAM-3D ss_encoder, posterior mean) + also records max_extent (= GT isotropic scale s) and CAD-frame AABB center + (used to build GT translation per frame). + +Round-trip encode->decode IoU is reported through the REAL standalone SAM-3D +ss_decoder (the authoritative SAM-3D SS latent space; expect ~1.0). For +contrast, the IoU through the generator-EMBEDDED decoder (TRELLIS-identical) is +also reported — it should be ~0, proving the latents live in SAM-3D's space and +NOT TRELLIS's. + +Outputs (cache_dir): + .pt : {latent[8,16,16,16], occupancy[1,64,64,64] uint8, + max_extent float, center_cad (3,)} + manifest.json : per-object {max_extent, center_cad, latent_path, n_occupied, + roundtrip_iou, roundtrip_iou_embedded} + +GPU is used only for the encoder/decoder -> run via sbatch. +""" +from __future__ import annotations + +import argparse +import json +import os +import sys + +import numpy as np +import torch +import trimesh + +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, os.path.dirname(_HERE)) # mv-sam3d-for-6d-v2/ + +from mvsam3d.data.canonical import voxelize_canonical # noqa: E402 +from mvsam3d.model.ss_vae import ( # noqa: E402 + decode_latent, encode_occupancy, load_ss_decoder, load_ss_encoder, voxel_iou, +) +from sam3d_objects.model.backbone.tdfy_dit.models.sparse_structure_vae import ( # noqa: E402 + SparseStructureDecoder, +) +from mvsam3d import paths as _paths # noqa: E402 + +SS_GENERATOR_CKPT = _paths.SS_GENERATOR_CKPT + + +def load_embedded_decoder(device: str) -> SparseStructureDecoder: + """Decoder embedded inside ss_generator.ckpt (_base_models.decoder.*). + + This is TRELLIS-identical and pairs with the TRELLIS encoder, NOT the SAM-3D + encoder; decoding SAM-3D latents through it gives IoU~0. Loaded ONLY to + demonstrate the latent-space distinction (never used to build/validate GT).""" + dec = SparseStructureDecoder(out_channels=1, latent_channels=8, channels=[512, 128, 32], + num_res_blocks=2, num_res_blocks_middle=2, + reshape_input_to_cube=False) + ckpt = torch.load(SS_GENERATOR_CKPT, map_location="cpu", weights_only=False) + sd = ckpt.get("state_dict", ckpt) if isinstance(ckpt, dict) else ckpt + prefix = "_base_models.decoder." + sub = {k[len(prefix):]: v for k, v in sd.items() + if k.startswith(prefix) and hasattr(v, "shape")} + missing, unexpected = dec.load_state_dict(sub, strict=False) + print(f"[build] embedded decoder: loaded {len(sub)} tensors " + f"({len(missing)} missing, {len(unexpected)} unexpected)") + return dec.to(device).eval() + + +def load_mesh(obj_path: str) -> trimesh.Trimesh: + m = trimesh.load(obj_path, process=False, force="mesh") + if isinstance(m, trimesh.Scene): + m = trimesh.util.concatenate([g for g in m.geometry.values()]) + return m + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--models_dir", default=os.path.join(_paths.DATA_ROOT, "dexycb", "models")) + ap.add_argument("--cache_dir", default=os.path.join(os.path.dirname(_HERE), "cache", "canonical")) + ap.add_argument("--mesh_name", default="textured_simple.obj") + args = ap.parse_args() + + os.makedirs(args.cache_dir, exist_ok=True) + device = "cuda" + # fp32 everywhere -> precise, deterministic GT latents (posterior mean). + enc = load_ss_encoder(device) # REAL SAM-3D encoder + dec = load_ss_decoder(device) # REAL standalone SAM-3D decoder + dec_emb = load_embedded_decoder(device) # gen-embedded (== TRELLIS) decoder + + objects = sorted(d for d in os.listdir(args.models_dir) + if os.path.isdir(os.path.join(args.models_dir, d))) + manifest = {} + print(f"[build] {len(objects)} objects -> {args.cache_dir}", flush=True) + for name in objects: + obj_path = os.path.join(args.models_dir, name, args.mesh_name) + mesh = load_mesh(obj_path) + occ, max_extent, center_cad = voxelize_canonical(mesh, resolution=64) + occ_f = occ.float().unsqueeze(0) # [1,1,64,64,64] + latent = encode_occupancy(enc, occ_f) # [1,8,16,16,16] + iou = voxel_iou(decode_latent(dec, latent), occ_f) # SAM-3D standalone decoder + iou_emb = voxel_iou(decode_latent(dec_emb, latent), occ_f) # gen-embedded (TRELLIS) decoder + + out = { + "latent": latent.squeeze(0).cpu(), + "occupancy": occ.squeeze(0).to(torch.uint8).cpu(), + "max_extent": float(max_extent), + "center_cad": np.asarray(center_cad, np.float64), + } + path = os.path.join(args.cache_dir, f"{name}.pt") + torch.save(out, path) + manifest[name] = { + "max_extent": float(max_extent), + "center_cad": [float(c) for c in center_cad], + "latent_path": path, + "n_occupied": int(occ.sum().item()), + "roundtrip_iou": float(iou), # SAM-3D standalone decoder + "roundtrip_iou_embedded": float(iou_emb), # gen-embedded (TRELLIS) decoder + } + print(f" {name:28s} ext={max_extent:.4f}m occ={int(occ.sum()):6d} " + f"IoU[SAM3D]={iou:.4f} IoU[embedded]={iou_emb:.4f}", flush=True) + + with open(os.path.join(args.cache_dir, "manifest.json"), "w") as f: + json.dump(manifest, f, indent=2) + ious = [m["roundtrip_iou"] for m in manifest.values()] + print(f"[build] done. SAM-3D-decoder roundtrip IoU mean={np.mean(ious):.4f} " + f"min={np.min(ious):.4f} (low IoU => encoder/decoder mismatch)", flush=True) + + +if __name__ == "__main__": + main() diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/build_canonical_latents_multi.py b/migrator/code/mv-sam3d-for-6d-v2/tools/build_canonical_latents_multi.py new file mode 100644 index 0000000000000000000000000000000000000000..dd17965f3c4c711e5d98a4d5ae64ae66321caea9 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/build_canonical_latents_multi.py @@ -0,0 +1,56 @@ +"""Generalized canonical SS-latent precompute for ANY adapter's objects. + +Enumerates each adapter's unique obj_keys, voxelizes the mesh (CPU) and encodes +to the SS latent (GPU, REAL SAM-3D encoder), writing cache/canonical/.pt +(namespaced keys create subdirs). Dedups by obj_key so YCB objects shared across +datasets are built once. GPU-only -> run via sbatch. Requires free disk. + + python tools/build_canonical_latents_multi.py --dataset pace --root .../pace + python tools/build_canonical_latents_multi.py --dataset housecat6d --root .../housecat6d/extracted +""" +import argparse, os, sys +import numpy as np, torch +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +from mvsam3d import paths as _paths +from mvsam3d.data.canonical import voxelize_canonical +from mvsam3d.data.adapters import ADAPTERS +from mvsam3d.model.ss_vae import encode_occupancy, load_ss_encoder + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--dataset", required=True) + ap.add_argument("--root", required=True) + ap.add_argument("--cache", default=_paths.CANONICAL_CACHE) + ap.add_argument("--adapter_kwargs", default="{}") + args = ap.parse_args() + import json + ad = ADAPTERS[args.dataset](args.root, args.cache, **json.loads(args.adapter_kwargs)) + # one representative sample per obj_key (mesh_for needs a sample) + rep = {} + for s in ad.list_samples(): + rep.setdefault(s.obj_key, s) + print(f"[latents] {args.dataset}: {len(rep)} unique objects", flush=True) + enc = load_ss_encoder("cuda") + done = skip = 0 + for key, s in rep.items(): + out_path = os.path.join(args.cache, f"{key}.pt") + if os.path.exists(out_path): + skip += 1; continue + os.makedirs(os.path.dirname(out_path), exist_ok=True) + try: + mesh = ad.mesh_for(s) + occ, ext, cen = voxelize_canonical(mesh, resolution=64) + occ_f = occ.float().unsqueeze(0) if occ.ndim == 4 else occ.float().unsqueeze(0).unsqueeze(0) + latent = encode_occupancy(enc, occ_f) + torch.save({"latent": latent.squeeze(0).cpu(), + "occupancy": (occ.squeeze(0) if occ.ndim == 4 else occ).to(torch.uint8).cpu(), + "max_extent": float(ext), + "center_cad": np.asarray(cen, np.float64)}, out_path) + done += 1 + if done % 20 == 0: print(f" {done} built...", flush=True) + except Exception as e: + print(f" [skip] {key}: {e}", flush=True) + print(f"[latents] {args.dataset}: built {done}, skipped {skip} (existing)", flush=True) + +if __name__ == "__main__": + main() diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/diagnose_reproject.py b/migrator/code/mv-sam3d-for-6d-v2/tools/diagnose_reproject.py new file mode 100644 index 0000000000000000000000000000000000000000..70c2a076617c8035be1d247eb3b20b7344b6a2c6 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/diagnose_reproject.py @@ -0,0 +1,167 @@ +"""Diagnose GT-frame reproject quality and hunt for wrong-convention samples. + +Improves on verify_gt_reproject.py in two ways: + 1. Fills the projected silhouette with an ADAPTIVE dilation sized to the + projected voxel footprint (fx * s/64 / Z), instead of a fixed 1px. This + removes the sparse-surface-voxel artifact that otherwise depresses IoU. + 2. Reports per-sample IoU *and* RECALL = |proj & seg| / |seg|. Recall is the + convention-sensitive metric: correct geometry -> recall ~ 1.0 for every + sample regardless of occlusion (occlusion only shrinks seg, lowering + precision/IoU, never recall). Low recall => real geometry/convention bug. + +Breaks results down by object class and by camera serial so any clustering of +bad samples (a specific object, camera, or sequence) is visible. CPU-only. +""" +from __future__ import annotations + +import argparse +import glob +import os +from collections import defaultdict + +import numpy as np +import torch + +from mvsam3d.data import canonical, dexycb_calib +from mvsam3d.data.dexycb_dataset import YCB_ID_TO_NAME + +_D = np.diag([-1.0, -1.0, 1.0]) +ROOT = os.path.expandvars("${MIGRATOR_DATA_ROOT}")+"/dexycb" +CACHE = os.path.expandvars("${MIGRATOR_CACHE}")+"/canonical" + +_CANON = {} + + +def load_canon(name): + if name not in _CANON: + c = torch.load(os.path.join(CACHE, f"{name}.pt"), map_location="cpu", + weights_only=False) + occ = c["occupancy"].numpy() + idx = np.argwhere(occ > 0) + v = (idx + 0.5) / occ.shape[0] - 0.5 + _CANON[name] = (v, np.asarray(c["center_cad"]), float(c["max_extent"])) + return _CANON[name] + + +def reproject(seq_dir, serial, frame, obj_id, obj_index, intr): + v, center_cad, max_extent = load_canon(YCB_ID_TO_NAME[obj_id]) + + lab = np.load(os.path.join(seq_dir, serial, f"labels_{frame:06d}.npz")) + py = lab["pose_y"][obj_index] + R_cv, t_cv = py[:, :3].astype(np.float64), py[:, 3].astype(np.float64) + if not np.isfinite(R_cv).all() or np.allclose(py, 0): + return None + R_p3d, t_p3d = _D @ R_cv, _D @ t_cv + R, t, s = canonical.gt_layout_anchor_frame(R_p3d, t_p3d, center_cad, max_extent) + + X_p3d = (R @ (s * v).T).T + t + X_cv = X_p3d @ _D + Z = X_cv[:, 2] + front = Z > 1e-3 + Zf = Z[front] + u = intr["fx"] * X_cv[front, 0] / Zf + intr["ppx"] + vv = intr["fy"] * X_cv[front, 1] / Zf + intr["ppy"] + ui = u.round().astype(int) + vi = vv.round().astype(int) + + seg = (lab["seg"] == obj_id) + H, W = seg.shape + proj = np.zeros((H, W), bool) + ok = (ui >= 0) & (ui < W) & (vi >= 0) & (vi < H) + proj[vi[ok], ui[ok]] = True + + # adaptive dilation: a canonical voxel has world size s/64; its projected + # footprint at the median depth is fx*(s/64)/Z px. Dilating by ~half that + # closes the gaps between adjacent surface voxels -> a filled silhouette. + medZ = float(np.median(Zf)) if Zf.size else 1.0 + foot = intr["fx"] * (s / 64.0) / max(medZ, 1e-3) + rad = int(max(1, round(foot / 2.0))) + for _ in range(rad): + proj[1:] |= proj[:-1].copy(); proj[:-1] |= proj[1:].copy() + proj[:, 1:] |= proj[:, :-1].copy(); proj[:, :-1] |= proj[:, 1:].copy() + + inter = (proj & seg).sum() + union = (proj | seg).sum() + seg_px = int(seg.sum()) + iou = inter / union if union > 0 else 0.0 + recall = inter / seg_px if seg_px > 0 else 0.0 + return iou, recall, seg_px, int(proj.sum()), rad + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n", type=int, default=300) + ap.add_argument("--min_seg", type=int, default=400) + ap.add_argument("--recall_thr", type=float, default=0.85) + args = ap.parse_args() + rng = np.random.RandomState(0) + seqs = sorted(glob.glob(os.path.join(ROOT, "2020*-subject-*", "2020*"))) + + rows = [] + tried = 0 + while len(rows) < args.n and tried < args.n * 30: + tried += 1 + seq = seqs[rng.randint(len(seqs))] + meta = dexycb_calib._load_yaml(os.path.join(seq, "meta.yml")) + serials = [str(x) for x in meta["serials"]] + ycb_ids = list(meta["ycb_ids"]) + frame = rng.randint(int(meta["num_frames"])) + serial = serials[rng.randint(len(serials))] + oi = rng.randint(len(ycb_ids)) + oid = int(ycb_ids[oi]) + if oid not in YCB_ID_TO_NAME: + continue + intr = dexycb_calib.load_intrinsics(os.path.join(ROOT, "calibration"), serial) + r = reproject(seq, serial, frame, oid, oi, intr) + if r is None: + continue + iou, recall, seg_px, proj_px, rad = r + if seg_px < args.min_seg: + continue + rows.append(dict(obj=YCB_ID_TO_NAME[oid], serial=serial, seq=os.path.basename(seq), + frame=frame, iou=iou, recall=recall, seg_px=seg_px, + proj_px=proj_px, rad=rad)) + + iou = np.array([r["iou"] for r in rows]) + rec = np.array([r["recall"] for r in rows]) + print(f"\n[diag] N={len(rows)}") + print(f"[diag] IoU : mean={iou.mean():.3f} median={np.median(iou):.3f} " + f"min={iou.min():.3f} >0.85: {(iou>0.85).mean()*100:.0f}%") + print(f"[diag] RECALL: mean={rec.mean():.3f} median={np.median(rec):.3f} " + f"min={rec.min():.3f} >0.85: {(rec>0.85).mean()*100:.0f}%") + + # per-object breakdown + by_obj = defaultdict(list) + for r in rows: + by_obj[r["obj"]].append(r) + print(f"\n{'object':<26}{'n':>3} {'IoU':>6} {'recall':>7} {'recall_min':>11}") + print("-" * 56) + for obj in sorted(by_obj): + rs = by_obj[obj] + i = np.mean([x["iou"] for x in rs]); rc = np.mean([x["recall"] for x in rs]) + rcm = np.min([x["recall"] for x in rs]) + print(f"{obj:<26}{len(rs):>3} {i:>6.3f} {rc:>7.3f} {rcm:>11.3f}") + + # per-camera breakdown (cluster => extrinsics/convention issue per camera) + by_cam = defaultdict(list) + for r in rows: + by_cam[r["serial"]].append(r) + print(f"\n{'camera serial':<18}{'n':>3} {'IoU':>6} {'recall':>7}") + print("-" * 36) + for cam in sorted(by_cam): + rs = by_cam[cam] + print(f"{cam:<18}{len(rs):>3} {np.mean([x['iou'] for x in rs]):>6.3f} " + f"{np.mean([x['recall'] for x in rs]):>7.3f}") + + # the convention-error suspects: low recall + bad = sorted([r for r in rows if r["recall"] < args.recall_thr], + key=lambda r: r["recall"]) + print(f"\n[diag] LOW-RECALL suspects (recall < {args.recall_thr}): {len(bad)}") + for r in bad[:40]: + print(f" recall={r['recall']:.3f} iou={r['iou']:.3f} {r['obj']:<24} " + f"cam {r['serial']} {r['seq']} f{r['frame']:03d} " + f"seg_px={r['seg_px']:6d} proj_px={r['proj_px']:6d} rad={r['rad']}") + + +if __name__ == "__main__": + main() diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/infer_demo.py b/migrator/code/mv-sam3d-for-6d-v2/tools/infer_demo.py new file mode 100644 index 0000000000000000000000000000000000000000..2cc0c18aa024a9fa9315f4285ba94a88426648ef --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/infer_demo.py @@ -0,0 +1,98 @@ +"""Minimal inference demo for the multi-view 6D-pose flow model (v2). + +Style mirrors sam3d-for-6d/demo.py: load a config + ONE checkpoint, build and +restore the model, take ONE input sample, run the sampler, print the predicted +pose. No metrics, no GT comparison — just "load model, run, print pose". + +The model build + checkpoint restore + sampler call are delegated to the +reusable ``mvsam3d.infer.predict`` module, so this stays a thin, readable demo +that exercises the SAME inference path as ``validate.run_validation``: + + build MVSam3DFlow(pretrained_ckpt=...) -> vendored load_checkpoint(step_*.pt) + -> eval() + AMP autocast + no_grad -> inner.predict(batch, num_steps). + +Run (GPU; sbatch only — see sbatch/infer_smoke.sh): + python -m tools.infer_demo \ + --checkpoint /.../cache/ckpt/step_0014000.pt \ + --idx 0 --num-steps 50 --subjects subject-01 --frame-stride 100 +""" +from __future__ import annotations + +import argparse + +import torch + +from mvsam3d.infer.predict import ( + DEFAULT_CONFIG, + build_inference_model, + load_config, + predict_pose, + _build_dexycb, +) +from mvsam3d.data.dexycb_dataset import collate_multiview + + +def main(): + ap = argparse.ArgumentParser(description="Minimal v2 6D-pose inference demo.") + ap.add_argument("--config", default=DEFAULT_CONFIG, help="training YAML config") + ap.add_argument("--checkpoint", required=True, help="v2 checkpoint (cache/ckpt/step_*.pt)") + ap.add_argument("--idx", type=int, default=0, help="sample index to predict (default 0)") + ap.add_argument("--num-steps", type=int, default=None, + help="Euler ODE steps (default: config train.val_num_steps)") + ap.add_argument("--subjects", type=str, nargs="*", default=None, + help="subject substrings (e.g. subject-01); default = config") + ap.add_argument("--frame-stride", type=int, default=None, help="override data.frame_stride") + ap.add_argument("--num-aux-views", type=int, default=None, help="override aux-view count") + ap.add_argument("--seed", type=int, default=1234, help="RNG seed for the sampler noise z0") + args = ap.parse_args() + + cfg = load_config(args.config) + data_cfg = cfg["data"] + num_steps = args.num_steps if args.num_steps is not None else cfg["train"]["val_num_steps"] + amp_dtype = cfg["train"]["amp_dtype"] + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Deterministic sampler noise (validate.py seeds before sampling). + torch.manual_seed(args.seed) + if device.type == "cuda": + torch.cuda.manual_seed_all(args.seed) + + # ---- build + restore (warm-start FIRST, then vendored load_checkpoint) ---- + model, report = build_inference_model(cfg, args.checkpoint, device, verbose=True) + + # ---- one input sample (simplest: pull sample --idx from MultiViewDexYCB) ---- + subjects = args.subjects if args.subjects else data_cfg["subjects"] + frame_stride = args.frame_stride if args.frame_stride is not None else data_cfg["frame_stride"] + num_aux = (args.num_aux_views if args.num_aux_views is not None + else (data_cfg["max_aux_views"] if data_cfg["random_views"] + else data_cfg["num_aux_views"])) + ds = _build_dexycb(data_cfg, subjects=subjects, frame_stride=frame_stride, + num_aux_views=num_aux, split="val") + print(f"[demo] dataset: {len(ds)} samples; using idx={args.idx}") + sample = ds[args.idx] + batch = collate_multiview([sample]) # B=1 + print(f"[demo] object={batch['object_name'][0]} n_views={len(batch['views'])}") + + # ---- run the sampler (same path as validation) ---- + pred = predict_pose(model, batch, num_steps=num_steps, amp_dtype=amp_dtype, device=device) + + R = pred["rotation"][0] # [3,3] + t = pred["translation"][0] # [3] metric + s = float(pred["scale"][0]) # metric scalar + det = float(torch.det(R)) + ortho_err = float((R.T @ R - torch.eye(3, device=R.device)).abs().max()) + + print("\n========== PREDICTED POSE ==========") + print(f"checkpoint step : {report['step']} (num_steps={num_steps})") + print(f"rotation [3,3]:\n{R.cpu().numpy()}") + print(f" det(R) : {det:+.6f} (orthonormal => +1)") + print(f" ||RᵀR - I||_max : {ortho_err:.3e} (orthonormal => ~0)") + print(f" finite : {bool(torch.isfinite(R).all())}") + print(f"translation [3] (metric) : {t.cpu().numpy()} finite={bool(torch.isfinite(t).all())}") + print(f"scale (metric) : {s:.6f} finite={bool(torch.isfinite(pred['scale'][0]))}") + print(f"shape_latent shape : {tuple(pred['shape_latent'].shape)}") + print("====================================") + + +if __name__ == "__main__": + main() diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/scan_corrupt_pngs.py b/migrator/code/mv-sam3d-for-6d-v2/tools/scan_corrupt_pngs.py new file mode 100644 index 0000000000000000000000000000000000000000..71d892ae0c9160933f0912b85e4046d1cc22d6df --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/scan_corrupt_pngs.py @@ -0,0 +1,61 @@ +"""[custom] Scan a dataset tree for corrupt PNGs (full decode via PIL, matching +exactly what the DataLoader does). Writes one offending path per line. + +Usage: + python tools/scan_corrupt_pngs.py --root --pattern "*/depth/*.png" \ + --out corrupt_pngs.txt --workers 32 +CPU-only. +""" +from __future__ import annotations + +import argparse +import fnmatch +import os +from concurrent.futures import ProcessPoolExecutor + +import numpy as np +from PIL import Image + + +def check(path: str): + try: + # np.asarray forces a full decode — same failure mode as the train loader + # (Image.verify() alone misses truncated IDAT payloads). + np.asarray(Image.open(path)) + return None + except Exception as e: + return f"{path}\t{type(e).__name__}: {e}" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--root", required=True) + ap.add_argument("--pattern", default="*/depth/*.png") + ap.add_argument("--out", required=True) + ap.add_argument("--workers", type=int, default=32) + args = ap.parse_args() + + paths = [] + for dirpath, _, files in os.walk(args.root): + for f in files: + p = os.path.join(dirpath, f) + if fnmatch.fnmatch(p, args.pattern): + paths.append(p) + print(f"[scan] {len(paths)} files matching {args.pattern!r} under {args.root}", flush=True) + + bad = [] + with ProcessPoolExecutor(max_workers=args.workers) as ex: + for i, res in enumerate(ex.map(check, paths, chunksize=256)): + if res: + bad.append(res) + print(f"[scan][BAD] {res}", flush=True) + if (i + 1) % 50000 == 0: + print(f"[scan] {i + 1}/{len(paths)} checked, {len(bad)} bad", flush=True) + + with open(args.out, "w") as f: + f.write("\n".join(bad) + ("\n" if bad else "")) + print(f"[scan] DONE: {len(bad)} corrupt / {len(paths)} total -> {args.out}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/validate_model.py b/migrator/code/mv-sam3d-for-6d-v2/tools/validate_model.py new file mode 100644 index 0000000000000000000000000000000000000000..ac0a8505ee0f127b68783778564fbdcc93d3ad2d --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/validate_model.py @@ -0,0 +1,120 @@ +"""Model + pretrained-load validation (GPU). + +Asserts (task BUILD step 4): + * the condition embedder IS the real EmbedderFuser class (vendored fork class); + * the condition embedder loaded 722 tensors from ss_generator.ckpt; + * the DiT torso loaded for shape / 6drotation_normalized / translation streams; + * ONLY mlp_cam is from-scratch (no pretrained source); + * one flow_step fwd+bwd on a REAL multi-view DexYCB batch (NaN-carrying + pointmaps) yields a finite loss; + * frozen modules (EmbedderFuser) receive 0 grad; mlp_cam receives grad. +""" +from __future__ import annotations + +import argparse + +import torch + +from sam3d_objects.model.backbone.dit.embedder.embedder_fuser import EmbedderFuser +from mvsam3d.data.dexycb_dataset import MultiViewDexYCB, collate_multiview +from mvsam3d.model.flow_model import MVSam3DFlow +from mvsam3d import paths as _paths + +GEN_CKPT = _paths.SS_GENERATOR_CKPT + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--cache", default=_paths.CANONICAL_CACHE) + ap.add_argument("--subjects", default="subject-01") + args = ap.parse_args() + device = "cuda" + torch.manual_seed(0) + + # Real multi-view batch (random 1..3 aux -> 2..4 views), NaN-carrying pointmaps. + ds = MultiViewDexYCB( + canonical_cache=args.cache, subjects=[args.subjects], + frame_stride=20, random_views=True, min_aux_views=1, max_aux_views=3, + ) + print(f"[validate] dataset samples: {len(ds)}") + batch = collate_multiview([ds[0], ds[1]]) # B=2 + n_views = len(batch["views"]) + # Confirm NaN actually present in a pointmap (the invalid_xyz_token pathway). + has_nan = any(torch.isnan(batch["views"][v]["pointmap_cropped"]).any().item() + for v in range(n_views)) + print(f"[validate] batch views={n_views} pointmap carries NaN: {has_nan}") + assert has_nan, "expected NaN in pointmaps (invalid/background pixels)" + + model = MVSam3DFlow(max_views=4, pretrained_ckpt=GEN_CKPT).to(device) + + # (1) EmbedderFuser is the REAL vendored class. + assert isinstance(model.condition.fuser, EmbedderFuser), type(model.condition.fuser) + print(f"[validate] condition.fuser is EmbedderFuser: OK ({type(model.condition.fuser).__module__})") + + # (2) condition embedder loaded 722 tensors. + ckpt = torch.load(GEN_CKPT, map_location="cpu", weights_only=False) + sd = ckpt.get("state_dict", ckpt) + cond_prefix = "_base_models.condition_embedder." + own = model.condition.fuser.state_dict() + loaded = sum(1 for nk in own if (cond_prefix + nk) in sd + and own[nk].shape == sd[cond_prefix + nk].shape) + n_src = sum(1 for k in sd if k.startswith(cond_prefix) and hasattr(sd[k], "shape")) + print(f"[validate] condition embedder: {loaded}/{len(own)} fuser tensors loaded; " + f"{n_src} source tensors under {cond_prefix}") + assert n_src == 722, f"expected 722 source tensors, got {n_src}" + assert loaded == len(own), f"fuser not fully loaded: {loaded}/{len(own)}" + + # (3) torso loaded per pose/shape stream. + own_bb = model.backbone.state_dict() + torso = "_base_models.generator.reverse_fn.backbone." + for stream in ("shape", "6drotation_normalized", "translation"): + got = sum(1 for nk in own_bb if nk.startswith(f"latent_mapping.{stream}.") + and (torso + nk) in sd and own_bb[nk].shape == sd[torso + nk].shape) + tot = sum(1 for nk in own_bb if nk.startswith(f"latent_mapping.{stream}.")) + print(f"[validate] torso stream {stream}: {got}/{tot} loaded") + assert got > 0, f"stream {stream} got no pretrained weights" + + # (4) ONLY mlp_cam from-scratch. Frozen fuser params have requires_grad=False; + # mlp_cam requires grad; backbone trains. + n_cam = sum(p.numel() for p in model.condition.mlp_cam.parameters()) + cam_trainable = all(p.requires_grad for p in model.condition.mlp_cam.parameters()) + fuser_trainable = sum(p.numel() for p in model.condition.fuser.parameters() if p.requires_grad) + cond_train_modules = sorted({n.split(".")[0] for n, p in model.condition.named_parameters() + if p.requires_grad}) + print(f"[validate] mlp_cam trainable={cam_trainable} ({n_cam} params); " + f"fuser trainable params={fuser_trainable}; " + f"condition trainable modules={cond_train_modules}") + assert cam_trainable, "mlp_cam must be trainable" + assert fuser_trainable == 0, "EmbedderFuser must be frozen" + assert cond_train_modules == ["mlp_cam"], cond_train_modules + + # (5) one fwd+bwd, finite loss. + batch = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()} + batch["views"] = [{k: t.to(device) for k, t in vw.items()} for vw in batch["views"]] + model.train() + losses = model.flow_step(batch) + print("[validate] losses:", {k: round(float(v), 5) for k, v in losses.items()}) + assert torch.isfinite(losses["loss"]), "loss is not finite" + losses["loss"].backward() + + # (6) grad audit. + fuser_grad = sum((p.grad.abs().sum().item() if p.grad is not None else 0.0) + for p in model.condition.fuser.parameters()) + cam_grad = sum((p.grad.abs().sum().item() if p.grad is not None else 0.0) + for p in model.condition.mlp_cam.parameters()) + bb_grad = sum((p.grad.abs().sum().item() if p.grad is not None else 0.0) + for p in model.backbone.parameters()) + print(f"[validate] grad: fuser={fuser_grad:.3e} (expect 0) " + f"mlp_cam={cam_grad:.3e} (expect >0) backbone={bb_grad:.3e} (expect >0)") + assert fuser_grad == 0.0, "frozen fuser received gradient" + assert cam_grad > 0.0, "mlp_cam received no gradient" + assert bb_grad > 0.0, "backbone received no gradient" + + n_train = sum(p.numel() for p in model.parameters() if p.requires_grad) + n_total = sum(p.numel() for p in model.parameters()) + print(f"[validate] params: total={n_total/1e6:.1f}M trainable={n_train/1e6:.1f}M") + print("[validate] ALL CHECKS PASSED") + + +if __name__ == "__main__": + main() diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/verify_adapter_reproject.py b/migrator/code/mv-sam3d-for-6d-v2/tools/verify_adapter_reproject.py new file mode 100644 index 0000000000000000000000000000000000000000..586635cd571e9ddadcc34d039a1e77ade036af9e --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/verify_adapter_reproject.py @@ -0,0 +1,85 @@ +import os +"""Adapter-driven GT-as-pred reprojection gate (CPU, in-memory). + +For random samples of an adapter: build canonical (R,t,s) via the SAME +gt_layout_anchor_frame the dataloader uses, map the mesh's canonical occupancy +voxels into the anchor camera, project with the anchor intrinsics, and IoU vs +the anchor object mask. High IoU == the adapter's convention + the base +canonical/D math are mutually consistent (a coordinate bug tanks IoU).""" +import sys, argparse, numpy as np +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +from mvsam3d.data import canonical +from mvsam3d.data.adapters import ADAPTERS +_D = np.diag([-1.0, -1.0, 1.0]) + +def occ_centers(occ): + occ = np.asarray(occ) + if occ.ndim == 4: occ = occ[0] + idx = np.argwhere(occ > 0) + return (idx + 0.5) / occ.shape[0] - 0.5 + +def gate(name, root, cache, n=12, stride=None, seed=0): + kw = {} + if stride is not None: kw["frame_stride"] = stride + ad = ADAPTERS[name](root, cache, **kw) + samples = ad.list_samples() + rng = np.random.RandomState(seed) + order = rng.permutation(len(samples)) + ious, tried = [], 0 + voxcache = {} + for si in order: + if len(ious) >= n or tried > n * 40: break + tried += 1 + s = samples[si] + try: + v = ad.load_view(s.anchor) + R_cv, t_cv = ad.object_pose_cv(s) + if s.obj_key not in voxcache: + occ, ext, cen = canonical.voxelize_canonical(ad.mesh_for(s), resolution=64) + voxcache[s.obj_key] = (occ_centers(occ), ext, np.asarray(cen)) + vc, ext, cen = voxcache[s.obj_key] + except Exception as e: + if tried <= 2: print(f" [skip] {e}") + continue + R_p3d, t_p3d = _D @ R_cv, _D @ t_cv + R, t, s_ext = canonical.gt_layout_anchor_frame(R_p3d, t_p3d, cen, float(ext)) + Xp = (R @ (s_ext * vc).T).T + t + Xcv = Xp @ _D + Z = Xcv[:, 2]; f = Z > 1e-3 + intr = v["intr"] + u = (intr["fx"] * Xcv[f, 0] / Z[f] + intr["ppx"]).round().astype(int) + w = (intr["fy"] * Xcv[f, 1] / Z[f] + intr["ppy"]).round().astype(int) + seg = v["mask"] > 0.5; H, W = seg.shape + proj = np.zeros((H, W), bool) + ok = (u >= 0) & (u < W) & (w >= 0) & (w < H) + proj[w[ok], u[ok]] = True + for _ in range(2): + proj[1:] |= proj[:-1].copy(); proj[:-1] |= proj[1:].copy() + proj[:, 1:] |= proj[:, :-1].copy(); proj[:, :-1] |= proj[:, 1:].copy() + if seg.sum() < 150: continue + inter = (proj & seg).sum(); union = (proj | seg).sum() + ious.append(inter / union if union else 0.0) + ious = np.array(ious) if ious else np.array([np.nan]) + print(f"[{name:11s}] N={len(ious):2d} meanIoU={np.nanmean(ious):.3f} " + f"median={np.nanmedian(ious):.3f} min={np.nanmin(ious):.3f} " + f">0.5:{np.mean(ious>0.5)*100:.0f}% (samples={len(samples)})") + return np.nanmean(ious) + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("--only", default=None) + args = ap.parse_args() + CACHE = os.path.expandvars("${MIGRATOR_CACHE}")+"/canonical" + DS = os.path.expandvars("${MIGRATOR_DATA_ROOT}")+"" + jobs = [ + ("dexycb", f"{DS}/dexycb"), + ("housecat6d",f"{DS}/housecat6d/extracted"), + ("pace", f"{DS}/pace"), + ("graspnet", f"{DS}/graspnet"), + ("ycbv", f"{DS}/ycbv"), + ] + for name, root in jobs: + if args.only and name != args.only: continue + try: gate(name, root, CACHE) + except Exception as e: + import traceback; print(f"[{name}] ERROR: {e}"); traceback.print_exc() diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/verify_gt_reproject.py b/migrator/code/mv-sam3d-for-6d-v2/tools/verify_gt_reproject.py new file mode 100644 index 0000000000000000000000000000000000000000..9f98c80d6789558e1cac3a78c5bc97b7703e19d6 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/verify_gt_reproject.py @@ -0,0 +1,109 @@ +"""Verify the GT-frame math by reprojecting the canonical voxel. + +For random (sequence, frame, anchor-camera, object) samples: + 1. Build canonical (R, t, s) in the anchor camera PyTorch3D frame exactly as + the dataloader does (pose_y -> OpenCV->P3D -> gt_layout_anchor_frame). + 2. Take the cached canonical occupancy, map occupied voxel centers + v_canon in [-0.5,0.5]^3 to the camera: X_p3d = R (s v_canon) + t. + 3. Convert back to OpenCV, project with the anchor intrinsics, rasterize a + mask, and compute IoU against the GT segmentation mask. + +High IoU confirms R = R_dex M^T, t = R_dex c_cad + t_dex, s = max_extent and the +D=diag(-1,-1,1) conversions are all consistent. CPU-only. +""" +from __future__ import annotations + +import argparse +import glob +import os + +import numpy as np + +from mvsam3d.data import canonical, dexycb_calib +from mvsam3d.data.dexycb_dataset import YCB_ID_TO_NAME + +import torch + +_D = np.diag([-1.0, -1.0, 1.0]) +ROOT = os.path.expandvars("${MIGRATOR_DATA_ROOT}")+"/dexycb" +CACHE = os.path.expandvars("${MIGRATOR_CACHE}")+"/canonical" + + +def occupied_centers(occ: np.ndarray) -> np.ndarray: + """occ [64,64,64] -> (N,3) canonical centers in [-0.5,0.5].""" + idx = np.argwhere(occ > 0) # (N,3) in grid coords + return (idx + 0.5) / occ.shape[0] - 0.5 + + +def reproject_iou(seq_dir, serial, frame, obj_id, obj_index, intr, dilate=1): + canon = torch.load(os.path.join(CACHE, f"{YCB_ID_TO_NAME[obj_id]}.pt"), + map_location="cpu", weights_only=False) + occ = canon["occupancy"].numpy() + v = occupied_centers(occ) # (N,3) canonical + + lab = np.load(os.path.join(seq_dir, serial, f"labels_{frame:06d}.npz")) + py = lab["pose_y"][obj_index] + R_cv, t_cv = py[:, :3].astype(np.float64), py[:, 3].astype(np.float64) + R_p3d, t_p3d = _D @ R_cv, _D @ t_cv + R, t, s = canonical.gt_layout_anchor_frame( + R_p3d, t_p3d, np.asarray(canon["center_cad"]), float(canon["max_extent"]) + ) + + X_p3d = (R @ (s * v).T).T + t # (N,3) anchor P3D + X_cv = X_p3d @ _D # back to OpenCV + Z = X_cv[:, 2] + front = Z > 1e-3 + u = (intr["fx"] * X_cv[front, 0] / Z[front] + intr["ppx"]).round().astype(int) + vv = (intr["fy"] * X_cv[front, 1] / Z[front] + intr["ppy"]).round().astype(int) + + seg = (lab["seg"] == obj_id) + H, W = seg.shape + proj = np.zeros((H, W), bool) + ok = (u >= 0) & (u < W) & (vv >= 0) & (vv < H) + proj[vv[ok], u[ok]] = True + # small dilation so sparse voxel hits fill the silhouette + for _ in range(dilate): + proj[1:] |= proj[:-1].copy(); proj[:-1] |= proj[1:].copy() + proj[:, 1:] |= proj[:, :-1].copy(); proj[:, :-1] |= proj[:, 1:].copy() + inter = (proj & seg).sum() + union = (proj | seg).sum() + return inter / union if union > 0 else 0.0, seg.sum(), proj.sum() + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n", type=int, default=30) + args = ap.parse_args() + rng = np.random.RandomState(0) + seqs = sorted(glob.glob(os.path.join(ROOT, "2020*-subject-*", "2020*"))) + + ious = [] + tried = 0 + while len(ious) < args.n and tried < args.n * 20: + tried += 1 + seq = seqs[rng.randint(len(seqs))] + meta = dexycb_calib._load_yaml(os.path.join(seq, "meta.yml")) + serials = [str(x) for x in meta["serials"]] + ycb_ids = list(meta["ycb_ids"]) + frame = rng.randint(int(meta["num_frames"])) + serial = serials[rng.randint(len(serials))] + oi = rng.randint(len(ycb_ids)) + oid = int(ycb_ids[oi]) + if oid not in YCB_ID_TO_NAME: + continue + intr = dexycb_calib.load_intrinsics(os.path.join(ROOT, "calibration"), serial) + iou, seg_px, proj_px = reproject_iou(seq, serial, frame, oid, oi, intr) + if seg_px < 200: # object barely visible / occluded + continue + ious.append(iou) + print(f"{YCB_ID_TO_NAME[oid]:24s} cam {serial} f{frame:03d} " + f"IoU={iou:.3f} seg_px={seg_px:6d} proj_px={proj_px:6d}") + + ious = np.array(ious) + print(f"\n[verify] N={len(ious)} mean IoU={ious.mean():.3f} " + f"median={np.median(ious):.3f} min={ious.min():.3f} " + f">0.5: {(ious > 0.5).mean()*100:.0f}%") + + +if __name__ == "__main__": + main() diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/verify_reproject_eval.py b/migrator/code/mv-sam3d-for-6d-v2/tools/verify_reproject_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..53c0cd6600a0ea162a0069b4d0b00549e06397b2 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/verify_reproject_eval.py @@ -0,0 +1,326 @@ +"""RIGOROUS GT-as-pred reprojection validator for the EVAL / HELD-OUT loaders. + +Parallel to tools/verify_reproject_v2.py, but instead of the integrated-dataloader +ADAPTERS it pulls (R_oc, t_oc, K, mask, mesh, sensor-depth) DIRECTLY from the +ho3d_eval loaders that are used at EVAL time, so it validates the EXACT eval code +path (loader.object_pose/gt_pose + intrinsics + mesh + mask). + +render_depth + the three occlusion-fair metrics (amodal-IoU / coverage / +depth-agreement) are the SAME as verify_reproject_v2.py (render_depth copied +verbatim; nvdiffrast bottom-origin -> flipud to register to the image frame). + +Datasets (each validates its own eval mesh + pose convention): + * ho3d : ho3d_loader.gt_pose (glcam diag(1,-1,-1) flip on R & t); dexycb YCB + mesh (m); mask = render_gt_mask (rendered AMODAL silhouette) -> the + fair gate is amodal-IoU (full-vs-full). Corroborated by depth-agree + against the real HO3D sensor depth. + * ycbm : ycbm_loader.object_pose (quat_xyzw AS-IS, _POSE_FIX F=I); mesh = + models_aligned/ycb_models_aligned_cm (the EVAL mesh, /100 -> m), NOT + dexycb. Real per-pixel seg mask (visible) -> gate = coverage, + corroborated by depth-agree (real Kinect depth, unit/handedness + sensitive). Also tests rotation variants (as-is / F-flips / R^T). + * hope : hope_video_loader.object_pose (t cm -> /100 m); mesh = meshes/eval + (cm -> /100 m); rendered visible mask -> gate = coverage, corroborated + by depth-agree (real RealSense depth). + +Run in the any6d env (torch + nvdiffrast + cv2 + trimesh): + python tools/verify_reproject_eval.py --only ycbm +""" +from __future__ import annotations +import os, sys, argparse, json +import numpy as np + +HO3D_EVAL = os.path.expandvars("${MIGRATOR_DATA_ROOT}/ho3d_eval")+"" +sys.path.insert(0, HO3D_EVAL) + +DEPTH_TOL = 0.02 # 2 cm + + +# --------------------------------------------------------------------------- # +# render_depth: copied VERBATIM from tools/verify_reproject_v2.py (which copied +# it verbatim from ho3d_eval/ycbv_vsd_scorer.py). +# --------------------------------------------------------------------------- # +def render_depth(verts_m, faces, R, t, K, H, W, ctx, torch, dr): + fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2] + Rt = torch.from_numpy(R.astype(np.float32)).cuda() + tt = torch.from_numpy(t.reshape(3).astype(np.float32)).cuda() + vt = torch.from_numpy(verts_m).cuda() + vc = (Rt @ vt.T).T + tt + z = vc[:, 2].clamp(min=1e-6) + u = vc[:, 0] / z * fx + cx; v = vc[:, 1] / z * fy + cy + x = (u / W) * 2 - 1; y = 1 - (v / H) * 2; zb = z / (z.max() + 1e-6) + w = torch.ones(len(verts_m), device="cuda") + clip = torch.stack([x, y, zb, w], 1).unsqueeze(0) + ft = torch.from_numpy(faces).int().cuda() + rast, _ = dr.rasterize(ctx, clip, ft, resolution=[H, W]) + zi, _ = dr.interpolate(z.unsqueeze(1).unsqueeze(0), rast, ft) + d = zi[0, :, :, 0].cpu().numpy() + msk = (rast[0, :, :, 3] > 0).cpu().numpy() + d[~msk] = 0.0 + return d + + +def _mesh_load_m(path, scale): + """Load a mesh (concatenate scenes), return (verts_m float32, faces int32).""" + import trimesh + m = trimesh.load(path, process=False) + if hasattr(m, "geometry"): + m = trimesh.util.concatenate(list(m.geometry.values())) + v = (np.asarray(m.vertices, np.float64) * scale).astype(np.float32) + f = np.asarray(m.faces, np.int32) + return v, f + + +# --------------------------------------------------------------------------- # +# metric core: given (verts, faces, R, t, K, visible_mask, sensor_depth[opt], +# amodal_mask[opt]) -> row dict. Identical math to verify_reproject_v2.gate. +# --------------------------------------------------------------------------- # +def eval_one(verts, faces, R, t, K, seg, depth_s, amodal, ctx, torch, dr): + H, W = seg.shape + R = np.asarray(R, np.float64); t = np.asarray(t, np.float64) + d_r = render_depth(verts, faces, R, t, K, H, W, ctx, torch, dr) + d_r = np.flipud(d_r).copy() + sil = d_r > 1e-6 + if sil.sum() < 50: + return None + row = {} + if amodal is not None and amodal.shape == sil.shape and amodal.sum() > 50: + inter = np.logical_and(sil, amodal).sum(); union = np.logical_or(sil, amodal).sum() + row["amodal_iou"] = inter / union if union else 0.0 + vi = np.logical_and(sil, seg).sum(); vu = np.logical_or(sil, seg).sum() + row["vis_iou"] = vi / vu if vu else 0.0 + row["coverage"] = np.logical_and(sil, seg).sum() / seg.sum() + if depth_s is not None: + valid = seg & (depth_s > 1e-3) + nden = int(valid.sum()) + if nden >= 50: + both = valid & sil + err = np.abs(d_r[both] - depth_s[both]) + row["depth_agree"] = float((err < DEPTH_TOL).sum()) / nden + if both.sum() >= 20: + row["depth_med"] = float(np.median(np.abs(d_r[both] - depth_s[both]))) + row["depth_medsigned"] = float(np.median(d_r[both] - depth_s[both])) + off = row["depth_medsigned"] + row["depth_agree_debias"] = float( + (np.abs((d_r[both] - depth_s[both]) - off) < DEPTH_TOL).sum()) / nden + return row + + +# --------------------------------------------------------------------------- # +# Per-dataset sample enumerators: yield dict(verts,faces,R,t,K,seg,depth,amodal) +# --------------------------------------------------------------------------- # +def iter_ho3d(n, rng, meshcache): + import ho3d_loader as L + seqs = L.EVAL_SEQUENCES + # build a flat (seq, fid) list + pairs = [] + for seq in seqs: + for fid in L.list_frame_ids(seq): + pairs.append((seq, fid)) + order = rng.permutation(len(pairs)) + got = 0 + for oi in order: + if got >= n: + break + seq, fid = pairs[oi] + meta = L.load_meta(seq, fid) + pose = L.gt_pose(meta) if meta is not None else None + if pose is None: + continue + R, t = pose + K = np.asarray(meta["camMat"], np.float64).reshape(3, 3) + depth = L.load_depth_m(seq, fid) + H, W = depth.shape + amodal = L.render_gt_mask(meta, K, H, W) # rendered AMODAL silhouette + seg = amodal.copy() # HO3D ships only this (full) mask + if seg.sum() < 150: + continue + obj = meta["objName"] + if obj not in meshcache: + meshcache[obj] = L._load_ycb_mesh(obj) # (V,3) m, (F,3) + v, f = meshcache[obj] + yield dict(verts=v.astype(np.float32), faces=f.astype(np.int32), + R=R, t=t, K=K, seg=seg, depth=depth, amodal=amodal, + tag=f"{seq}/{fid}:{obj}") + got += 1 + + +YCBM_ALIGNED = os.path.expandvars("${MIGRATOR_DATA_ROOT}")+"/ycbm/models_aligned/ycb_models_aligned_cm" + + +def _ycbm_mesh(obj, meshcache): + if obj in meshcache: + return meshcache[obj] + for sub in ("google_16k/textured.obj", "textured.obj", "poisson/textured.obj"): + p = os.path.join(YCBM_ALIGNED, obj, sub) + if os.path.isfile(p): + break + v, f = _mesh_load_m(p, 1.0 / 100.0) # cm -> m (matches ycbm_metric._load_cad_m) + meshcache[obj] = (v, f) + return v, f + + +def iter_ycbm(n, rng, meshcache, rot_variant="as_is"): + import ycbm_loader as L + scenes = L.list_scenes() + tasks = [] + for sc in scenes: + for obj in L.scene_objects(sc): + for fid in L.list_frame_ids(sc): + tasks.append((sc, fid, obj)) + order = rng.permutation(len(tasks)) + got = 0 + for oi in order: + if got >= n: + break + sc, fid, obj = tasks[oi] + pose = L.object_pose(sc, fid, obj) + if pose is None: + continue + R, t = pose + R, t = _apply_rot_variant(R, t, rot_variant) + K = L.intrinsics_K(sc) + seg = L.load_mask(sc, fid, obj) + if seg.sum() < 150: + continue + depth = L.load_depth_m(sc, fid) + v, f = _ycbm_mesh(obj, meshcache) + yield dict(verts=v, faces=f, R=R, t=t, K=K, seg=seg, depth=depth, + amodal=None, tag=f"{sc}/{fid}:{obj}") + got += 1 + + +HOPE_MESH_DIR = os.path.expandvars("${MIGRATOR_DATA_ROOT}")+"/hope_video_repo/meshes/eval" + + +def iter_hope(n, rng, meshcache): + import hope_video_loader as L + scenes = L.list_scenes() + tasks = [] + for sc in scenes: + for obj in L.scene_objects(sc): + for fid in L.list_frame_ids(sc): + tasks.append((sc, fid, obj)) + order = rng.permutation(len(tasks)) + got = 0 + for oi in order: + if got >= n: + break + sc, fid, obj = tasks[oi] + pose = L.object_pose(sc, fid, obj) + if pose is None: + continue + R, t = pose + K = L.intrinsics_K(sc, fid) + seg = L.load_mask(sc, fid, obj) + if seg.sum() < 150: + continue + depth = L.load_depth_m(sc, fid) + if obj not in meshcache: + meshcache[obj] = _mesh_load_m(os.path.join(HOPE_MESH_DIR, f"{obj}.obj"), 1.0 / 100.0) + v, f = meshcache[obj] + yield dict(verts=v, faces=f, R=R, t=t, K=K, seg=seg, depth=depth, + amodal=None, tag=f"{sc}/{fid}:{obj}") + got += 1 + + +# --------------------------------------------------------------------------- # +# YCB-M rotation-variant probe. t is held FIXED (translation was independently +# validated by ycbm_calib); only the ORIENTATION R is varied to test handedness. +# --------------------------------------------------------------------------- # +_Fyz = np.diag([1.0, -1.0, -1.0]) +_Fxz = np.diag([-1.0, 1.0, -1.0]) + + +def _apply_rot_variant(R, t, variant): + if variant == "as_is": + return R, t + if variant == "transpose": + return R.T, t + if variant == "conj_yz": # diag(1,-1,-1) R diag(1,-1,-1) + return _Fyz @ R @ _Fyz, t + if variant == "conj_xz": + return _Fxz @ R @ _Fxz, t + if variant == "leftmul_yz": # reflect cam axes on R only + return _Fyz @ R, t + raise ValueError(variant) + + +PRIMARY = {"ho3d": "amodal_iou", "ycbm": "coverage", "hope": "coverage"} +ITERS = {"ho3d": iter_ho3d, "ycbm": iter_ycbm, "hope": iter_hope} + + +def run(name, n, seed, ctx, torch, dr, rot_variant="as_is", verbose=False): + rng = np.random.RandomState(seed) + meshcache = {} + rows = [] + it = ITERS[name] + kw = {"rot_variant": rot_variant} if name == "ycbm" else {} + for s in it(n, rng, meshcache, **kw): + row = eval_one(s["verts"], s["faces"], s["R"], s["t"], s["K"], + s["seg"], s["depth"], s["amodal"], ctx, torch, dr) + if row is None: + continue + rows.append(row) + if verbose and len(rows) <= 4: + print(f" {s['tag']}: {json.dumps({k: round(x,3) for k,x in row.items()})}") + + def col(key): + vals = [r[key] for r in rows if key in r] + return np.array(vals) if vals else np.array([np.nan]) + + out = {"name": name, "N": len(rows), "rot_variant": rot_variant} + for key in ["amodal_iou", "coverage", "vis_iou", "depth_agree", + "depth_agree_debias", "depth_med", "depth_medsigned"]: + c = col(key) + out[key] = float(np.nanmean(c)); out[key + "_n"] = int(np.sum(~np.isnan(c))) + prim = PRIMARY[name] + out["primary_metric"] = prim; out["primary_value"] = out.get(prim, float("nan")) + return out + + +def fmt(o): + def g(k): + v = o.get(k, float("nan")) + return " n/a " if (v != v) else f"{v:.3f}" + return (f"[{o['name']:6s}|{o.get('rot_variant',''):9s}] N={o['N']:2d} " + f"amodalIoU={g('amodal_iou')}({o['amodal_iou_n']}) coverage={g('coverage')} " + f"vis_iou={g('vis_iou')} depth_agree={g('depth_agree')}(debias={g('depth_agree_debias')}) " + f"med|dz|={g('depth_med')}m PRIMARY[{o['primary_metric']}]={g('primary_value')}") + + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("--only", default=None) + ap.add_argument("--n", type=int, default=25) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--ycbm_variants", action="store_true", + help="for ycbm, sweep rotation variants (handedness probe)") + ap.add_argument("--verbose", action="store_true") + args = ap.parse_args() + + import torch + import nvdiffrast.torch as dr + ctx = dr.RasterizeCudaContext() + print("[eval-v2] nvdiffrast ctx ok", flush=True) + + names = ["ho3d", "ycbm", "hope"] + results = [] + for nm in names: + if args.only and nm != args.only: + continue + try: + if nm == "ycbm" and args.ycbm_variants: + for var in ["as_is", "transpose", "conj_yz", "conj_xz", "leftmul_yz"]: + o = run(nm, args.n, args.seed, ctx, torch, dr, rot_variant=var, + verbose=args.verbose) + results.append(o); print(fmt(o), flush=True) + else: + o = run(nm, args.n, args.seed, ctx, torch, dr, verbose=args.verbose) + results.append(o); print(fmt(o), flush=True) + except Exception as e: + import traceback + print(f"[{nm}] ERROR: {e}"); traceback.print_exc() + print("\n=== SUMMARY ===") + for o in results: + print(fmt(o)) diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/verify_reproject_v2.py b/migrator/code/mv-sam3d-for-6d-v2/tools/verify_reproject_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..8382300d1292f93c4c31f7fd0d8e31deaa52b519 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/verify_reproject_v2.py @@ -0,0 +1,265 @@ +"""RIGOROUS GT-as-pred reprojection validator for the integrated-dataloader adapters. + +Unlike tools/verify_adapter_reproject.py (coarse voxel-point splat of the mesh vs the +VISIBLE occlusion-aware mask -> undershoots for correct conventions), this tool: + + 1. RASTERIZES the mesh at the dataset GT pose (adapter.object_pose_cv + adapter.mesh_for) + with the GT intrinsics using a REAL z-buffer. `render_depth` is copied VERBATIM from + ho3d_eval/ycbv_vsd_scorer.py (same nvdiffrast rasterizer used by the VSD scorer); + it returns a metric depth map + a clean silhouette. + + 2. Computes THREE occlusion-fair metrics per sample: + (A) amodal-IoU : IoU(rendered silhouette, AMODAL full-object mask). BOP datasets + (YCB-V, PACE) ship mask/ (amodal) next to mask_visib/ -> the + full-silhouette-vs-visible mismatch disappears. + (B) depth-agree : fraction of visible-GT-mask pixels (with valid sensor depth) where + |rendered_depth - sensor_depth| < DEPTH_TOL (default 2cm). + Occlusion-robust (only where the object is actually seen) AND + catches depth-scale/unit errors the silhouette cannot. + (C) coverage : |visible_mask ∩ rendered_silhouette| / |visible_mask| + (a correct pose renders a silhouette that CONTAINS the visible mask). + Also reported for diagnosis: visible-mask silhouette IoU (the OLD metric), and the + MEDIAN signed/abs depth error (to distinguish a systematic sensor bias from a bug). + + 3. The PRIMARY pass/fail per dataset = the most appropriate occlusion-fair metric >= 0.80: + amodal available (ycbv, pace) -> amodal-IoU + visible-only (dexycb, graspnet, housecat6d, hocap, contactpose) -> depth-agree + +Run in the any6d env (torch + nvdiffrast): + python tools/verify_reproject_v2.py --only dexycb +""" +from __future__ import annotations +import os, sys, argparse, json, glob +import numpy as np + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +from mvsam3d.data.adapters import ADAPTERS +from PIL import Image + +DEPTH_TOL = 0.02 # 2 cm + + +# --------------------------------------------------------------------------- # +# render_depth: copied VERBATIM from ho3d_eval/ycbv_vsd_scorer.py (lines 78-95) +# (importing the module directly would drag in bop_toolkit + ycbv_loader). +# --------------------------------------------------------------------------- # +def render_depth(verts_m, faces, R, t, K, H, W, ctx, torch, dr): + fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2] + Rt = torch.from_numpy(R.astype(np.float32)).cuda() + tt = torch.from_numpy(t.reshape(3).astype(np.float32)).cuda() + vt = torch.from_numpy(verts_m).cuda() + vc = (Rt @ vt.T).T + tt + z = vc[:, 2].clamp(min=1e-6) + u = vc[:, 0] / z * fx + cx; v = vc[:, 1] / z * fy + cy + x = (u / W) * 2 - 1; y = 1 - (v / H) * 2; zb = z / (z.max() + 1e-6) + w = torch.ones(len(verts_m), device="cuda") + clip = torch.stack([x, y, zb, w], 1).unsqueeze(0) + ft = torch.from_numpy(faces).int().cuda() + rast, _ = dr.rasterize(ctx, clip, ft, resolution=[H, W]) + zi, _ = dr.interpolate(z.unsqueeze(1).unsqueeze(0), rast, ft) + d = zi[0, :, :, 0].cpu().numpy() + msk = (rast[0, :, :, 3] > 0).cpu().numpy() + d[~msk] = 0.0 + return d + + +def _mesh_vf(mesh): + v = np.asarray(mesh.vertices, np.float32) + f = np.asarray(mesh.faces, np.int32) + return v, f + + +# --------------------------------------------------------------------------- # +# amodal-mask loaders (BOP mask/ dir, same gt-index as mask_visib/) +# --------------------------------------------------------------------------- # +def _amodal_bop(adapter, ref): + """YCB-V / PACE: scene/mask/{imid:06d}_{gi:06d}.png (amodal full silhouette).""" + gi = adapter._gi(ref) + p = os.path.join(ref["scene"], "mask", f"{ref['imid']:06d}_{gi:06d}.png") + if not os.path.exists(p): + return None + return (np.asarray(Image.open(p)) > 0) + + +AMODAL = {"ycbv": _amodal_bop, "pace": _amodal_bop} + +# PRIMARY occlusion-fair gate per dataset: +# * amodal mask shipped -> amodal-IoU (strict full-silhouette vs full-silhouette; +# not foolable by over-projection, occlusion-immune). [ycbv, pace] +# * visible-only masks -> COVERAGE = |visible ∩ rendered| / |visible| +# (the occlusion-robust silhouette metric: a correct pose renders a full +# silhouette that CONTAINS the visible mask). Corroborated (not replaced) by +# depth-agree (rules out silhouette bloat / scale-unit errors -> reported too). +PRIMARY = { + "ycbv": "amodal_iou", "pace": "amodal_iou", + "dexycb": "coverage", "graspnet": "coverage", + "housecat6d": "coverage", "hocap": "coverage", "contactpose": "coverage", + # H2O ships NO masks -> the adapter renders the mask from mesh+pose, so + # `coverage` (rendered_sil vs that same rendered mask) is trivially ~1.0 and + # NOT trustworthy. The honest, non-gameable gate is depth-agree: rendered + # metric depth (GT pose) vs the independent RGB-D SENSOR depth within 2 cm. + "h2o": "depth_agree", + # HOGraspNet: grasp = heavy hand occlusion, obj_mask is a PSEUDO (rendered) + # mask -> visible-silhouette IoU undershoots. `coverage` (rendered full + # silhouette CONTAINS the visible pseudo-mask) is the occlusion-robust gate, + # corroborated by depth-agree (independent Kinect depth, unit/scale-sensitive). + "hograspnet": "coverage", + # DexH2R: human->robot handover => heavy hand/robot occlusion, and masks are + # rendered from mesh+pose (so `coverage` is trivially ~1.0 and not trustworthy, + # same as H2O). The honest, non-gameable gate is depth-agree: rendered metric + # depth (GT pose) vs the independent Kinect SENSOR depth within 2 cm. + "dexh2r": "depth_agree", +} + + +def gate(name, root, cache, ctx, torch, dr, n=25, stride=None, seed=0, verbose=False): + kw = {} + if stride is not None: + kw["frame_stride"] = stride + ad = ADAPTERS[name](root, cache, **kw) + samples = ad.list_samples() + rng = np.random.RandomState(seed) + order = rng.permutation(len(samples)) + amodal_fn = AMODAL.get(name) + + rows = [] # per-sample metric dict + meshcache = {} + tried = 0 + for si in order: + if len(rows) >= n or tried > n * 40: + break + tried += 1 + s = samples[si] + try: + v = ad.load_view(s.anchor) + R_cv, t_cv = ad.object_pose_cv(s) + R_cv = np.asarray(R_cv, np.float64); t_cv = np.asarray(t_cv, np.float64) + if s.obj_key not in meshcache: + meshcache[s.obj_key] = _mesh_vf(ad.mesh_for(s)) + verts, faces = meshcache[s.obj_key] + except Exception as e: + if tried <= 3: + print(f" [skip:{name}] {type(e).__name__}: {e}") + continue + + seg = np.asarray(v["mask"]) > 0.5 + H, W = seg.shape + if seg.sum() < 150: + continue + depth_s = np.asarray(v["depth_m"], np.float64) + intr = v["intr"] + K = np.array([[intr["fx"], 0, intr["ppx"]], [0, intr["fy"], intr["ppy"]], [0, 0, 1]], np.float64) + + try: + d_r = render_depth(verts, faces, R_cv, t_cv, K, H, W, ctx, torch, dr) + except Exception as e: + if tried <= 3: + print(f" [render-skip:{name}] {type(e).__name__}: {e}") + continue + # nvdiffrast rasterizes with a bottom-origin (OpenGL) y-axis, so its output + # is vertically flipped relative to top-left-origin image arrays (masks/depth). + # Flip back so the rendered depth registers to the dataset image frame. + # (Verified on DexYCB, correct by construction: raw IoU 0.0 -> flipud IoU 0.98.) + d_r = np.flipud(d_r).copy() + sil = d_r > 1e-6 # rendered full silhouette + if sil.sum() < 50: + continue + + row = {} + # (A) amodal IoU + if amodal_fn is not None: + ma = amodal_fn(ad, s.anchor) + if ma is not None and ma.shape == sil.shape and ma.sum() > 50: + inter = np.logical_and(sil, ma).sum(); union = np.logical_or(sil, ma).sum() + row["amodal_iou"] = inter / union if union else 0.0 + + # visible-mask silhouette IoU (OLD metric, for reference) + vi = np.logical_and(sil, seg).sum(); vu = np.logical_or(sil, seg).sum() + row["vis_iou"] = vi / vu if vu else 0.0 + # (C) coverage + row["coverage"] = np.logical_and(sil, seg).sum() / seg.sum() + + # (B) depth agreement (occlusion-robust, unit/scale sensitive) + valid = seg & (depth_s > 1e-3) # object visible with valid sensor depth + nden = int(valid.sum()) + if nden >= 50: + both = valid & sil # also rendered here + err = np.abs(d_r[both] - depth_s[both]) + row["depth_agree"] = float((err < DEPTH_TOL).sum()) / nden + if both.sum() >= 20: + row["depth_med"] = float(np.median(np.abs(d_r[both] - depth_s[both]))) + row["depth_medsigned"] = float(np.median(d_r[both] - depth_s[both])) + # bias-corrected agreement: subtract median signed offset (sensor bias diag) + off = row["depth_medsigned"] + row["depth_agree_debias"] = float((np.abs((d_r[both] - depth_s[both]) - off) < DEPTH_TOL).sum()) / nden + rows.append(row) + if verbose and len(rows) <= 3: + print(f" sample {len(rows)}: {json.dumps({k: round(x,3) for k,x in row.items()})}") + + def col(key): + vals = [r[key] for r in rows if key in r] + return np.array(vals) if vals else np.array([np.nan]) + + out = {"name": name, "N": len(rows), "n_samples": len(samples)} + for key in ["amodal_iou", "depth_agree", "depth_agree_debias", "coverage", "vis_iou", + "depth_med", "depth_medsigned"]: + c = col(key) + out[key] = float(np.nanmean(c)) + out[key + "_n"] = int(np.sum(~np.isnan(c))) + prim = PRIMARY.get(name, "depth_agree") + out["primary_metric"] = prim + out["primary_value"] = out.get(prim, float("nan")) + return out + + +def fmt(o): + def g(k): + v = o.get(k, float("nan")) + return " n/a " if (v != v) else f"{v:.3f}" + return (f"[{o['name']:11s}] N={o['N']:2d} amodalIoU={g('amodal_iou')}({o['amodal_iou_n']}) " + f"depth_agree={g('depth_agree')} (debias={g('depth_agree_debias')}) " + f"coverage={g('coverage')} vis_iou={g('vis_iou')} " + f"med|dz|={g('depth_med')}m PRIMARY[{o['primary_metric']}]={g('primary_value')}") + + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("--only", default=None) + ap.add_argument("--n", type=int, default=25) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--verbose", action="store_true") + args = ap.parse_args() + + import torch + import nvdiffrast.torch as dr + ctx = dr.RasterizeCudaContext() + print("[v2] nvdiffrast ctx ok", flush=True) + + CACHE = os.path.expandvars("${MIGRATOR_CACHE}")+"/canonical" + DS = os.path.expandvars("${MIGRATOR_DATA_ROOT}")+"" + jobs = [ + ("dexycb", f"{DS}/dexycb"), + ("housecat6d", f"{DS}/housecat6d/extracted"), + ("pace", f"{DS}/pace"), + ("graspnet", f"{DS}/graspnet"), + ("ycbv", f"{DS}/ycbv"), + ("hocap", f"{DS}/hocap"), + ("contactpose",f"{DS}/contactpose"), + ("h2o", f"{DS}/h2o"), + ("hograspnet", f"{DS}/hograspnet"), + ("dexh2r", f"{DS}/dexh2r"), + ] + results = [] + for nm, root in jobs: + if args.only and nm != args.only: + continue + try: + o = gate(nm, root, CACHE, ctx, torch, dr, n=args.n, seed=args.seed, verbose=args.verbose) + results.append(o) + print(fmt(o), flush=True) + except Exception as e: + import traceback + print(f"[{nm}] ERROR: {e}"); traceback.print_exc() + print("\n=== SUMMARY ===") + for o in results: + print(fmt(o)) diff --git a/migrator/code/mv-sam3d-for-6d-v2/tools/verify_rotation_roundtrip.py b/migrator/code/mv-sam3d-for-6d-v2/tools/verify_rotation_roundtrip.py new file mode 100644 index 0000000000000000000000000000000000000000..eec73c67a6361da3e33b51d88cce76d73342edd2 --- /dev/null +++ b/migrator/code/mv-sam3d-for-6d-v2/tools/verify_rotation_roundtrip.py @@ -0,0 +1,47 @@ +"""Rotation 6D round-trip test (CPU-only). + +Verifies the dataset encode and model decode are exact inverses on SO(3): + R_hat = rot6d_to_matrix(denormalize_rot6d(normalize_rot6d(matrix_to_rot6d(R)))) + assert max|R_hat - R| < 1e-5 + +Uses the VENDORED column-convention utilities from +``sam3d_objects.geometry.geometry`` (the same functions the +dataset and model import), so this checks the actual code path, not a reimpl. +""" +from __future__ import annotations + +import argparse + +import torch +from pytorch3d.transforms import random_rotations + +from sam3d_objects.geometry.geometry import ( + denormalize_rot6d, + matrix_to_rot6d, + normalize_rot6d, + rot6d_to_matrix, +) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n", type=int, default=10000) + ap.add_argument("--tol", type=float, default=1e-5) + args = ap.parse_args() + + torch.manual_seed(0) + R = random_rotations(args.n).double() # [N,3,3] uniform SO(3) + enc = normalize_rot6d(matrix_to_rot6d(R)) # encode (dataset path) + R_hat = rot6d_to_matrix(denormalize_rot6d(enc)) # decode (model path) + err = (R_hat - R).abs().max().item() + # Also confirm decoded matrices are proper rotations. + det = torch.linalg.det(R_hat) + orth = (R_hat @ R_hat.transpose(-1, -2) - torch.eye(3, dtype=R_hat.dtype)).abs().max().item() + print(f"[rot-roundtrip] N={args.n} max|R_hat-R|={err:.3e} " + f"det in [{det.min():.6f},{det.max():.6f}] max|RR^T-I|={orth:.3e}") + assert err < args.tol, f"round-trip error {err:.3e} >= tol {args.tol:.3e}" + print(f"[rot-roundtrip] PASS (< {args.tol:.0e})") + + +if __name__ == "__main__": + main() diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/__init__.py b/migrator/code/sam3d-for-6d/sam3d_objects/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a700544f5c8388939eff112e419a851dbabf7671 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import os + +# Allow skipping initialization for lightweight tools +if not os.environ.get('LIDRA_SKIP_INIT'): + import sam3d_objects.init diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/pipeline_use_scene_scale_false.yaml b/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/pipeline_use_scene_scale_false.yaml new file mode 100644 index 0000000000000000000000000000000000000000..132d0fac1f7599b5da974fd459ea18eed9a60b3b --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/pipeline_use_scene_scale_false.yaml @@ -0,0 +1,106 @@ +_target_: sam3d_objects.pipeline.inference_pipeline_pointmap.InferencePipelinePointMap +ss_generator_config_path: ss_generator.yaml +ss_generator_ckpt_path: ss_generator.ckpt +slat_generator_config_path: slat_generator.yaml +slat_generator_ckpt_path: slat_generator.ckpt +ss_decoder_config_path: ss_decoder.yaml +ss_decoder_ckpt_path: ss_decoder.ckpt +slat_decoder_gs_config_path: slat_decoder_gs.yaml +slat_decoder_gs_ckpt_path: slat_decoder_gs.ckpt +slat_decoder_gs_4_config_path: slat_decoder_gs_4.yaml +slat_decoder_gs_4_ckpt_path: slat_decoder_gs_4.ckpt +slat_decoder_mesh_config_path: slat_decoder_mesh.yaml +slat_decoder_mesh_ckpt_path: slat_decoder_mesh.ckpt +pad_size: 1.0 +dtype: float16 +version: 3dfy_v9 +slat_cfg_strength: 1 +slat_rescale_t: 1 +downsample_ss_dist: 1 +compile_model: true +ss_condition_input_mapping: [] +ss_preprocessor: + _target_: sam3d_objects.data.dataset.tdfy.preprocessor.PreProcessor + img_mask_joint_transform: [] + img_mask_pointmap_joint_transform: + - _partial_: true + _target_: sam3d_objects.data.dataset.tdfy.img_and_mask_transforms.resize_all_to_same_size + - _partial_: true + _target_: sam3d_objects.data.dataset.tdfy.img_and_mask_transforms.crop_around_mask_with_padding + box_size_factor: 1.2 + padding_factor: 0.0 + img_transform: + _target_: torchvision.transforms.Compose + transforms: + - _partial_: true + _target_: sam3d_objects.data.dataset.tdfy.img_processing.pad_to_square_centered + - _target_: torchvision.transforms.Resize + size: 518 + mask_transform: + _target_: torchvision.transforms.Compose + transforms: + - _partial_: true + _target_: sam3d_objects.data.dataset.tdfy.img_processing.pad_to_square_centered + - _target_: torchvision.transforms.Resize + interpolation: 0 + size: 518 + normalize_pointmap: true + pointmap_normalizer: + _target_: sam3d_objects.data.dataset.tdfy.img_and_mask_transforms.ObjectCentricSSI + allow_scale_and_shift_override: true + use_scene_scale: false + pointmap_transform: + _target_: torchvision.transforms.Compose + transforms: + - _partial_: true + _target_: sam3d_objects.data.dataset.tdfy.img_processing.pad_to_square_centered + - _target_: torchvision.transforms.Resize + interpolation: 0 + size: 518 +pose_decoder_name: ScaleShiftInvariant +depth_model: + _target_: sam3d_objects.pipeline.depth_models.moge.MoGe + model: + _target_: moge.model.v1.MoGeModel.from_pretrained + pretrained_model_name_or_path: Ruicheng/moge-vitl +slat_condition_input_mapping: [] +slat_preprocessor: + _target_: sam3d_objects.data.dataset.tdfy.preprocessor.PreProcessor + img_transform: + _target_: torchvision.transforms.Compose + transforms: + - _target_: sam3d_objects.data.dataset.tdfy.img_processing.pad_to_square_centered + _partial_: true + - _target_: torchvision.transforms.Resize + size: 518 + mask_transform: + _target_: torchvision.transforms.Compose + transforms: + - _target_: sam3d_objects.data.dataset.tdfy.img_processing.pad_to_square_centered + _partial_: true + - _target_: torchvision.transforms.Resize + size: 518 + interpolation: 0 + img_mask_joint_transform: + - _target_: sam3d_objects.data.dataset.tdfy.img_and_mask_transforms.crop_around_mask_with_padding + _partial_: true + box_size_factor: 1.2 + padding_factor: 0.0 +slat_mean: +- 0.12211431 +- 0.37204156 +- -1.26521907 +- -2.05276058 +- -3.10432536 +- -0.11294304 +- -0.85146744 +- 0.45506954 +slat_std: +- 2.37326008 +- 2.13174402 +- 2.2413953 +- 2.30589401 +- 2.1191894 +- 1.8969511 +- 2.41684989 +- 2.08374642 diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/slat_decoder_gs.yaml b/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/slat_decoder_gs.yaml new file mode 100644 index 0000000000000000000000000000000000000000..10683b1bd592e1db70eb7070a3365b536ad0ba9e --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/slat_decoder_gs.yaml @@ -0,0 +1,25 @@ +_target_: sam3d_objects.model.backbone.tdfy_dit.models.structured_latent_vae.decoder_gs.SLatGaussianDecoderTdfyWrapper +resolution: 64 +model_channels: 768 +latent_channels: 8 +num_blocks: 12 +num_heads: 12 +mlp_ratio: 4 +attn_mode: swin +window_size: 8 +representation_config: + lr: + _xyz: 1.0 + _features_dc: 1.0 + _opacity: 1.0 + _scaling: 1.0 + _rotation: 0.1 + perturb_offset: true + voxel_size: 1.5 + num_gaussians: 32 + 2d_filter_kernel_size: 0.1 + 3d_filter_kernel_size: 0.0009 + scaling_bias: 0.004 + opacity_bias: 0.1 + scaling_activation: softplus +use_fp16: true diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/slat_decoder_mesh.yaml b/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/slat_decoder_mesh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..74735a4b9ba1fbd76b8d1a56f27b6ebe63aeafc9 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/slat_decoder_mesh.yaml @@ -0,0 +1,12 @@ +_target_: sam3d_objects.model.backbone.tdfy_dit.models.structured_latent_vae.decoder_mesh.SLatMeshDecoderTdfyWrapper +resolution: 64 +model_channels: 768 +latent_channels: 8 +num_blocks: 12 +num_heads: 12 +mlp_ratio: 4 +attn_mode: swin +window_size: 8 +representation_config: + use_color: true +use_fp16: true diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/ss_encoder.safetensors b/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/ss_encoder.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/ss_encoder.yaml b/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/ss_encoder.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/ss_generator.yaml b/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/ss_generator.yaml new file mode 100644 index 0000000000000000000000000000000000000000..35df0ea780c4c6f6eaf490c1a828ad054d4f9f0b --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/checkpoints/hf/ss_generator.yaml @@ -0,0 +1,141 @@ +module: + condition_embedder: + backbone: + _target_: sam3d_objects.model.backbone.dit.embedder.embedder_fuser.EmbedderFuser + drop_modalities_weight: + - - - pointmap + - rgb_pointmap + - 1.0 + dropout_prob: 0.1 + embedder_list: + - - _target_: sam3d_objects.model.backbone.dit.embedder.dino.Dino + dino_model: dinov2_vitl14_reg + input_size: 518 + normalize_images: true + - - - image + - cropped + - - rgb_image + - full + - - _target_: sam3d_objects.model.backbone.dit.embedder.dino.Dino + dino_model: dinov2_vitl14_reg + input_size: 518 + normalize_images: true + - - - mask + - cropped + - - rgb_image_mask + - full + - - _target_: sam3d_objects.model.backbone.dit.embedder.pointmap.PointPatchEmbed + embed_dim: 512 + input_size: 256 + patch_size: 8 + remap_output: linear + - - - pointmap + - cropped + - - rgb_pointmap + - full + force_drop_modalities: null + freeze: true + projection_net_hidden_dim_multiplier: 4.0 + use_pos_embedding: learned + generator: + backbone: + _target_: sam3d_objects.model.backbone.generator.shortcut.model.ShortCut + batch_mode: true + cfg_modalities: + - shape + inference_steps: 2 + loss_weights: + 6drotation_normalized: 0.1 + _target_: sam3d_objects.config.utils.make_dict + scale: 0.1 + shape: 0 + translation: 1.0 + translation_scale: 0.0 + ratio_cfg_samples_in_self_consistency_target: 0.25 + rescale_t: 1 + reverse_fn: + _target_: sam3d_objects.model.backbone.generator.classifier_free_guidance.ClassifierFreeGuidanceWithExternalUnconditionalProbability + backbone: + _target_: sam3d_objects.model.backbone.tdfy_dit.models.mot_sparse_structure_flow.SparseStructureFlowTdfyWrapper + cond_channels: 1024 + condition_embedder: null + force_zeros_cond: true + freeze_d_time_embedder: true + freeze_shared_parameters: true + in_channels: 8 + is_shortcut_model: true + latent_mapping: + 6drotation_normalized: + _target_: sam3d_objects.model.backbone.tdfy_dit.models.mm_latent.Latent + in_channels: 6 + model_channels: 1024 + pos_embedder: + _target_: sam3d_objects.model.backbone.tdfy_dit.models.mm_latent.LearntPositionEmbedder + model_channels: 1024 + token_len: 1 + scale: + _target_: sam3d_objects.model.backbone.tdfy_dit.models.mm_latent.Latent + in_channels: 3 + model_channels: 1024 + pos_embedder: + _target_: sam3d_objects.model.backbone.tdfy_dit.models.mm_latent.LearntPositionEmbedder + model_channels: 1024 + token_len: 1 + shape: + _target_: sam3d_objects.model.backbone.tdfy_dit.models.mm_latent.Latent + in_channels: 8 + model_channels: 1024 + pos_embedder: + _target_: sam3d_objects.model.backbone.tdfy_dit.models.mm_latent.ShapePositionEmbedder + model_channels: 1024 + patch_size: 1 + resolution: 16 + translation: + _target_: sam3d_objects.model.backbone.tdfy_dit.models.mm_latent.Latent + in_channels: 3 + model_channels: 1024 + pos_embedder: + _target_: sam3d_objects.model.backbone.tdfy_dit.models.mm_latent.LearntPositionEmbedder + model_channels: 1024 + token_len: 1 + translation_scale: + _target_: sam3d_objects.model.backbone.tdfy_dit.models.mm_latent.Latent + in_channels: 1 + model_channels: 1024 + pos_embedder: + _target_: sam3d_objects.model.backbone.tdfy_dit.models.mm_latent.LearntPositionEmbedder + model_channels: 1024 + token_len: 1 + latent_share_transformer: + 6drotation_normalized: + - 6drotation_normalized + - translation + - scale + - translation_scale + mlp_ratio: 4 + model_channels: 1024 + num_blocks: 24 + num_heads: 16 + out_channels: 8 + patch_size: 1 + pe_mode: ape + qk_rms_norm: true + resolution: 16 + use_checkpoint: false + use_fp16: false + interval: + - 0 + - 500 + p_unconditional: 0.1 + strength: 2.0 + unconditional_handling: add_flag + self_consistency_cfg_strength: 2.0 + self_consistency_prob: 0.25 + shortcut_loss_weight: 1.0 + sigma_min: 0.0 + time_scale: 1000.0 + training_time_sampler_fn: + _partial_: true + _target_: sam3d_objects.model.backbone.generator.flow_matching.model.lognorm_sampler + mean: -1.0 + std: 1.0 diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/data/__init__.py b/migrator/code/sam3d-for-6d/sam3d_objects/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..71ca4b12c770afea62d06f97064cdf0c97d40ed7 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/data/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/__init__.py b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..71ca4b12c770afea62d06f97064cdf0c97d40ed7 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/__init__.py b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..71ca4b12c770afea62d06f97064cdf0c97d40ed7 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/img_and_mask_transforms.py b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/img_and_mask_transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..6b5439416a059f67a985fe0bf8b8379108833085 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/img_and_mask_transforms.py @@ -0,0 +1,1058 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from __future__ import annotations +from collections import namedtuple +import random +from typing import Optional, Dict + +import numpy as np +import matplotlib.pyplot as plt +import torchvision.transforms.functional +from sam3d_objects.data.dataset.tdfy.img_processing import pad_to_square_centered +from sam3d_objects.model.backbone.dit.embedder.point_remapper import PointRemapper +from typing import Optional, Dict +from loguru import logger +import torch +import torch.nn.functional as F +import torchvision +import torchvision.transforms as tv_transforms +import torchvision.transforms.functional +import torchvision.transforms.functional as TF + +from sam3d_objects.data.dataset.tdfy.img_processing import pad_to_square_centered + + +def UNNORMALIZE(mean, std): + mean = torch.tensor(mean).reshape((3, 1, 1)) + std = torch.tensor(std).reshape((3, 1, 1)) + + def unnormalize_img(img): + assert img.ndim == 3 and img.shape[0] == 3 + + return img * std.to(img.device) + mean.to(img.device) + + return unnormalize_img + + +IMAGENET_MEAN = (0.485, 0.456, 0.406) +IMAGENET_STD = (0.229, 0.224, 0.225) + + +IMAGENET_NORMALIZATION = tv_transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD) +IMAGENET_UNNORMALIZATION = UNNORMALIZE(IMAGENET_MEAN, IMAGENET_STD) + + +class BoundingBoxError(Exception): + pass + + +def check_bounding_box(bbox_w, bbox_h): + if bbox_w < 0 or bbox_h < 0: + raise BoundingBoxError("Bounding box dimensions must be non-negative.") + + +class RGBAImageProcessor: + def __init__( + self, + resize_and_make_square_kwargs: Optional[Dict] = None, + object_crop_kwargs: Optional[Dict] = None, + remove_background: bool = False, + imagenet_normalization: bool = False, + ): + self.remove_background = remove_background + self.resize_and_pad_kwargs = resize_and_make_square_kwargs + self.object_crop_kwargs = object_crop_kwargs + self.imagenet_normalization = imagenet_normalization + if resize_and_make_square_kwargs is not None: + self.transforms = resize_and_make_square(**resize_and_make_square_kwargs) + + def __call__( + self, image: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: + if mask is None: + assert ( + image.shape[0] == 4 + ), f"Requires 4 channels (RGB + alpha), got {image.shape[0]=}" + image, mask = split_rgba(image) + else: + assert ( + image.shape[0] == 3 + ), f"Requires 3 channels (RGB), got {image.shape[0]=}" + assert mask.dim() == 2, f"Requires 2D mask, got {mask.dim()=}" + + if not self.object_crop_kwargs in [None, False]: + image, mask = crop_around_mask_with_padding( + image, mask, **self.object_crop_kwargs + ) + + if self.remove_background: + image, mask = rembg(image, mask) + + image = self.transforms["img_transform"](image) + mask = self.transforms["mask_transform"](mask.unsqueeze(0)) + + if self.imagenet_normalization: + image = IMAGENET_NORMALIZATION(image) + return image, mask + + +def load_rgb(fpath: str) -> torch.Tensor: + """ + Load a RGB(A) image from a file path. + """ + image = plt.imread(fpath) # Why use matplotlib? + if image.dtype == "uint8": + image = image / 255 + image = image.astype(np.float32) + image = torch.from_numpy(image) + image = image.permute(2, 0, 1).contiguous() + return image + + +def concat_rgba( + rgb_image: torch.Tensor, + mask: torch.Tensor, +) -> torch.Tensor: + """ + Create a 4-channel RGBA image from a 3-channel RGB image and a mask. + """ + assert rgb_image.dim() == 3, f"{rgb_image.shape=}" + assert mask.dim() == 2, f"{mask.shape=}" + assert rgb_image.shape[0] == 3, f"{rgb_image.shape[0]=}" + assert rgb_image.shape[1:] == mask.shape, f"{rgb_image.shape[1:]=} != {mask.shape=}" + return torch.cat((rgb_image, mask[None, ...]), dim=0) + + +def split_rgba(rgba_image: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Split a 4-channel RGBA image into a 3-channel RGB image and a 1-channel mask. + + Args: + rgba_image: A 4-channel RGBA image. + + Returns: + A tuple of (rgb_image, mask). + """ + assert rgba_image.dim() == 3, f"{rgba_image.shape=}" + assert rgba_image.shape[0] == 4, f"{rgba_image.shape[0]=}" + return rgba_image[:3], rgba_image[3] + + +def get_mask( + rgb_image: torch.Tensor, + depth_image: torch.Tensor, + mask_source: str, +) -> torch.Tensor: + """ + Extract a mask from either the alpha channel of an RGB image or a depth image. + + Args: + rgb_image: Tensor of shape (B, C, H, W) or (C, H, W) where C >= 4 if using alpha channel + depth_image: Tensor of shape (B, 1, H, W) or (1, H, W) containing depth information + mask_source: Source of the mask, either "ALPHA_CHANNEL" or "DEPTH" + + Returns: + mask: Tensor of shape (B, 1, H, W) or (1, H, W) containing the extracted mask + """ + # Handle unbatched inputs (add batch dimension if needed) + is_batched = len(rgb_image.shape) == 4 + + if not is_batched: + rgb_image = rgb_image.unsqueeze(0) + if depth_image is not None: + depth_image = depth_image.unsqueeze(0) + + if mask_source == "ALPHA_CHANNEL": + if rgb_image.shape[1] != 4: + logger.warning(f"No ALPHA CHANNEL for the image, cannot read mask.") + mask = None + else: + mask = rgb_image[:, 3:4, :, :] + elif mask_source == "DEPTH": + mask = depth_image + else: + raise ValueError(f"Invalid mask source: {mask_source}") + + # Remove batch dimension if input was unbatched + if not is_batched: + mask = mask.squeeze(0) + + return mask + + +def rembg(image, mask, pointmap=None): + """ + Remove the background from an image using a mask. + For pointmaps, sets background regions to NaN. + + This function follows the standard transform pattern: + - If called with (image, mask), returns (image, mask) + - If called with (image, mask, pointmap), returns (image, mask, pointmap) + """ + masked_image = image * mask + + if pointmap is not None: + masked_pointmap = torch.where(mask > 0, pointmap, torch.nan) + return masked_image, mask, masked_pointmap + + return masked_image, mask + + +def resize_and_make_square( + img_size: int, + make_square: bool | str = False, +): + """ + Create image and mask transforms based on configuration. + + Returns: + dict: {"img_transform": img_transform, "mask_transform": mask_transform} + """ + if isinstance(make_square, str): + make_square = make_square.lower() + assert make_square in ["pad", "crop", False] + pre_resize_transform = tv_transforms.Lambda(lambda x: x) + post_resize_transform = tv_transforms.Lambda(lambda x: x) + if make_square == "pad": + pre_resize_transform = pad_to_square_centered + elif make_square == "crop": + post_resize_transform = tv_transforms.CenterCrop(img_size) + + img_resize = tv_transforms.Resize(img_size) + mask_resize = tv_transforms.Resize( + img_size, + interpolation=tv_transforms.InterpolationMode.BILINEAR, + ) + + img_transform = tv_transforms.Compose( + [ + pre_resize_transform, + img_resize, + post_resize_transform, + ] + ) + + mask_transform = tv_transforms.Compose( + [ + pre_resize_transform, + mask_resize, + post_resize_transform, + ] + ) + + return { + "img_transform": img_transform, + "mask_transform": mask_transform, + } + + +def crop_around_mask_with_random_box_size_factor( + loaded_image: torch.Tensor, + mask: torch.Tensor, + random_box_size_factor: float = 1.0, + pointmap: Optional[torch.Tensor] = None, +) -> np.ndarray: + return crop_around_mask_with_padding( + loaded_image, + mask, + box_size_factor=1.0 + random.uniform(0, 1) * random_box_size_factor, + padding_factor=0.0, + pointmap=pointmap, + ) + + +def crop_around_mask_with_padding( + loaded_image: torch.Tensor, + mask: torch.Tensor, + box_size_factor: float = 1.6, + padding_factor: float = 0.1, + pointmap: Optional[torch.Tensor] = None, +) -> np.ndarray: + # cast to ensure the function can be called normally + cast_mask = False + if mask.dim() == 3: + assert mask.shape[0] == 1, "cannot take mask with channel dimension not 1" + mask = mask[0] + cast_mask = True + loaded_image = concat_rgba(loaded_image, mask) + + try: + bbox = compute_mask_bbox(mask, box_size_factor) + except BoundingBoxError as exc: + logger.warning( + "Invalid mask bbox encountered during crop; falling back to full-frame crop. {}", + exc, + ) + height, width = mask.shape + bbox = (0, 0, max(int(width), 2), max(int(height), 2)) + + crop_h = max(int(bbox[3] - bbox[1]), 2) + crop_w = max(int(bbox[2] - bbox[0]), 2) + loaded_image = torchvision.transforms.functional.crop( + loaded_image, int(bbox[1]), int(bbox[0]), crop_h, crop_w + ) + + # Crop pointmap if provided + if pointmap is not None: + pointmap = torchvision.transforms.functional.crop( + pointmap, int(bbox[1]), int(bbox[0]), crop_h, crop_w + ) + + C, H, W = loaded_image.shape + max_dim = max(H, W) # Get the larger dimension + + # Step 1: Pad to square shape + pad_h = (max_dim - H) // 2 + pad_w = (max_dim - W) // 2 + pad_h_extra = (max_dim - H) - pad_h # To ensure even padding + pad_w_extra = (max_dim - W) - pad_w + + loaded_image = torch.nn.functional.pad( + loaded_image, (pad_w, pad_w_extra, pad_h, pad_h_extra), mode="constant", value=0 + ) + if pointmap is not None: + pointmap = torch.nn.functional.pad( + pointmap, + (pad_w, pad_w_extra, pad_h, pad_h_extra), + mode="constant", + value=float("nan"), + ) + + # Step 2: Extend by 10% on each side; idk but this seems to have better results overall + if padding_factor > 0: + extend_size = int(max_dim * padding_factor) # 10% extension on each side + loaded_image = torch.nn.functional.pad( + loaded_image, + (extend_size, extend_size, extend_size, extend_size), + mode="constant", + value=0, + ) + + if pointmap is not None: + pointmap = torch.nn.functional.pad( + pointmap, + (extend_size, extend_size, extend_size, extend_size), + mode="constant", + value=float("nan"), + ) + + rgb_image, mask = split_rgba(loaded_image) + if cast_mask: + mask = mask[None] + + if pointmap is not None: + return rgb_image, mask, pointmap + return rgb_image, mask + + +def compute_mask_bbox( + mask: torch.Tensor, box_size_factor: float = 1.0 +) -> tuple[float, float, float, float]: + """ + Compute a bounding box around a binary mask with optional size adjustment. + + Args: + mask: A 2D binary tensor where non-zero values represent the object of interest. + box_size_factor: Factor to scale the bounding box size. Values > 1.0 create a larger box. + Default is 1.0 (tight bounding box). + + Returns: + A tuple of (x1, y1, x2, y2) coordinates representing the bounding box, + where (x1, y1) is the top-left corner and (x2, y2) is the bottom-right corner. + + Raises: + ValueError: If mask is not a torch.Tensor or not a 2D tensor. + """ + if not isinstance(mask, torch.Tensor): + raise ValueError("Mask must be a torch.Tensor") + if not mask.dim() == 2: + raise ValueError("Mask must be a 2D tensor") + bbox_indices = torch.nonzero(mask) + if bbox_indices.numel() == 0: + logger.warning("Empty mask encountered while computing bbox; falling back to full-frame crop.") + height, width = mask.shape + return (0, 0, max(int(width), 2), max(int(height), 2)) + + y_indices = bbox_indices[:, 0] + x_indices = bbox_indices[:, 1] + + min_x = torch.min(x_indices).item() + min_y = torch.min(y_indices).item() + max_x = torch.max(x_indices).item() + max_y = torch.max(y_indices).item() + + bbox = (min_x, min_y, max_x, max_y) + + center_x = (bbox[0] + bbox[2]) / 2 + center_y = (bbox[1] + bbox[3]) / 2 + + bbox_w, bbox_h = bbox[2] - bbox[0], bbox[3] - bbox[1] + + check_bounding_box(bbox_w, bbox_h) + + size = max(bbox_w, bbox_h, 2) + size = max(int(np.ceil(size * box_size_factor)), 2) + + bbox = ( + int(center_x - size // 2), + int(center_y - size // 2), + int(center_x + size // 2), + int(center_y + size // 2), + ) + if bbox[2] - bbox[0] < 2: + bbox = (bbox[0], bbox[1], bbox[0] + 2, bbox[3]) + if bbox[3] - bbox[1] < 2: + bbox = (bbox[0], bbox[1], bbox[2], bbox[1] + 2) + return bbox + + +def crop_and_pad(image, bbox): + """ + Crop an image using a bounding box and pad with zeros if out of bounds. + + Args: + image (torch.Tensor): CxHxW image. + bbox (tuple): (x1, y1, x2, y2) bounding box. + + Returns: + torch.Tensor: Cropped and zero-padded image. + """ + C, H, W = image.shape + x1, y1, x2, y2 = bbox + + # Ensure coordinates are integers + x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) + + # Compute cropping coordinates + x1_pad, y1_pad = max(0, -x1), max(0, -y1) + x2_pad, y2_pad = max(0, x2 - W), max(0, y2 - H) + + # Compute valid region in the original image + x1_crop, y1_crop = max(0, x1), max(0, y1) + x2_crop, y2_crop = min(W, x2), min(H, y2) + + # Extract the valid part + cropped = image[:, y1_crop:y2_crop, x1_crop:x2_crop] + + # Create a zero-padded output + padded = torch.zeros((C, y2 - y1, x2 - x1), dtype=image.dtype) + + # Place the cropped image into the zero-padded array + padded[ + :, y1_pad : y1_pad + cropped.shape[1], x1_pad : x1_pad + cropped.shape[2] + ] = cropped + + return padded + + +def resize_all_to_same_size( + rgb_image: torch.Tensor, + mask: torch.Tensor, + pointmap: Optional[torch.Tensor] = None, + target_size: Optional[tuple[int, int]] = None, +) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """ + Resize RGB image, mask, and pointmap to the same size. + + This is crucial when pointmaps have different resolution than RGB images, + which must be done BEFORE any cropping operations. + + Args: + rgb_image: RGB image tensor of shape (C, H, W) + mask: Mask tensor of shape (H, W) or (1, H, W) + pointmap: Optional pointmap tensor of shape (C_p, H_p, W_p) + target_size: Target size as (H, W). If None, uses RGB image size. + + Returns: + Tuple of (resized_rgb, resized_mask, resized_pointmap) + """ + squeeze_mask = (mask.dim() == 2) + if squeeze_mask: + mask = mask.unsqueeze(0) + + if target_size is None: + target_size = (rgb_image.shape[1], rgb_image.shape[2]) # (H, W) + + rgb_needs_resize = (rgb_image.shape[1], rgb_image.shape[2]) != target_size + if rgb_needs_resize: + rgb_image = torchvision.transforms.functional.resize( + rgb_image, target_size, interpolation=torchvision.transforms.InterpolationMode.BILINEAR + ) + mask = torchvision.transforms.functional.resize( + mask, target_size, interpolation=torchvision.transforms.InterpolationMode.NEAREST + ) + + if pointmap is not None: + pointmap_size = (pointmap.shape[1], pointmap.shape[2]) + if pointmap_size != target_size: + # Handle NaN values in pointmap during resizing + # Direct resize would propagate NaN values, so we need special handling + nan_mask = torch.isnan(pointmap).any(dim=0) + pointmap_clean = torch.where(torch.isnan(pointmap), torch.zeros_like(pointmap), pointmap) + pointmap_resized = torchvision.transforms.functional.resize( + pointmap_clean, target_size, interpolation=torchvision.transforms.InterpolationMode.BILINEAR + ) + + # Resize the nan mask to identify which regions should remain invalid + nan_mask_resized = torchvision.transforms.functional.resize( + nan_mask.unsqueeze(0).float(), target_size, + interpolation=torchvision.transforms.InterpolationMode.NEAREST + ).squeeze(0) > 0.5 + + # Restore NaN values in regions that were originally invalid + pointmap = torch.where( + nan_mask_resized.unsqueeze(0).expand_as(pointmap_resized), + torch.full_like(pointmap_resized, float('nan')), + pointmap_resized + ) + + if squeeze_mask: + mask = mask.squeeze(0) + + if pointmap is not None: + return rgb_image, mask, pointmap + return rgb_image, mask + + +SSINormalizedPointmap = namedtuple("SSINormalizedPointmap", ["pointmap", "scale", "shift"]) +class SSIPointmapNormalizer: + + def normalize(self, pointmap: torch.Tensor, mask: torch.Tensor, + scale: Optional[torch.Tensor] = None, shift: Optional[torch.Tensor] = None, + ) -> SSINormalizedPointmap: + if scale is None or shift is None: + normalized_pointmap, scale, shift = normalize_pointmap_ssi(pointmap) + else: + assert scale.shape == (3,) and shift.shape == (3,), "scale and shift must be in (3,) format" + scale, shift = _sanitize_ssi_scale_and_shift(scale, shift) + normalized_pointmap = _apply_metric_to_ssi(pointmap, scale, shift) + return SSINormalizedPointmap(normalized_pointmap, scale, shift) + + def denormalize(self, pointmap: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: + pointmap = _apply_metric_to_ssi(pointmap, scale, shift, apply_inverse=True) + return pointmap + + + +class ObjectCentricSSI(SSIPointmapNormalizer): + def __init__(self, + use_scene_scale: bool = True, + quantile_drop_threshold: float = 0.1, + clip_beyond_scale: Optional[float] = None, + # scale_factor: float = 3.8076, # e^(1.337); empirical mean of R3+Artist train + scale_factor: float = 1.0, # e^(1.337); empirical mean of R3+Artist train + allow_scale_and_shift_override: bool = False, + raise_on_no_valid_points: bool = False, + ): + self.use_scene_scale = use_scene_scale + self.quantile_drop_threshold = quantile_drop_threshold + self.clip_beyond_scale = clip_beyond_scale + self.scale_factor = scale_factor + self.allow_scale_and_shift_override = allow_scale_and_shift_override + self.raise_on_no_valid_points = raise_on_no_valid_points + + def _compute_scale_and_shift(self, pointmap: torch.Tensor, mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + pointmap_size = (pointmap.shape[1], pointmap.shape[2]) + + + mask_resized = torchvision.transforms.functional.resize( + mask, pointmap_size, + interpolation=torchvision.transforms.InterpolationMode.NEAREST + ).squeeze(0) + + pointmap_flat = pointmap.reshape(3, -1) + # Get valid points from the mask + mask_bool = mask_resized.reshape(-1) > 0.5 + mask_points = pointmap_flat[:, mask_bool] + mask_points_finite = torch.isfinite(mask_points).all(dim=0) + mask_points_nonzero = mask_points.abs().sum(dim=0) > 1e-6 + valid_mask_points = mask_points[:, mask_points_finite & mask_points_nonzero] + + if valid_mask_points.numel() == 0: + if self.raise_on_no_valid_points: + raise ValueError(f"No valid points found in mask") + logger.warning(f"No valid points found in mask; setting scale to {self.scale_factor} and shift to 0") + return torch.ones_like(pointmap_flat[:,0]) * self.scale_factor, torch.zeros_like(pointmap_flat[:,0]) + + # Compute median for shift + shift = valid_mask_points.nanmedian(dim=-1).values + # logger.info(f"{pointmap.shape=} {mask_resized.shape=} {shift.shape=}") + + + if self.use_scene_scale == True: + # Normalize by the scene scale + points_centered = pointmap_flat - shift.unsqueeze(-1) + scene_valid = torch.isfinite(pointmap_flat).all(dim=0) & (pointmap_flat.abs().sum(dim=0) > 1e-6) + scene_points_centered = points_centered[:, scene_valid] + if scene_points_centered.numel() == 0: + logger.warning( + "No valid scene points found while computing SSI scene scale; " + "falling back to masked object points." + ) + scene_points_centered = valid_mask_points - shift.unsqueeze(-1) + max_dims = scene_points_centered.abs().max(dim=0).values + scale = max_dims.nanmedian(dim=-1).values + elif self.use_scene_scale == False: + # Normalize by the object scale + shifted_mask_points = valid_mask_points - shift.unsqueeze(-1) + norm = shifted_mask_points.norm(dim=0) + quantiles = torch.nanquantile(norm, + torch.tensor([self.quantile_drop_threshold, 1. - self.quantile_drop_threshold], + device=shifted_mask_points.device), + dim=-1) + scale = (quantiles[1] - quantiles[0]).max(dim=-1).values * 2.0 + elif self.use_scene_scale.upper() == "OBJECT_NORM_MEDIAN": + # Normalize by the object scale + shifted_mask_points = valid_mask_points - shift.unsqueeze(-1) + norm = shifted_mask_points.norm(dim=0) + scale = norm.nanmedian(dim=-1).values + else: + raise ValueError(f"Invalid use_scene_scale: {self.use_scene_scale}") + scale = scale.expand_as(shift) # per-dim scaling + scale = scale * self.scale_factor + return scale, shift + + def normalize(self, pointmap: torch.Tensor, mask: torch.Tensor, + scale: Optional[torch.Tensor] = None, shift: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # 1. resize mask to size of pointmap using nearest interpolation + # 2. get mask points: pointmap[mask > 0.5] + # 3. shift = mask_points.median() # xyz + # 4. scale = # filter. If no points, then + # logger.info(f"{pointmap.shape=} {mask.shape=}") + assert pointmap.shape[0] == 3, "pointmap must be in (3, H, W) format" + pointmap_size = (pointmap.shape[1], pointmap.shape[2]) + + _scale, _shift = self._compute_scale_and_shift(pointmap, mask) + if scale is not None and self.allow_scale_and_shift_override: + _scale = scale + if shift is not None and self.allow_scale_and_shift_override: + _shift = shift + return_scale, return_shift = _sanitize_ssi_scale_and_shift(_scale, _shift) + + # Apply normalization + pointmap_normalized = _apply_metric_to_ssi(pointmap, return_scale, return_shift) + + if self.clip_beyond_scale is not None and self.clip_beyond_scale > 0: + new_norm = pointmap_normalized.norm(dim=0) + pointmap_normalized = torch.where( + new_norm > self.clip_beyond_scale, + torch.full_like(pointmap_normalized, float('nan')), + pointmap_normalized + ) + + return SSINormalizedPointmap(pointmap_normalized, return_scale, return_shift) + + +class ObjectApparentSizeSSI(SSIPointmapNormalizer): + def __init__(self, + clip_beyond_scale: Optional[float] = None, + use_scene_scale: bool = True, + scale_factor: float = 1.0, # e^(1.337); empirical mean of R3+Artist train + ): + self.clip_beyond_scale = clip_beyond_scale + self.use_scene_scale = use_scene_scale + self.scale_factor = scale_factor + + def _get_scale_and_shift(self, pointmap: torch.Tensor, mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + pointmap_size = (pointmap.shape[1], pointmap.shape[2]) + pointmap_flat = pointmap.reshape(3, -1) + + if not self.use_scene_scale: + # Get valid points from the mask + mask_resized = torchvision.transforms.functional.resize( + mask, pointmap_size, + interpolation=torchvision.transforms.InterpolationMode.NEAREST + ).squeeze(0) + mask_bool = mask_resized.reshape(-1) > 0.5 + pointmap_flat = pointmap_flat[:, mask_bool] + + # Median z-distance + median_z = pointmap_flat[-1, ...].nanmedian().unsqueeze(0) + scale = median_z.expand(3) * self.scale_factor + shift = torch.zeros_like(scale) + # logger.info(f'median z = {median_z}') + return scale, shift + + def normalize(self, + pointmap: torch.Tensor, + mask: torch.Tensor, + scale: Optional[torch.Tensor] = None, + shift: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + assert pointmap.shape[0] == 3, "pointmap must be in (3, H, W) format" + pointmap_size = (pointmap.shape[1], pointmap.shape[2]) + + if scale is None or shift is None: + scale, shift = self._get_scale_and_shift(pointmap, mask) + else: + assert scale.shape == (3,) and shift.shape == (3,), "scale and shift must be in (3,) format" + scale, shift = _sanitize_ssi_scale_and_shift(scale, shift) + + # Apply normalization and clip + pointmap_normalized = _apply_metric_to_ssi(pointmap, scale, shift) + # logger.info(f"{pointmap_normalized.shape=}") + + if self.clip_beyond_scale is not None and self.clip_beyond_scale > 0: + pointmap_normalized = torch.where( + pointmap_normalized[-1, ...] > self.clip_beyond_scale, + torch.full_like(pointmap_normalized, float('nan')), + pointmap_normalized + ) + + # return pointmap_normalized, scale, shift + return SSINormalizedPointmap(pointmap_normalized, scale, shift) + + +class NormalizedDisparitySpaceSSI(SSIPointmapNormalizer): + def __init__(self, + clip_beyond_scale: Optional[float] = None, + use_scene_scale: bool = True, + log_disparity_shift: float = 0.0, + ): + self.clip_beyond_scale = clip_beyond_scale + self.use_scene_scale = use_scene_scale + self.point_remapper = PointRemapper(remap_type="exp_disparity") + self.log_disparity_shift = log_disparity_shift + + def normalize(self, pointmap: torch.Tensor, mask: torch.Tensor, + scale: Optional[torch.Tensor] = None, shift: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + assert pointmap.shape[0] == 3, "pointmap must be in (3, H, W) format" + + + disparity_space_pointmap = self.point_remapper.forward(pointmap.permute(1, 2, 0)).permute(2, 0, 1) + if scale is None or shift is None: + scale, shift = self._get_scale_and_shift(disparity_space_pointmap, mask) + else: + assert scale.shape == (3,) and shift.shape == (3,), "scale and shift must be in (3,) format" + scale, shift = _sanitize_ssi_scale_and_shift(scale, shift) + + # pointmap_normalized = pointmap.clone().detach() + pointmap_normalized = _apply_metric_to_ssi(disparity_space_pointmap, scale, shift) + # logger.info(f"{pointmap_normalized.shape=}") + + if self.clip_beyond_scale is not None and self.clip_beyond_scale > 0: + pointmap_normalized = torch.where( + pointmap_normalized[2, ...].abs() > self.clip_beyond_scale, + torch.full_like(pointmap_normalized, float('nan')), + pointmap_normalized + ) + + # return pointmap_normalized, scale, shift + return SSINormalizedPointmap(pointmap_normalized, scale, shift) + + def denormalize(self, pointmap: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: + pointmap = _apply_metric_to_ssi(pointmap, scale, shift, apply_inverse=True) + pointmap = self.point_remapper.inverse(pointmap.permute(1, 2, 0)).permute(2, 0, 1) + return pointmap + + def _get_scale_and_shift(self, pointmap: torch.Tensor, mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + pointmap_size = (pointmap.shape[1], pointmap.shape[2]) + mask_resized = torchvision.transforms.functional.resize( + mask, pointmap_size, + interpolation=torchvision.transforms.InterpolationMode.NEAREST + ).squeeze(0) + + pointmap_flat = pointmap.reshape(3, -1) + if self.use_scene_scale: + median_z = pointmap_flat[-1, ...].nanmedian().unsqueeze(0) + shift = torch.zeros_like(median_z.expand(3)) + shift[-1, ...] = median_z[0] + self.log_disparity_shift + else: + # Get valid points from the mask (shift, x/z, y/z, log(z)) + mask_bool = mask_resized.reshape(-1) > 0.5 + pointmap_flat = pointmap_flat[:, mask_bool] + shift = pointmap_flat.nanmedian(dim=-1).values + + scale = torch.ones_like(shift) + # logger.info(f'median z = {median_z}') + return scale, shift + +def normalize_pointmap_ssi(pointmap: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Normalize pointmap using Scale-Shift Invariant (SSI) normalization. + + Args: + pointmap: Pointmap tensor of shape (H, W, 3) or (3, H, W) + + Returns: + Tuple of (normalized_pointmap, scale, shift) + """ + from sam3d_objects.data.dataset.tdfy.pose_target import ScaleShiftInvariant + + # Convert to (H, W, 3) if needed for get_scale_and_shift + if pointmap.shape[0] == 3: + pointmap_hw3 = pointmap.permute(1, 2, 0) + original_format = 'chw' + else: + pointmap_hw3 = pointmap + original_format = 'hwc' + + # Get scale and shift using existing method + scale, shift = ScaleShiftInvariant.get_scale_and_shift(pointmap_hw3) + scale, shift = _sanitize_ssi_scale_and_shift(scale, shift) + + pointmap_normalized = _apply_metric_to_ssi(pointmap, scale, shift) + return pointmap_normalized, scale, shift + +def _sanitize_ssi_scale_and_shift( + scale: torch.Tensor, + shift: torch.Tensor, + *, + min_scale: float = 1e-6, +) -> tuple[torch.Tensor, torch.Tensor]: + scale = torch.as_tensor(scale).clone() + shift = torch.as_tensor(shift).clone() + + if scale.shape != shift.shape: + raise ValueError(f"Expected scale and shift to share shape, got {scale.shape} and {shift.shape}") + + invalid_shift = ~torch.isfinite(shift) + if invalid_shift.any(): + logger.warning( + "Non-finite SSI shift detected; replacing {} entries with 0.", + int(invalid_shift.sum().item()), + ) + shift = torch.where(invalid_shift, torch.zeros_like(shift), shift) + + invalid_scale = ~torch.isfinite(scale) + if invalid_scale.any(): + logger.warning( + "Non-finite SSI scale detected; replacing {} entries with 1 before clamping.", + int(invalid_scale.sum().item()), + ) + scale = torch.where(invalid_scale, torch.ones_like(scale), scale) + + near_zero_scale = scale.abs() < min_scale + if near_zero_scale.any(): + logger.warning( + "Near-zero SSI scale detected; clamping {} entries to {}.", + int(near_zero_scale.sum().item()), + min_scale, + ) + scale = scale.abs().clamp_min(min_scale) + + return scale, shift + +def _apply_metric_to_ssi(pointmap: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, apply_inverse: bool = False) -> torch.Tensor: + """ + Normalize pointmap using Scale-Shift Invariant (SSI) normalization. + + Args: + pointmap: Pointmap tensor of shape (H, W, 3) or (3, H, W) + + Returns: + Tuple of (normalized_pointmap, scale, shift) + """ + from sam3d_objects.data.dataset.tdfy.pose_target import ScaleShiftInvariant + + # Convert to (H, W, 3) if needed for get_scale_and_shift + if pointmap.shape[0] == 3: + pointmap_hw3 = pointmap.permute(1, 2, 0) + original_format = 'chw' + else: + pointmap_hw3 = pointmap + original_format = 'hwc' + + scale, shift = _sanitize_ssi_scale_and_shift(scale, shift) + + # Apply normalization + ssi_to_metric = ScaleShiftInvariant.ssi_to_metric(scale, shift) + metric_to_ssi = ssi_to_metric.inverse() + transform_to_apply = metric_to_ssi + + if apply_inverse: + transform_to_apply = ssi_to_metric + + pointmap_flat = pointmap_hw3.reshape(-1, 3) + pointmap_normalized = transform_to_apply.transform_points(pointmap_flat) + + # Reshape back to original format + if original_format == 'chw': + pointmap_normalized = pointmap_normalized.reshape(pointmap.shape[1], pointmap.shape[2], 3).permute(2, 0, 1) + else: + pointmap_normalized = pointmap_normalized.reshape(pointmap_hw3.shape) + + return pointmap_normalized + + +def perturb_mask_translation( + image: torch.Tensor, + mask: torch.Tensor, + max_px_delta: int = 5, +): + """ + Applies data augmentation to the mask by randomly translating the mask. + + Args: + image: (C, H, W) float32 [0, 1] tensor. + mask: (1, H, W) float32 [0, 1] tensor. + max_px_delta: The maximum number of pixels we will randomly shift by in each 2D direction. + """ + dx = random.randint(-max_px_delta, max_px_delta) + dy = random.randint(-max_px_delta, max_px_delta) + + mask = mask.squeeze(0) + mask = torch.roll(mask, shifts=(dy, dx), dims=(0, 1)) + + # Zero out wrapped regions + if dy > 0: + mask[:dy, :] = 0 + elif dy < 0: + mask[dy:, :] = 0 + if dx > 0: + mask[:, :dx] = 0 + elif dx < 0: + mask[:, dx:] = 0 + + mask = mask.unsqueeze(0) + return image, mask + + +def perturb_mask_boundary( + image: torch.Tensor, + mask: torch.Tensor, + kernel_range: tuple[int, int] = (2, 5), + p_erode: float = 0.1, + p_dilate: float = 0.8, + **kwargs, +): + """ + Applies data augmentation to the mask by randomly eroding or dilating the mask. + + Args: + image: (C, H, W) float32 [0, 1] tensor. + mask: (1, H, W) float32 [0, 1] tensor. + kernel_range: Range of kernel sizes to sample from. + p_erode: Probability of erosion. + p_dilate: Probability of dilation. + kwargs: Kwargs for the cv2 erode/dilate function. + """ + import cv2 + + C, H, W = image.shape + assert mask.shape == (1, H, W) + assert mask.dtype == torch.float32 + assert torch.all((mask == 0) | (mask == 1)), "Mask must be binary (0 or 1)" + + p_none = 1.0 - p_erode - p_dilate + assert 0 <= p_none <= 1, "Probabilities must sum to 1 and be valid." + + # Sample operation. + op = random.choices(["erode", "dilate", "none"], weights=[p_erode, p_dilate, p_none], k=1)[0] + + if op == "none": + pass + else: + # Sample kernel size + ksize = random.randint(*kernel_range) + kernel = np.ones((ksize, ksize), np.uint8) + + mask = mask.squeeze().cpu().numpy().astype(np.uint8) # (H, W) + + if op == "erode": + mask = cv2.erode(mask, kernel, **kwargs) + elif op == "dilate": + mask = cv2.dilate(mask, kernel, **kwargs) + else: + raise NotImplementedError + + mask = torch.from_numpy(mask).float()[None] # (1, H, W) + + return image, mask + + +def resolution_blur( + image: torch.Tensor, + mask: torch.Tensor, + scale_range=(0.05, 0.95), + interpolation_down=tv_transforms.InterpolationMode.BICUBIC, + interpolation_up=tv_transforms.InterpolationMode.BICUBIC, +): + """ + Blur the input image by applying upsample(downsample(x)). + + Args: + image (torch.Tensor): Image tensor of shape (C, H, W), float32, with values in [0, 1]. + mask (torch.Tensor): Mask tensor of shape (1, H, W), float32, with values in [0, 1]. The mask is returned unchanged. + scale_range: Tuple of (min_scale, max_scale) for downsampling. + interpolation_down: Interpolation mode for downsampling. + interpolation_up: Interpolation mode for upsampling. + """ + C, H, W = image.shape + scale = random.uniform(*scale_range) + new_H, new_W = max(1, int(H * scale)), max(1, int(W * scale)) + + # Downsample + image = TF.resize(image, size=[new_H, new_W], interpolation=interpolation_down) + + # Upsample back to original size + image = TF.resize(image, size=[H, W], interpolation=interpolation_up) + + return image, mask + + +def gaussian_blur( + image: torch.Tensor, + mask: torch.Tensor, + kernel_range: tuple[int, int] = (3, 15), + sigma_range: tuple[int, int] = (0.1, 4.0), +): + """ + Apply gaussian blur to the input image. + + Args: + image (torch.Tensor): Image tensor of shape (C, H, W), float32, with values in [0, 1]. + mask (torch.Tensor): Mask tensor of shape (1, H, W), float32, with values in [0, 1]. The mask is returned unchanged. + kernel_range (tuple): Range of odd kernel sizes to sample from for the Gaussian blur (min, max). + sigma_range (tuple): Range of sigma values (standard deviation) to sample from for the Gaussian kernel (min, max). + """ + kernel_size = random.choice([k for k in range(kernel_range[0], kernel_range[1]+1) if k % 2 == 1]) + sigma = random.uniform(*sigma_range) + pad = kernel_size // 2 + + # Step 1: Pad the image + image = F.pad(image.unsqueeze(0), (pad, pad, pad, pad), mode='replicate') + + # Step 2: Apply gaussian blur + image = TF.gaussian_blur(image, kernel_size=[kernel_size, kernel_size], sigma=sigma) + + # Step 3: Unpad to get back to original size + image = image[:, :, pad:-pad, pad:-pad] + + return image.squeeze(0), mask + + +def apply_blur_augmentation( + image: torch.Tensor, + mask: torch.Tensor, + p_resolution: float = 0.33, + p_gaussian: float = 0.33, + gaussian_kwargs: dict = None, + resolution_kwargs: dict = None, +): + """Apply blur augmentation with configurable parameters""" + + # Handle None defaults BEFORE unpacking + if gaussian_kwargs is None: + gaussian_kwargs = {} + if resolution_kwargs is None: + resolution_kwargs = {} + + p_none = 1.0 - p_gaussian - p_resolution + assert 0 <= p_none <= 1, "Probabilities must sum to 1 and be valid." + + operation = random.choices( + ["gaussian", "resolution", "none"], + weights=[p_gaussian, p_resolution, p_none], + k=1 + )[0] + + if operation == "gaussian": + return gaussian_blur(image, mask, **gaussian_kwargs) + elif operation == "resolution": + return resolution_blur(image, mask, **resolution_kwargs) + elif operation == "none": + return image, mask + else: + raise NotImplementedError diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/img_processing.py b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/img_processing.py new file mode 100644 index 0000000000000000000000000000000000000000..e4e7fa4b9c2c5f833c1ea68d5f21953e6e243485 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/img_processing.py @@ -0,0 +1,189 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import math + +import random + +import torch +import torch.nn.functional as F + +from torchvision import transforms +from torchvision.transforms import functional as tv_F + + +class RandomResizedCrop(transforms.RandomResizedCrop): + """ + RandomResizedCrop for matching TF/TPU implementation: no for-loop is used. + This may lead to results different with torchvision's version. + Following BYOL's TF code: + https://github.com/deepmind/deepmind-research/blob/master/byol/utils/dataset.py#L206 + """ + + @staticmethod + def get_params(img, scale, ratio): + width, height = tv_F._get_image_size(img) + area = height * width + + target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item() + log_ratio = torch.log(torch.tensor(ratio)) + aspect_ratio = torch.exp( + torch.empty(1).uniform_(log_ratio[0], log_ratio[1]) + ).item() + + w = int(round(math.sqrt(target_area * aspect_ratio))) + h = int(round(math.sqrt(target_area / aspect_ratio))) + + w = min(w, width) + h = min(h, height) + + i = torch.randint(0, height - h + 1, size=(1,)).item() + j = torch.randint(0, width - w + 1, size=(1,)).item() + + return i, j, h, w + + +# following PT3D CO3D data to pad image +def pad_to_square(image, value=0): + _, _, h, w = image.shape # Assuming image is in (B, C, H, W) format + if h == w: + return image # The image is already square + + # Calculate the padding + diff = abs(h - w) + pad2 = diff + + # Pad the image to make it square + if h > w: + padding = (0, pad2, 0, 0) # Pad width (left, right, top, bottom) + else: + padding = (0, 0, 0, pad2) # Pad height + # Apply padding + padded_image = torch.nn.functional.pad(image, padding, mode="constant", value=value) + return padded_image + + +def preprocess_img( + x, + mask=None, + img_target_shape=224, + mask_target_shape=256, + normalize=False, +): + if x.shape[1] != x.shape[2]: + x = pad_to_square(x) + if mask is not None and mask.shape[1] != mask.shape[2]: + mask = pad_to_square(mask) + if x.shape[2] != img_target_shape: + x = F.interpolate( + x, + size=(img_target_shape, img_target_shape), + # scale_factor=float(img_target_shape)/x.shape[2], + mode="bilinear", + ) + if mask is not None and mask.shape[2] != mask_target_shape: + if mask is not None: + mask = F.interpolate( + mask, + size=(mask_target_shape, mask_target_shape), + # scale_factor=float(mask_target_shape)/mask.shape[2], + mode="nearest", + ) + if normalize: + imgs_normed = resnet_img_normalization(x) + else: + imgs_normed = x + return imgs_normed, mask + + +def resnet_img_normalization(x): + resnet_mean = torch.tensor([0.485, 0.456, 0.406], device=x.device).reshape( + (3, 1, 1) + ) + resnet_std = torch.tensor([0.229, 0.224, 0.225], device=x.device).reshape((3, 1, 1)) + if x.ndim == 4: + resnet_mean = resnet_mean[None] + resnet_std = resnet_std[None] + x = (x - resnet_mean) / resnet_std + return x + + +# pad image to be centered for unprojecting depth +def pad_to_square_centered(image, value=0, pointmap=None): + h, w = image.shape[-2], image.shape[-1] # Assuming image is in (B, C, H, W) format + if h == w: + if pointmap is not None: + return image, pointmap + return image # The image is already square + + # Calculate the padding + diff = abs(h - w) + pad1 = diff // 2 + pad2 = diff - pad1 + + # Pad the image to make it square + if h > w: + padding = (pad1, pad2, 0, 0) # Pad width (left, right, top, bottom) + else: + padding = (0, 0, pad1, pad2) # Pad height + # Apply padding to image + padded_image = F.pad(image, padding, mode="constant", value=value) + + # Apply padding to pointmap if provided + if pointmap is not None: + # Pad pointmap using torch functional with NaN fill value + padded_pointmap = F.pad(pointmap, padding, mode="constant", value=float("nan")) + + return padded_image, padded_pointmap + return padded_image + + +def crop_img_to_obj(mask, context_size): + nonzeros = torch.nonzero(mask) + if len(nonzeros) > 0: + r_max, c_max = nonzeros.max(dim=0)[0] + r_min, c_min = nonzeros.min(dim=0)[0] + box_h = max(1, r_max - r_min) + box_w = max(1, c_max - c_min) + left = max(0, c_min - int(box_w * context_size)) + right = min(mask.shape[-1], c_max + int(box_w * context_size)) + top = max(0, r_min - int(box_h * context_size)) + bot = min(mask.shape[-2], r_max + int(box_h * context_size)) + return left, right, top, bot + return None, None, None, None + + +def random_pad(img, mask=None, max_ratio=0.0, pointmap=None): + max_size = int(max(img.shape) * max_ratio) + padding = tuple([random.randint(0, max_size) for _ in range(4)]) + img = F.pad(img, padding) + if mask is not None: + mask = F.pad(mask, padding) + + if pointmap is not None: + pointmap = F.pad(pointmap, padding, mode="constant", value=float("nan")) + return img, mask, pointmap + return img, mask + + +def get_img_color_augmentation( + color_jit_prob=0.5, + gaussian_blur_prob=0.1, +): + transform = transforms.Compose( + [ + # (a) Random Color Jitter + transforms.RandomApply( + [ + transforms.ColorJitter( + brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1 + ) + ], + p=color_jit_prob, + ), + # (b) Randomly apply GaussianBlur + transforms.RandomApply( + [transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0))], + p=gaussian_blur_prob, + ), + ] + ) + return transform diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/pose_target.py b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/pose_target.py new file mode 100644 index 0000000000000000000000000000000000000000..f6e221f25c4b437f415fb49725a9a5322df14c98 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/pose_target.py @@ -0,0 +1,784 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import torch +from typing import Dict, Optional, Tuple, Any +from dataclasses import dataclass, asdict, field +from loguru import logger + +from sam3d_objects.data.utils import expand_as_right, tree_tensor_map +from sam3d_objects.data.dataset.tdfy.transforms_3d import compose_transform, decompose_transform +from pytorch3d.transforms import Transform3d, quaternion_to_matrix, matrix_to_quaternion + + +@dataclass +class InstancePose: + """ + Stores the pose of an object. + Also, stores some information about the scene that was used to normalize the pose. + """ + + instance_scale_l2c: torch.Tensor + instance_position_l2c: torch.Tensor + instance_quaternion_l2c: torch.Tensor + scene_scale: torch.Tensor + scene_shift: torch.Tensor + + @classmethod + def _broadcast_postcompose( + cls, + scale: torch.Tensor, + rotation: torch.Tensor, + translation: torch.Tensor, + transform_to_postcompose: Transform3d, + ) -> Transform3d: + """ + Assumes scale, rotation, translation are of shape: + B, K, C + --- + B: batch size + K: number of objects + C: number of channels + + Takes a transform where + get_matrix() has shape (B, 3, 3) + + Returns pose.compose(transform_to_postcompose) + """ + scale_c = scale.shape[-1] + ndim_orig = scale.ndim + if ndim_orig == 3: + b, k, _ = scale.shape + elif ndim_orig == 2: + b = scale.shape[0] + k = 1 + elif ndim_orig == 1: + b = 1 + k = 1 + else: + raise ValueError(f"Invalid scale shape: {scale.shape}") + + # Create transform of shape (B * K) + wide = {"scale": scale, "rotation": rotation, "translation": translation} + shapes_orig = {k: v.shape for k, v in wide.items()} + long = tree_tensor_map(lambda x: x.reshape(b * k, x.shape[-1]), wide) + long["rotation"] = quaternion_to_matrix(long["rotation"]) + if scale_c == 1: + long["scale"] = long["scale"].expand(b * k, 3) + + composed = compose_transform(**long) + + # Apply transform to shape (B * K) + pc_transform = transform_to_postcompose.get_matrix() + pc_transform = pc_transform.repeat(k, 1, 1) + stacked_pc_transform = Transform3d(matrix=pc_transform) + assert stacked_pc_transform.get_matrix().shape == composed.get_matrix().shape + postcomposed = composed.compose(stacked_pc_transform) + + # Decompose transform to shape (B, K, C) + scale, rotation, translation = decompose_transform(postcomposed) + rotation = matrix_to_quaternion(rotation) + pc_long = {"scale": scale, "rotation": rotation, "translation": translation} + pc_wide = tree_tensor_map(lambda x: x.reshape(b, k, x.shape[-1]), pc_long) + if scale_c == 1: + pc_wide["scale"] = pc_wide["scale"][..., 0].unsqueeze(-1) + for k, shape in shapes_orig.items(): + pc_wide[k] = pc_wide[k].reshape(*shape) + return pc_wide["scale"], pc_wide["rotation"], pc_wide["translation"] + + +@dataclass +class PoseTarget: + x_instance_scale: torch.Tensor + x_instance_rotation: torch.Tensor + x_instance_translation: torch.Tensor + x_scene_scale: torch.Tensor + x_scene_center: torch.Tensor + x_translation_scale: torch.Tensor + pose_target_convention: str = field(default="unknown") + + +@dataclass +class InvariantPoseTarget: + """ + This is the canonical representation of pose targets, used for computing metrics. + instance_pose <-> invariant_pose_targets <-> all other pose_target_conventions + + Background: + --- + We want to estimate a transformation T: R³ → R³ despite scene scale ambiguity. + + The transformation taking object points to scene points is defined as + T(x) = s · R(q) · x + t + where: + - x is a point in the object coordinate frame, + - q is a unit quaternion representing rotation, + - s is the object-to-scene scale, and + - t is the translation. + + However, there is an inherent scale ambiguity in the scene, denoted as s_scene; + This ambiguity introduces irreducible error that complicates both evaluation and training. + + To decouple the scene scale from the invariant quantities, we define: + T(x) = s_scene · |t_rel| [ s_tilde · R(q) · x + t_unit ] + where we define + t_rel = t / s_scene + s_rel = s / s_scene + s_tilde = s_rel / |t_rel| + t_unit = t_rel / |t_rel| + + During training, you would predict (q, s_tilde, t_unit), leaving s_scene separate. + + + Hand-wavy error analysis: + --- + 1. Naive (coupled) estimate: + T(x) = s_scene [ s_rel · R(q) · x + t_rel ] + + We can define: + U = ln(s_rel) + V = ln(|t_rel|) + so that the error is governed by Var(U + V). + + 2. In the decoupled case, we have: + T(x) = s_scene · |t_rel| [ s_tilde · R(q) · x + t_unit ] + = s_scene · |t_rel| [ (s_rel / |t_rel|) R(q) · x + t_unit ] + Then ln(s_tilde) = ln(s_rel) - ln(|t_rel|) = U - V, and the error is + Var(U - V) = Var(U) + Var(V) - 2Cov(U, V). + + """ + + # These are invariant + q: torch.Tensor + t_unit: torch.Tensor + s_scene: torch.Tensor + t_scene_center: Optional[torch.Tensor] = None + t_rel_norm: Optional[torch.Tensor] = None + s_tilde: Optional[torch.Tensor] = None + s_rel: Optional[torch.Tensor] = None + + def __post_init__(self): + # Check that fields that are required always have values. + if self.q is None: + raise ValueError("Field 'q' (quaternion) must be provided.") + if self.s_scene is None: + raise ValueError("Field 's_scene' must be provided.") + if self.s_rel is None: + if self.s_tilde is not None: + self.s_rel = self.s_tilde * self.t_rel_norm + else: + raise ValueError("Field 's_rel' or 's_tilde' must be provided.") + if self.t_unit is None: + raise ValueError("Field 't_unit' must be provided.") + + if self.t_scene_center is None: + self.t_scene_center = torch.zeros_like(self.t_unit) + + # There is a simple relationship between s_tilde and t_rel_norm: + # s_tilde = s_rel / t_rel_norm + # + # If one of these is missing and the other is provided, we can compute the missing field. + if self.s_tilde is None and self.t_rel_norm is not None: + self.s_tilde = self.s_rel / self.t_rel_norm + elif self.t_rel_norm is None and self.s_tilde is not None: + self.t_rel_norm = self.s_rel / self.s_tilde + + # If both are provided, we check for consistency. + if self.s_tilde is not None and self.t_rel_norm is not None: + computed_s_tilde = self.s_rel / self.t_rel_norm + # If the provided s_tilde deviates from what is computed, update it. + if not torch.allclose(self.s_tilde, computed_s_tilde, atol=1e-6): + logger.warning( + f"s_tilde and t_rel_norm are provided, but they are not consistent. " + f"Updating s_tilde to {computed_s_tilde}." + ) + self.s_tilde = computed_s_tilde + + self._validate_fields() + + def _validate_fields(self): + for field in self.__dict__: + if self.__dict__[field] is None: + raise ValueError(f"Field '{field}' must be provided.") + + + @staticmethod + def from_instance_pose(instance_pose: InstancePose) -> "InvariantPoseTarget": + q = instance_pose.instance_quaternion_l2c + s_obj_to_scene = instance_pose.instance_scale_l2c # (..., 1) or (..., 3) + t_obj_to_scene = instance_pose.instance_position_l2c # (..., 3) + s_scene = instance_pose.scene_scale # (..., 1) or scalar-broadcastable + t_scene_center = instance_pose.scene_shift # (..., 3) + + # Normalize to scene scale (per the derivation) + if not ( s_obj_to_scene.ndim == (s_scene.ndim + 1)): + raise ValueError(f"s_scene should be ND [...,3] and s_obj_to_scene should be (N+1)D [...,K,3], but got {s_scene.shape=} {s_obj_to_scene.shape=}") + if not (t_obj_to_scene.ndim == (s_scene.ndim + 1)): + raise ValueError(f"t_scene_center should be ND [B,3] and t_obj_to_scene should be (N+1)D [B,K,3], but got {t_scene_center.shape=} {t_obj_to_scene.shape=}") + s_scene_exp = s_scene.unsqueeze(-2) + + s_rel = s_obj_to_scene / s_scene_exp + t_rel = t_obj_to_scene / s_scene_exp + + # Robust norms + eps = 1e-8 + t_rel_norm = t_rel.norm(dim=-1, keepdim=True).clamp_min(eps) + + s_tilde = s_rel / t_rel_norm + t_unit = t_rel / t_rel_norm + + return InvariantPoseTarget( + q=q, + s_scene=s_scene, + t_scene_center=t_scene_center, + s_rel=s_rel, + s_tilde=s_tilde, + t_unit=t_unit, + t_rel_norm=t_rel_norm, + ) + + + @staticmethod + def to_instance_pose(invariant_targets: "InvariantPoseTarget") -> InstancePose: + # scale factor per the derivation: s_scene * |t_rel| + # Normalize to scene scale (per the derivation) + t_rel_norm_ndim = invariant_targets.t_rel_norm.ndim + if not (invariant_targets.s_scene.ndim == (t_rel_norm_ndim - 1)) : + raise ValueError(f"s_scene should be ND [...,3] and t_rel_norm should be (N+1)D [...,K,3], but got {invariant_targets.s_scene.shape=} {invariant_targets.t_rel_norm.shape=}") + + scale = invariant_targets.s_scene.unsqueeze(-2) * invariant_targets.t_rel_norm + return InstancePose( + instance_scale_l2c=invariant_targets.s_tilde * scale, + instance_position_l2c=invariant_targets.t_unit * scale, + instance_quaternion_l2c=invariant_targets.q, + scene_scale=invariant_targets.s_scene, + scene_shift=invariant_targets.t_scene_center, + ) + + +class PoseTargetConvention: + """ + Converts pose_targets <-> instance_pose <-> invariant_pose_targets + """ + + pose_target_convention: str + + @classmethod + def from_invariant(cls, invariant_targets: InvariantPoseTarget) -> PoseTarget: + raise NotImplementedError("Implement this in a subclass") + + @classmethod + def to_invariant(cls, instance_pose: InstancePose) -> InvariantPoseTarget: + raise NotImplementedError("Implement this in a subclass") + + @classmethod + def from_instance_pose(cls, instance_pose: InstancePose) -> PoseTarget: + invariant_targets = InvariantPoseTarget.from_instance_pose(instance_pose) + return cls.from_invariant(invariant_targets) + + @classmethod + def to_instance_pose(cls, pose_target: PoseTarget) -> InstancePose: + invariant_targets = cls.to_invariant(pose_target) + return InvariantPoseTarget.to_instance_pose(invariant_targets) + + +class ScaleShiftInvariant(PoseTargetConvention): + """ + + Midas eq. (6): https://arxiv.org/pdf/1907.01341v3 + But for pointmaps (see MoGe): https://arxiv.org/pdf/2410.19115 + """ + + pose_target_convention: str = "ScaleShiftInvariant" + scale_mean = torch.tensor([1.0232692956924438, 1.0232691764831543, 1.0232692956924438]).to(torch.float32) + scale_std = torch.tensor([1.3773751258850098, 1.3773752450942993, 1.3773750066757202]).to(torch.float32) + translation_mean = torch.tensor([0.003191213821992278, 0.017236359417438507, 0.9401122331619263]).to(torch.float32) + translation_std = torch.tensor([1.341888666152954, 0.7665449380874634, 3.175130605697632]).to(torch.float32) + + @classmethod + def from_instance_pose(cls, instance_pose: InstancePose, normalize: bool = False) -> PoseTarget: + metric_to_ssi = cls.ssi_to_metric( + instance_pose.scene_scale, instance_pose.scene_shift + ).inverse() + + ssi_scale, ssi_rotation, ssi_translation = InstancePose._broadcast_postcompose( + scale=instance_pose.instance_scale_l2c, + rotation=instance_pose.instance_quaternion_l2c, + translation=instance_pose.instance_position_l2c, + transform_to_postcompose=metric_to_ssi, + ) + # logger.info(f"{normalize=} {ssi_scale.shape=} {ssi_rotation.shape=} {ssi_translation.shape=}") + if normalize: + device = ssi_scale.device + ssi_scale = (ssi_scale - cls.scale_mean.to(device)) / cls.scale_std.to(device) + ssi_translation = (ssi_translation - cls.translation_mean.to(device)) / cls.translation_std.to(device) + + return PoseTarget( + x_instance_scale=ssi_scale, + x_instance_rotation=ssi_rotation, + x_instance_translation=ssi_translation, + x_scene_scale=instance_pose.scene_scale, + x_scene_center=instance_pose.scene_shift, + x_translation_scale=torch.ones_like(ssi_scale)[..., 0].unsqueeze(-1), + pose_target_convention=cls.pose_target_convention, + ) + + @classmethod + def to_instance_pose(cls, pose_target: PoseTarget, normalize: bool = False) -> InstancePose: + scene_scale = pose_target.x_scene_scale + scene_shift = pose_target.x_scene_center + ssi_to_metric = cls.ssi_to_metric(scene_scale, scene_shift) + + if normalize: + device = pose_target.x_instance_scale.device + pose_target.x_instance_scale = pose_target.x_instance_scale * cls.scale_std.to(device) + cls.scale_mean.to(device) + pose_target.x_instance_translation = pose_target.x_instance_translation * cls.translation_std.to(device) + cls.translation_mean.to(device) + + ins_scale, ins_rotation, ins_translation = InstancePose._broadcast_postcompose( + scale=pose_target.x_instance_scale, + rotation=pose_target.x_instance_rotation, + translation=pose_target.x_instance_translation, + transform_to_postcompose=ssi_to_metric, + ) + + return InstancePose( + instance_scale_l2c=ins_scale, + instance_position_l2c=ins_translation, + instance_quaternion_l2c=ins_rotation, + scene_scale=scene_scale, + scene_shift=scene_shift, + ) + + @classmethod + def to_invariant(cls, pose_target: PoseTarget, normalize: bool = False) -> InvariantPoseTarget: + instance_pose = cls.to_instance_pose(pose_target, normalize=normalize) + return InvariantPoseTarget.from_instance_pose(instance_pose) + + @classmethod + def from_invariant(cls, invariant_targets: InvariantPoseTarget, normalize: bool = False) -> PoseTarget: + instance_pose = InvariantPoseTarget.to_instance_pose(invariant_targets) + return cls.from_instance_pose(instance_pose, normalize=normalize) + + @classmethod + def get_scale_and_shift(cls, pointmap): + shift_z = pointmap[..., -1].nanmedian().unsqueeze(0) + shift = torch.zeros_like(shift_z.expand(1, 3)) + shift[..., -1] = shift_z + + shifted_pointmap = pointmap - shift + scale = shifted_pointmap.abs().nanmean().to(shift.device) + + shift = shift.reshape(3) + scale = scale.expand(3) + + return scale, shift + + @staticmethod + def ssi_to_metric(scale: torch.Tensor, shift: torch.Tensor): + if scale.ndim == 1: + scale = scale.unsqueeze(0) + if shift.ndim == 1: + shift = shift.unsqueeze(0) + return Transform3d().scale(scale).translate(shift).to(shift.device) + + +class ScaleShiftInvariantWTranslationScale(PoseTargetConvention): + """ + + Midas eq. (6): https://arxiv.org/pdf/1907.01341v3 + But for pointmaps (see MoGe): https://arxiv.org/pdf/2410.19115 + """ + + pose_target_convention: str = "ScaleShiftInvariantWTranslationScale" + scale_mean = torch.tensor([1.0232692956924438, 1.0232691764831543, 1.0232692956924438]).to(torch.float32) + scale_std = torch.tensor([1.3773751258850098, 1.3773752450942993, 1.3773750066757202]).to(torch.float32) + translation_mean = torch.tensor([0.003191213821992278, 0.017236359417438507, 0.9401122331619263]).to(torch.float32) + translation_std = torch.tensor([1.341888666152954, 0.7665449380874634, 3.175130605697632]).to(torch.float32) + + @classmethod + def from_instance_pose(cls, instance_pose: InstancePose, normalize: bool = False) -> PoseTarget: + metric_to_ssi = cls.ssi_to_metric( + instance_pose.scene_scale, instance_pose.scene_shift + ).inverse() + + ssi_scale, ssi_rotation, ssi_translation = InstancePose._broadcast_postcompose( + scale=instance_pose.instance_scale_l2c, + rotation=instance_pose.instance_quaternion_l2c, + translation=instance_pose.instance_position_l2c, + transform_to_postcompose=metric_to_ssi, + ) + + ssi_translation_scale = ssi_translation.norm(dim=-1, keepdim=True) + ssi_translation_unit = ssi_translation / ssi_translation_scale.clamp_min(1e-7) + + return PoseTarget( + x_instance_scale=ssi_scale, + x_instance_rotation=ssi_rotation, + x_instance_translation=ssi_translation_unit, + x_scene_scale=instance_pose.scene_scale, + x_scene_center=instance_pose.scene_shift, + x_translation_scale=ssi_translation_scale, + pose_target_convention=cls.pose_target_convention, + ) + + @classmethod + def to_instance_pose(cls, pose_target: PoseTarget, normalize: bool = False) -> InstancePose: + scene_scale = pose_target.x_scene_scale + scene_shift = pose_target.x_scene_center + ssi_to_metric = cls.ssi_to_metric(scene_scale, scene_shift) + + ins_translation_unit = pose_target.x_instance_translation / pose_target.x_instance_translation.norm(dim=-1, keepdim=True) + ins_translation = ins_translation_unit * pose_target.x_translation_scale + + + ins_scale, ins_rotation, ins_translation = InstancePose._broadcast_postcompose( + scale=pose_target.x_instance_scale, + rotation=pose_target.x_instance_rotation, + translation=ins_translation, + transform_to_postcompose=ssi_to_metric, + ) + + + return InstancePose( + instance_scale_l2c=ins_scale, + instance_position_l2c=ins_translation, + instance_quaternion_l2c=ins_rotation, + scene_scale=scene_scale, + scene_shift=scene_shift, + ) + + @classmethod + def to_invariant(cls, pose_target: PoseTarget) -> InvariantPoseTarget: + instance_pose = cls.to_instance_pose(pose_target) + return InvariantPoseTarget.from_instance_pose(instance_pose) + + @classmethod + def from_invariant(cls, invariant_targets: InvariantPoseTarget) -> PoseTarget: + instance_pose = InvariantPoseTarget.to_instance_pose(invariant_targets) + return cls.from_instance_pose(instance_pose) + + @classmethod + def get_scale_and_shift(cls, pointmap): + shift_z = pointmap[..., -1].nanmedian().unsqueeze(0) + shift = torch.zeros_like(shift_z.expand(1, 3)) + shift[..., -1] = shift_z + + shifted_pointmap = pointmap - shift + scale = shifted_pointmap.abs().nanmean().to(shift.device) + + shift = shift.reshape(3) + scale = scale.expand(3) + + return scale, shift + + @staticmethod + def ssi_to_metric(scale: torch.Tensor, shift: torch.Tensor): + if scale.ndim == 1: + scale = scale.unsqueeze(0) + if shift.ndim == 1: + shift = shift.unsqueeze(0) + return Transform3d().scale(scale).translate(shift).to(shift.device) + + +class DisparitySpace(PoseTargetConvention): + pose_target_convention: str = "DisparitySpace" + + @classmethod + def from_instance_pose(cls, instance_pose: InstancePose, normalize: bool = False) -> PoseTarget: + + # x_instance_scale = orig_scale / scene_scale + # x_instance_translation = [x/z, y/z, 0] / scene_scale + # x_translation_scale = z / scene_scale + assert torch.allclose(instance_pose.scene_scale, torch.ones_like(instance_pose.scene_scale)) + + if not instance_pose.scene_shift.ndim == instance_pose.instance_position_l2c.ndim - 1: + raise ValueError(f"scene_shift must be (N+1)D and instance_position_l2c must be (N+1)D, but got {instance_pose.scene_shift.ndim} and {instance_pose.instance_position_l2c.ndim}") + shift_xy, shift_z_log = instance_pose.scene_shift.unsqueeze(-2).split([2, 1], dim=-1) + + + pose_xy, pose_z = instance_pose.instance_position_l2c.split([2, 1], dim=-1) + # Handle batch dimensions properly + if shift_xy.ndim < pose_xy.ndim: + shift_xy = shift_xy.unsqueeze(-2) + pose_xy_scaled = pose_xy / pose_z - shift_xy + + pose_z_scaled_log = torch.log(pose_z) - shift_z_log + x_instance_scale_log = torch.log(instance_pose.instance_scale_l2c) - torch.log(pose_z) + + x_instance_translation = torch.cat([pose_xy_scaled, torch.zeros_like(pose_z)], dim=-1) + x_translation_scale = torch.exp(pose_z_scaled_log) + x_instance_scale = torch.exp(x_instance_scale_log) + + + + return PoseTarget( + x_instance_scale=x_instance_scale, + x_instance_translation=x_instance_translation, + x_instance_rotation=instance_pose.instance_quaternion_l2c, + x_scene_scale=instance_pose.scene_scale, + x_scene_center=instance_pose.scene_shift, + x_translation_scale=x_translation_scale, + pose_target_convention=cls.pose_target_convention, + ) + + @classmethod + def to_instance_pose(cls, pose_target: PoseTarget, normalize: bool = False) -> InstancePose: + scene_scale = pose_target.x_scene_scale + scene_shift = pose_target.x_scene_center + + if not pose_target.x_scene_center.ndim == pose_target.x_instance_translation.ndim - 1: + raise ValueError(f"x_scene_center must be (N+1)D and x_instance_translation must be (N+1)D, but got {pose_target.x_scene_center.ndim} and {pose_target.x_instance_translation.ndim}") + shift_xy, shift_z_log = pose_target.x_scene_center.unsqueeze(-2).split([2, 1], dim=-1) + scene_z_scale = torch.exp(shift_z_log) + + z = pose_target.x_translation_scale + ins_translation = pose_target.x_instance_translation.clone() + ins_translation[...,2] = 1.0 + ins_translation[...,:2] = ins_translation[...,:2] + shift_xy + ins_translation = ins_translation * z * scene_z_scale + + ins_scale = pose_target.x_instance_scale * z * scene_z_scale + + return InstancePose( + instance_scale_l2c=ins_scale * scene_scale, + instance_position_l2c=ins_translation * scene_scale, + instance_quaternion_l2c=pose_target.x_instance_rotation, + scene_scale=scene_scale, + scene_shift=scene_shift, + ) + + @classmethod + def to_invariant(cls, pose_target: PoseTarget, normalize: bool = False) -> InvariantPoseTarget: + instance_pose = cls.to_instance_pose(pose_target, normalize=normalize) + return InvariantPoseTarget.from_instance_pose(instance_pose) + + @classmethod + def from_invariant(cls, invariant_targets: InvariantPoseTarget, normalize: bool = False) -> PoseTarget: + instance_pose = InvariantPoseTarget.to_instance_pose(invariant_targets) + return cls.from_instance_pose(instance_pose, normalize=normalize) + + + +class NormalizedSceneScale(PoseTargetConvention): + """ + x_instance_scale and x_translation_scale are normalized to x_scene_scale + """ + + pose_target_convention: str = "NormalizedSceneScale" + + @classmethod + def from_invariant(cls, invariant_targets: InvariantPoseTarget): + translation = invariant_targets.t_unit * invariant_targets.t_rel_norm + return PoseTarget( + x_instance_scale=invariant_targets.s_rel, + x_instance_rotation=invariant_targets.q, + x_instance_translation=translation, + x_scene_scale=invariant_targets.s_scene, + x_scene_center=invariant_targets.t_scene_center, + x_translation_scale=torch.ones_like(invariant_targets.t_rel_norm), + pose_target_convention=cls.pose_target_convention, + ) + + @classmethod + def to_invariant(cls, pose_target: PoseTarget): + t_rel_norm = torch.norm( + pose_target.x_instance_translation, dim=-1, keepdim=True + ) + return InvariantPoseTarget( + s_scene=pose_target.x_scene_scale, + s_rel=pose_target.x_instance_scale, + q=pose_target.x_instance_rotation, + t_unit=pose_target.x_instance_translation / t_rel_norm, + t_rel_norm=t_rel_norm, + t_scene_center=pose_target.x_scene_center, + ) + + +class Naive(PoseTargetConvention): + pose_target_convention: str = "Naive" + + @classmethod + def from_invariant(cls, invariant_targets: InvariantPoseTarget): + s_scene = invariant_targets.s_rel * invariant_targets.s_scene + t_scene = invariant_targets.t_unit * invariant_targets.t_rel_norm + return PoseTarget( + x_instance_scale=s_scene, + x_instance_rotation=invariant_targets.q, + x_instance_translation=t_scene, + x_scene_scale=invariant_targets.s_scene, + x_scene_center=invariant_targets.t_scene_center, + x_translation_scale=torch.ones_like(invariant_targets.t_rel_norm), + pose_target_convention=cls.pose_target_convention, + ) + + @classmethod + def to_invariant(cls, pose_target: PoseTarget): + s_scene = pose_target.x_scene_scale + t_rel_norm = torch.norm( + pose_target.x_instance_translation, dim=-1, keepdim=True + ) + return InvariantPoseTarget( + s_scene=s_scene, + t_scene_center=pose_target.x_scene_center, + s_rel=pose_target.x_instance_scale / s_scene, + q=pose_target.x_instance_rotation, + t_unit=pose_target.x_instance_translation / t_rel_norm, + t_rel_norm=t_rel_norm, + ) + + +class NormalizedSceneScaleAndTranslation(PoseTargetConvention): + """ + x_instance_scale and x_translation_scale are normalized to x_scene_scale + x_instance_translation is unit + """ + + pose_target_convention: str = "NormalizedSceneScaleAndTranslation" + + @classmethod + def from_invariant(cls, invariant_targets: InvariantPoseTarget): + return PoseTarget( + x_instance_scale=invariant_targets.s_rel, + x_instance_rotation=invariant_targets.q, + x_instance_translation=invariant_targets.t_unit, + x_scene_scale=invariant_targets.s_scene, + x_scene_center=invariant_targets.t_scene_center, + x_translation_scale=invariant_targets.t_rel_norm, + pose_target_convention=cls.pose_target_convention, + ) + + @classmethod + def to_invariant(cls, pose_target: PoseTarget): + return InvariantPoseTarget( + s_scene=pose_target.x_scene_scale, + t_scene_center=pose_target.x_scene_center, + s_rel=pose_target.x_instance_scale, + q=pose_target.x_instance_rotation, + t_unit=pose_target.x_instance_translation, + t_rel_norm=pose_target.x_translation_scale, + ) + + +class ApparentSize(PoseTargetConvention): + pose_target_convention: str = "ApparentSize" + + @classmethod + def from_invariant(cls, invariant_targets: InvariantPoseTarget): + return PoseTarget( + x_instance_scale=invariant_targets.s_tilde, + x_instance_rotation=invariant_targets.q, + x_instance_translation=invariant_targets.t_unit, + x_scene_scale=invariant_targets.s_scene, + x_scene_center=invariant_targets.t_scene_center, + x_translation_scale=invariant_targets.t_rel_norm, + pose_target_convention=cls.pose_target_convention, + ) + + @classmethod + def to_invariant(cls, pose_target: PoseTarget): + return InvariantPoseTarget( + s_scene=pose_target.x_scene_scale, + t_scene_center=pose_target.x_scene_center, + s_tilde=pose_target.x_instance_scale, + q=pose_target.x_instance_rotation, + t_unit=pose_target.x_instance_translation, + t_rel_norm=pose_target.x_translation_scale, + ) + + +class Identity(PoseTargetConvention): + """ + Identity convention - no transformation applied. + Direct passthrough mapping between instance pose and pose target values. + This preserves all values including scene_scale and scene_shift. + """ + + pose_target_convention: str = "Identity" + + @classmethod + def from_instance_pose(cls, instance_pose: InstancePose) -> PoseTarget: + return PoseTarget( + x_instance_scale=instance_pose.instance_scale_l2c, + x_instance_rotation=instance_pose.instance_quaternion_l2c, + x_instance_translation=instance_pose.instance_position_l2c, + x_scene_scale=instance_pose.scene_scale, + x_scene_center=instance_pose.scene_shift, + x_translation_scale=torch.ones_like(instance_pose.instance_scale_l2c)[..., 0].unsqueeze(-1), + pose_target_convention=cls.pose_target_convention, + ) + + @classmethod + def to_instance_pose(cls, pose_target: PoseTarget) -> InstancePose: + return InstancePose( + instance_scale_l2c=pose_target.x_instance_scale, + instance_position_l2c=pose_target.x_instance_translation, + instance_quaternion_l2c=pose_target.x_instance_rotation, + scene_scale=pose_target.x_scene_scale, + scene_shift=pose_target.x_scene_center, + ) + + @classmethod + def to_invariant(cls, pose_target: PoseTarget) -> InvariantPoseTarget: + instance_pose = cls.to_instance_pose(pose_target) + return InvariantPoseTarget.from_instance_pose(instance_pose) + + @classmethod + def from_invariant(cls, invariant_targets: InvariantPoseTarget) -> PoseTarget: + instance_pose = InvariantPoseTarget.to_instance_pose(invariant_targets) + return cls.from_instance_pose(instance_pose) + + +class PoseTargetConverter: + @staticmethod + def pose_target_to_instance_pose(pose_target: PoseTarget, normalize: bool = False) -> InstancePose: + _convention_class = globals()[pose_target.pose_target_convention] + if _convention_class == ScaleShiftInvariant: + return _convention_class.to_instance_pose(pose_target, normalize=normalize) + else: + return _convention_class.to_instance_pose(pose_target) + + @staticmethod + def instance_pose_to_pose_target( + instance_pose: InstancePose, pose_target_convention: str, normalize: bool = False + ) -> PoseTarget: + _convention_class = globals()[pose_target_convention] + if _convention_class == ScaleShiftInvariant: + return _convention_class.from_instance_pose(instance_pose, normalize=normalize) + else: + return _convention_class.from_instance_pose(instance_pose) + + @staticmethod + def dicts_instance_pose_to_pose_target( + pose_target_convention: str, + **kwargs, + ): + instance_pose = InstancePose(**kwargs) + pose_target = PoseTargetConverter.instance_pose_to_pose_target( + instance_pose, pose_target_convention + ) + return asdict(pose_target) + + @staticmethod + def dicts_pose_target_to_instance_pose( + **kwargs, + ): + pose_target_convention = kwargs.get("pose_target_convention") + _convention_class = globals()[pose_target_convention] + assert ( + _convention_class.pose_target_convention == pose_target_convention + ), f"Normalization name mismatch: {_convention_class.pose_target_convention} != {pose_target_convention}" + + normalize = kwargs.pop("normalize", False) + pose_target = PoseTarget(**kwargs) + instance_pose = PoseTargetConverter.pose_target_to_instance_pose(pose_target, normalize) + return asdict(instance_pose) + + +class LogScaleShiftNormalizer: + def __init__(self, shift_log: torch.Tensor = 0.0, scale_log: torch.Tensor = 1.0): + self.shift_log = shift_log + self.scale_log = scale_log + + def normalize(self, value: torch.Tensor): + return torch.log(value) - self.shift_log / self.scale_log + + def denormalize(self, value: torch.Tensor): + return torch.exp(value * self.scale_log + self.shift_log) \ No newline at end of file diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/preprocessor.py b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/preprocessor.py new file mode 100644 index 0000000000000000000000000000000000000000..672ab23649b35cd337b77db2f570bdedb94dc41f --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/preprocessor.py @@ -0,0 +1,212 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import warnings +import torch +from loguru import logger +from dataclasses import dataclass +from typing import Callable, Optional +import warnings + +from .img_and_mask_transforms import ( + SSIPointmapNormalizer, +) + + +# Load and process data +@dataclass +class PreProcessor: + """ + Preprocessor configuration for image, mask, and pointmap transforms. + + Transform application order: + 1. Pointmap normalization (if normalize_pointmap=True) + 2. Joint transforms (img_mask_pointmap_joint_transform or img_mask_joint_transform) + 3. Individual transforms (img_transform, mask_transform, pointmap_transform) + + For backward compatibility, img_mask_joint_transform is preserved. When both + img_mask_pointmap_joint_transform and img_mask_joint_transform are present, + img_mask_pointmap_joint_transform takes priority. + """ + + img_transform: Callable = (None,) + mask_transform: Callable = (None,) + img_mask_joint_transform: list[Callable] = (None,) + rgb_img_mask_joint_transform: list[Callable] = (None,) + + # New fields for pointmap support + pointmap_transform: Callable = (None,) + img_mask_pointmap_joint_transform: list[Callable] = (None,) + + # Pointmap normalization option + normalize_pointmap: bool = False + pointmap_normalizer: Optional[Callable] = None + rgb_pointmap_normalizer: Optional[Callable] = None + + def __post_init__(self): + if self.pointmap_normalizer is None: + self.pointmap_normalizer = SSIPointmapNormalizer() + if self.normalize_pointmap == False: + warnings.warn("normalize_pointmap is also set to False, which means we will return the moments but not normalize the pointmap. This supports old unnormalized pointmap models, but this is dangerous behavior.", DeprecationWarning, stacklevel=2) + + if self.rgb_pointmap_normalizer is None: + logger.warning("No rgb pointmap normalizer provided, using scale + shift ") + self.rgb_pointmap_normalizer = self.pointmap_normalizer + + + def _normalize_pointmap( + self, pointmap: torch.Tensor, + mask: torch.Tensor, + pointmap_normalizer: Callable, + scale: Optional[torch.Tensor] = None, + shift: Optional[torch.Tensor] = None, + ): + if pointmap is None: + return pointmap, None, None + + if self.normalize_pointmap == False: + # old behavior: Pose is normalized to the pointmap center, but pointmap is not + _, pointmap_scale, pointmap_shift = pointmap_normalizer.normalize(pointmap, mask) + return pointmap, pointmap_scale, pointmap_shift + + if scale is not None or shift is not None: + return pointmap_normalizer.normalize(pointmap, mask, scale, shift) + + return pointmap_normalizer.normalize(pointmap, mask) + + def _process_image_mask_pointmap_mess( + self, rgb_image, rgb_image_mask, pointmap=None + ): + """Extended version that handles pointmaps""" + # Keep a pointmap-aligned copy of the full-image mask. Spatial image/mask + # transforms such as pad+resize are applied later to the output tensors, + # but normalization must use a mask that is still aligned with the + # untransformed pointmap values. + rgb_pointmap_mask = rgb_image_mask.clone() if rgb_image_mask is not None else None + + # Apply pointmap normalization if enabled + pointmap_for_crop, pointmap_scale, pointmap_shift = self._normalize_pointmap( + pointmap, rgb_image_mask, self.pointmap_normalizer + ) + + # Apply transforms to the original full rgb image and mask. + rgb_image, rgb_image_mask = self._preprocess_rgb_image_mask(rgb_image, rgb_image_mask) + + # These two are typically used for getting cropped images of the object + # : first apply joint transforms + processed_rgb_image, processed_mask, processed_pointmap = ( + self._preprocess_image_mask_pointmap(rgb_image, rgb_image_mask, pointmap_for_crop) + ) + # : then apply individual transforms on top of the joint transforms + processed_rgb_image = self._apply_transform( + processed_rgb_image, self.img_transform + ) + processed_mask = self._apply_transform(processed_mask, self.mask_transform) + if processed_pointmap is not None: + processed_pointmap = self._apply_transform( + processed_pointmap, self.pointmap_transform + ) + + # This version is typically the full version of the image + # : apply individual transforms only + rgb_image = self._apply_transform(rgb_image, self.img_transform) + rgb_image_mask = self._apply_transform(rgb_image_mask, self.mask_transform) + + rgb_pointmap, rgb_pointmap_scale, rgb_pointmap_shift = self._normalize_pointmap( + pointmap, + rgb_pointmap_mask, + self.rgb_pointmap_normalizer, + pointmap_scale, + pointmap_shift, + ) + + if rgb_pointmap is not None: + rgb_pointmap = self._apply_transform(rgb_pointmap, self.pointmap_transform) + + result = { + "mask": processed_mask, + "image": processed_rgb_image, + "rgb_image": rgb_image, + "rgb_image_mask": rgb_image_mask, + } + + # Add pointmap results if available + if processed_pointmap is not None: + result.update( + { + "pointmap": processed_pointmap, + "rgb_pointmap": rgb_pointmap, + } + ) + + # Add normalization parameters if normalization was applied + if pointmap_scale is not None and pointmap_shift is not None: + result.update( + { + "pointmap_scale": pointmap_scale, + "pointmap_shift": pointmap_shift, + "rgb_pointmap_scale": rgb_pointmap_scale, + "rgb_pointmap_shift": rgb_pointmap_shift, + } + ) + + return result + + def _process_image_and_mask_mess(self, rgb_image, rgb_image_mask): + """Original method - calls extended version without pointmap""" + return self._process_image_mask_pointmap_mess(rgb_image, rgb_image_mask, None) + + def _preprocess_rgb_image_mask(self, rgb_image: torch.Tensor, rgb_image_mask: torch.Tensor): + """Apply joint transforms to rgb_image and rgb_image_mask.""" + if ( + self.rgb_img_mask_joint_transform != (None,) + and self.rgb_img_mask_joint_transform is not None + ): + for trans in self.rgb_img_mask_joint_transform: + rgb_image, rgb_image_mask = trans(rgb_image, rgb_image_mask) + return rgb_image, rgb_image_mask + + def _preprocess_image_mask_pointmap(self, rgb_image, mask_image, pointmap=None): + """Apply joint transforms with priority: triple transforms > dual transforms.""" + # Priority: img_mask_pointmap_joint_transform when pointmap is provided + if ( + self.img_mask_pointmap_joint_transform != (None,) + and self.img_mask_pointmap_joint_transform is not None + and pointmap is not None + ): + for trans in self.img_mask_pointmap_joint_transform: + rgb_image, mask_image, pointmap = trans( + rgb_image, mask_image, pointmap=pointmap + ) + return rgb_image, mask_image, pointmap + + # Fallback: img_mask_joint_transform (existing behavior) + elif ( + self.img_mask_joint_transform != (None,) + and self.img_mask_joint_transform is not None + ): + for trans in self.img_mask_joint_transform: + rgb_image, mask_image = trans(rgb_image, mask_image) + return rgb_image, mask_image, pointmap + + return rgb_image, mask_image, pointmap + + def _preprocess_image_and_mask(self, rgb_image, mask_image): + """Backward compatibility wrapper - only applies dual transforms""" + rgb_image, mask_image, _ = self._preprocess_image_mask_pointmap( + rgb_image, mask_image, None + ) + return rgb_image, mask_image + + # keep here for backward compatibility + def _preprocess_image_and_mask_inference(self, rgb_image, mask_image): + warnings.warn( + "The _preprocess_image_and_mask_inference is deprecated! Please use _preprocess_image_and_mask", + category=DeprecationWarning, + stacklevel=2, + ) + return self._preprocess_image_and_mask(rgb_image, mask_image) + + def _apply_transform(self, input: torch.Tensor, transform): + if input is not None and transform is not None and transform != (None,): + input = transform(input) + + return input diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/transforms_3d.py b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/transforms_3d.py new file mode 100644 index 0000000000000000000000000000000000000000..e9f417c2be8ec6259ccf2f403adb3bad43c54c87 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/data/dataset/tdfy/transforms_3d.py @@ -0,0 +1,50 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from collections import namedtuple +import math +import torch + +from pytorch3d.transforms import ( + Rotate, + Translate, + Scale, + Transform3d, + quaternion_to_matrix, + axis_angle_to_quaternion, +) + +DecomposedTransform = namedtuple( + "DecomposedTransform", ["scale", "rotation", "translation"] +) + + +def compose_transform( + scale: torch.Tensor, rotation: torch.Tensor, translation: torch.Tensor +) -> Transform3d: + """ + Args: + scale: (..., 3) tensor of scale factors + rotation: (..., 3, 3) tensor of rotation matrices + translation: (..., 3) tensor of translation vectors + """ + tfm = Transform3d(dtype=scale.dtype, device=scale.device) + return tfm.scale(scale).rotate(rotation).translate(translation) + + +def decompose_transform(transform: Transform3d) -> DecomposedTransform: + """ + Returns: + scale: (..., 3) tensor of scale factors + rotation: (..., 3, 3) tensor of rotation matrices + translation: (..., 3) tensor of translation vectors + """ + matrices = transform.get_matrix() + scale = torch.norm(matrices[:, :3, :3], dim=-1) + rotation = matrices[:, :3, :3] / scale.unsqueeze(-1) # Normalize rotation matrix + translation = matrices[:, 3, :3] # Extract translation vector + return DecomposedTransform(scale, rotation, translation) + + +def get_rotation_about_x_axis(angle: float = math.pi / 2) -> torch.Tensor: + axis = torch.tensor([1.0, 0.0, 0.0]) + axis_angle = axis * angle + return axis_angle_to_quaternion(axis_angle) diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/data/utils.py b/migrator/code/sam3d-for-6d/sam3d_objects/data/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e990bafbd803b292124f18ddb5c9cb07e06fdbcb --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/data/utils.py @@ -0,0 +1,242 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +from typing import Any, Iterable, Tuple, Union, Dict, Sequence, Mapping, Container +import optree +import torch +from collections.abc import Iterable +import inspect +import ast +import astor +from torch.utils import _pytree + +# None = root, Iterable[Any] = path, Any = path of one +ChildPathType = Union[None, Iterable[Any], Any] +ArgsType = Iterable[ChildPathType] +KwargsType = Mapping[str, ChildPathType] +ArgsKwargsType = Tuple[ArgsType, KwargsType] +MappingType = Union[None, ArgsKwargsType, ArgsType, KwargsType] + + +def tree_transpose_level_one( + structure, + check_children=False, + map_fn=None, + is_leaf=None, +): + _, outer_spec = optree.tree_flatten( + structure, + is_leaf=lambda x: x is not structure, + none_is_leaf=True, + ) + + spec = optree.tree_structure(structure, none_is_leaf=True, is_leaf=is_leaf) + children_spec = spec.children() + if len(children_spec) > 0: + inner_spec = children_spec[0] + if check_children: + for child_spec in children_spec[1:]: + assert ( + inner_spec == child_spec + ), f"one child was found having a different tree structure ({inner_spec} != {child_spec})" + + structure = optree.tree_transpose(outer_spec, inner_spec, structure) + + if map_fn is not None: + structure = optree.tree_map( + map_fn, + structure, + is_leaf=lambda x: optree.tree_structure( + x, is_leaf=is_leaf, none_is_leaf=True + ) + == outer_spec, + none_is_leaf=True, + ) + + return structure + + +def tree_tensor_map(fn, tree, *rest): + return optree.tree_map( + fn, + tree, + *rest, + is_leaf=lambda x: isinstance(x, torch.Tensor), + none_is_leaf=False, + ) + + +def to_device(obj, device): + """Recursively moves all tensors in obj to the specified device. + + Args: + obj: Object to move to device - can be a tensor, list, tuple, dict or any nested combination + device: Target device (e.g. 'cuda', 'cpu', torch.device('cuda:0') etc.) + + Returns: + Same object structure with all contained tensors moved to specified device + """ + to_fn = lambda x: x.to(device) + return optree.tree_map(to_fn, obj, is_leaf=torch.is_tensor, none_is_leaf=False) + + +def expand_right(tensor, target_shape): + """ + e.g. Takes tensor of (a, b, c) and returns a tensor of (a, b, c, 1, 1, ...) + """ + current_shape = tensor.shape + dims_to_add = len(target_shape) - len(current_shape) + result = tensor + for _ in range(dims_to_add): + result = result.unsqueeze(-1) + expand_shape = list(current_shape) + [-1] * dims_to_add + for i in range(len(target_shape)): + if i < len(expand_shape) and expand_shape[i] == -1: + expand_shape[i] = target_shape[i] + return result.expand(*expand_shape) + + +def expand_as_right(tensor, target): + return expand_right(tensor, target.shape) + + +def as_keys(path: ChildPathType): + if isinstance(path, Iterable) and (not isinstance(path, str)): + return tuple(path) + elif path is None: + return () + return (path,) + + +def get_child(obj: Any, *keys: Iterable[Any]): + for key in keys: + obj = obj[key] + return obj + + +def set_child(obj: Any, value: Any, *keys: Iterable[Any]): + parent = None + for key in keys: + parent = obj + obj = obj[key] + if parent is None: + obj = value + else: + parent[key] = value + return obj + + +def build_args_batch_extractor(args_mapping: ArgsType): + def extract_fn(batch): + return tuple(get_child(batch, *as_keys(path)) for path in args_mapping) + + return extract_fn + + +def build_kwargs_batch_extractor(kwargs_mapping: KwargsType): + def extract_fn(batch): + return { + name: get_child(batch, *as_keys(path)) + for name, path in kwargs_mapping.items() + } + + return extract_fn + + +empty_mapping = object() +kwargs_identity_mapping = object() + + +def build_batch_extractor(mapping: MappingType): + extract_args_fn = lambda x: () + extract_kwargs_fn = lambda x: {} + + if mapping is None: + + def extract_args_fn(batch): + return (batch,) + + elif mapping is empty_mapping: + pass + elif mapping is kwargs_identity_mapping: + extract_kwargs_fn = lambda x: x + elif isinstance(mapping, Sequence) and (not isinstance(mapping, str)): + if ( + len(mapping) == 2 + and isinstance(mapping[0], Sequence) + and isinstance(mapping[1], Dict) + ): + extract_args_fn = build_args_batch_extractor(mapping[0]) + extract_kwargs_fn = build_kwargs_batch_extractor(mapping[1]) + else: + extract_args_fn = build_args_batch_extractor(mapping) + elif isinstance(mapping, Mapping): + extract_kwargs_fn = build_kwargs_batch_extractor(mapping) + else: + + def extract_args_fn(batch): + return (get_child(batch, *as_keys(mapping)),) + + def extract_fn(batch): + return extract_args_fn(batch), extract_kwargs_fn(batch) + + return extract_fn + + +# > + + +def right_broadcasting(arr, target): + return arr.reshape(arr.shape + (1,) * (target.ndim - arr.ndim)) + + +def get_stats(tensor: torch.Tensor): + float_tensor = tensor.float() + return { + "shape": tuple(tensor.shape), + "min": tensor.min().item(), + "max": tensor.max().item(), + "mean": float_tensor.mean().item(), + "median": tensor.median().item(), + "std": float_tensor.std().item(), + } + + +def _get_caller_arg_name(argnum=0, parent_frame=1): + try: + frame = inspect.currentframe() # current frame + frame = inspect.getouterframes(frame)[1 + parent_frame] # parent frame + code = inspect.getframeinfo(frame[0]).code_context[0].strip() # get code line + + tree = ast.parse(code) + + for node in ast.walk(tree): + if isinstance(node, ast.Call): + args = node.args + break # only get the first parent call + + # get first argument string (do not handle '=') + label = astor.to_source(args[argnum]).strip() + except: + # TODO(Pierre) log exception + label = "{label}" + return label + + +def print_stats(tensor, label=None): + if label is None: + label = _get_caller_arg_name(argnum=0) + stats = get_stats(tensor) + string = f"{label}:\n" + "\n".join(f"- {k}: {v}" for k, v in stats.items()) + print(string) + + +def tree_reduce_unique(fn, tree, ensure_unique=True, **kwargs): + values = _pytree.tree_flatten(tree, **kwargs)[0] + values = tuple(map(fn, values)) + first = values[0] + if ensure_unique: + for value in values[1:]: + if value != first: + raise RuntimeError( + f"different values found, {value} and {first} should be the same" + ) + return first diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/checkpointing.py b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/checkpointing.py new file mode 100644 index 0000000000000000000000000000000000000000..e0695ebe3a3b3e68858b645a498dfce96b9b623a --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/checkpointing.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch + + + +def unwrap_model(model: torch.nn.Module) -> torch.nn.Module: + return model.module if hasattr(model, "module") else model + + + +def trainable_state_dict(model: torch.nn.Module) -> dict[str, torch.Tensor]: + m = unwrap_model(model) + required = {name for name, p in m.named_parameters() if p.requires_grad} + full = m.state_dict() + return {k: v for k, v in full.items() if k in required} + + + +def save_checkpoint( + checkpoint_path: str, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + scheduler, + scaler, + step: int, + config_dict: dict[str, Any], + history: dict[str, Any] | None = None, +) -> None: + checkpoint_path = str(Path(checkpoint_path).resolve()) + Path(checkpoint_path).parent.mkdir(parents=True, exist_ok=True) + payload = { + "step": int(step), + "trainable_state_dict": trainable_state_dict(model), + "optimizer": optimizer.state_dict(), + "scheduler": scheduler.state_dict() if scheduler is not None else None, + "scaler": scaler.state_dict() if scaler is not None else None, + "config": config_dict, + "history": history if history is not None else {}, + } + torch.save(payload, checkpoint_path) + + + +def load_checkpoint( + checkpoint_path: str, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer | None = None, + scheduler=None, + scaler=None, + map_location: str = "cpu", +) -> tuple[int, dict[str, Any]]: + checkpoint_path = str(Path(checkpoint_path).resolve()) + ckpt = torch.load(checkpoint_path, map_location=map_location, weights_only=False) + m = unwrap_model(model) + m.load_state_dict(ckpt["trainable_state_dict"], strict=False) + + if optimizer is not None and ckpt.get("optimizer") is not None: + optimizer.load_state_dict(ckpt["optimizer"]) + if scheduler is not None and ckpt.get("scheduler") is not None: + scheduler.load_state_dict(ckpt["scheduler"]) + if scaler is not None and ckpt.get("scaler") is not None: + scaler.load_state_dict(ckpt["scaler"]) + return int(ckpt.get("step", 0)), ckpt.get("history", {}) diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/code_outline.txt b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/code_outline.txt new file mode 100644 index 0000000000000000000000000000000000000000..b66f735169e9a1dbf5a57f8d92c9c327f89b027c --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/code_outline.txt @@ -0,0 +1,12 @@ +CFG는 이렇게 작동함 + + +reverse_fn(z, t, cond, mode="first_frame_inject") + → CFG.forward(**kwargs_cond) # picks up mode from kwargs_cond + → _manage_inject_cache(backbone, kwargs_cond.get("mode")) ✓ + → backbone(x, t, *args, **kwargs) # mode in kwargs + → SparseStructureFlowTdfyWrapper.forward(mode=...) ✓ + → SparseStructureFlowModel.forward(mode=...) ✓ + → block(h, t_emb, cond, mode=mode) × 24 ✓ + → self.self_attn(h, mode=mode) ✓ + → mm_scale_dot_product_attention(q, k, v, mode) ✓ diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/config.py b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/config.py new file mode 100644 index 0000000000000000000000000000000000000000..02d3f146b52e1e8a333632eac1b9634fa4af6d98 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/config.py @@ -0,0 +1,273 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import List + +from omegaconf import OmegaConf + + +@dataclass +class ExperimentConfig: + name: str = "stage1_flow_lora_ho3d" + output_dir: str = "outputs/stage1_flow_lora_ho3d" + seed: int = 42 + resume_from: str = "" + + +@dataclass +class DataConfig: + # ── Dataset selector ───────────────────────────────────────────── + dataset_name: str = "ho3d" # "ho3d" or "dexycb" + # ── HO3D fields ────────────────────────────────────────────────── + ho3d_root: str = "/scratch/x3302a04/workspace/6d_pose_estimation/6data/data/ho3d_full" + split_file: str = "train.txt" + mask_mode: str = "blue_channel" + train_sequences: List[str] = field(default_factory=list) + val_sequences: List[str] = field(default_factory=lambda: ["MC1", "SMu1", "SS1"]) + # ── DexYCB fields ──────────────────────────────────────────────── + dexycb_root: str = "" + target_ycb_id: int = 13 + target_ycb_ids: List[int] = field(default_factory=list) + pointmap_dir: str = "" + # DexYCB pointmap on-disk format consumed by dataset_unified: + # "dense" (default): ///raw/.pt (N_cams,H,W,3) + # "compact" : ///raw/.npz (sparse per-view, + # tools/compact_dexycb_pointmaps.py). Bit-exact, ~38x smaller. + pointmap_format: str = "dense" + # manifest-driven DexYCB path (dataset_name="dexycb_jsonl"): per-frame GT for + # moving-object training. JSONL rows: subject/session/serial/ycb_id/obj_index/frame. + dexycb_manifest_jsonl: str = "" + # held-out s0 TEST subset used by SAM-3D-inference validation (relative R/t metric). + dexycb_val_manifest_jsonl: str = "" + train_subjects: List[str] = field(default_factory=list) + val_subjects: List[str] = field(default_factory=list) + subjects: List[str] = field(default_factory=list) # optional subject filter for unified builder + # ── RoboCasa / unified fields (dataset_unified.build_manifest_robocasa) ── + robocasa_pm_root: str = "" + robocasa_mv_root: str = "" + robocasa_dataset_root: str = "" + robocasa_task: str = "" + # ── Unified sampling ────────────────────────────────────────────── + frame_stride: int = 1 + # Validation split is SEQUENCE-LEVEL (whole sequences held out; see + # dataset_unified.build_train_val_datasets). val_fraction is the fraction of + # SEQUENCES to hold out; val_num_sequences (>0) overrides it with an absolute + # count; val_stride is the relative-pose pair stride used both to pick + # pair-able val sequences and by the relative-pose validation metric. + val_fraction: float = 0.05 + val_num_sequences: int = 0 + val_stride: int = 10 + # ── Shared fields ───────────────────────────────────────────────── + initial_pose_file: str = "" + min_keyframes: int = 3 + min_masked_pointmap_scale: float = 0.0 + prefilter_query_pointmap_validity: bool = False + max_initial_pose_translation_norm_m: float = 10.0 + # pairs_per_epoch: int = 40000 + # val_pairs_per_epoch: int = 2000 + num_workers: int = 8 + pin_memory: bool = True + persistent_workers: bool = True + + +@dataclass +class ModelConfig: + pipeline_yaml: str = "checkpoints/hf/pipeline.yaml" + # When use_lora=False, the layout-transformer parameter subset (same filter as + # the LoRA targets) is unfrozen and fully fine-tuned instead of injecting LoRA + # adapters. gradient_checkpointing trades compute for activation memory on the + # backbone blocks (needed for full fine-tune to fit a usable batch size). + use_lora: bool = True + gradient_checkpointing: bool = False + lora_r: int = 16 + lora_alpha: int = 32 + lora_dropout: float = 0.05 + target_module_tokens: List[str] = field( + default_factory=lambda: ["to_qkv", "to_q", "to_kv", "to_out", "mlp.0", "mlp.2"] + ) + target_name_must_start_with: str = "blocks." + target_name_must_contain: List[str] = field(default_factory=lambda: [".6drotation_normalized."]) + target_name_must_not_contain: List[str] = field( + default_factory=lambda: [".shape.", "latent_mapping", "d_embedder", "t_embedder", "adaLN_modulation"] + ) + strict_layout_only_filter: bool = True + pointmap_embedder_index: int = 2 + ss_cfg_strength: float = 7.0 + ss_cfg_interval: List[int] = field(default_factory=lambda: [0, 500]) + ss_cfg_strength_pm: float = 0.0 + ss_rescale_t: int = 3 + ss_inference_steps: int = 25 + flow_mode: str = "default" + batch_conditioning: bool = False + + +@dataclass +class TrainConfig: + batch_size_per_gpu: int = 4 + grad_accum_steps: int = 4 + max_steps: int = 30000 + lr: float = 1.0e-4 + weight_decay: float = 1.0e-2 + betas: List[float] = field(default_factory=lambda: [0.9, 0.95]) + max_grad_norm: float = 1.0 + amp_dtype: str = "bf16" + scheduler: str = "cosine" + warmup_steps: int = 1000 + min_lr: float = 1.0e-6 + log_every: int = 20 + val_every: int = 2000 + save_every: int = 3000 + # frame stride for the relative R/t metric in SAM-3D-inference validation + eval_rel_stride: int = 5 + # query-frame batch size for the SAM-3D-inference validation decode (shape K/V is + # primed at this batch too); lower it if validation OOMs alongside the training job. + eval_infer_batch: int = 32 + # distilled shortcut solver steps for validation (matches production inference) + eval_shortcut_steps: int = 2 + expected_global_batch_size: int = 32 + timestep_sampler: str = "uniform" + timestep_lognorm_mean: float = 0.0 + timestep_lognorm_std: float = 1.0 + cfg_prob_uncond: float = 0.1 + cfg_prob_pm_drop: float = 0.1 + forward_mode: str = "2pass" + eval_seed: int = 12345 + + +@dataclass +class LossConfig: + rotation_weight: float = 1.0 + translation_weight: float = 1.0 + scale_weight: float = 1.0 + min_valid_pointmap_scale: float = 1.0e-6 + max_valid_gt_translation_norm_m: float = 10.0 + skip_invalid_translation_supervision: bool = True + # Scale supervision: GT scale latent = log(gt_scale_metric / pm_scale) (SSI + # space, mirrors translation). Samples with gt_scale outside this metric + # longest-axis-extent range (meters) are excluded from the scale loss. + min_valid_gt_scale_m: float = 1.0e-4 + max_valid_gt_scale_m: float = 10.0 + + +@dataclass +class DistributedConfig: + backend: str = "ddp" + deepspeed_config: str = "configs/deepspeed/zero2_stage1_lora.json" + + +@dataclass +class WandbConfig: + enable_wandb: bool = False + # Backward-compatible typo alias requested by some launch configs. + enable_wanb: bool = False + project: str = "sam3d-objects" + entity: str = "" + run_name: str = "" + group: str = "" + job_type: str = "train" + notes: str = "" + mode: str = "online" + tags: List[str] = field(default_factory=list) + + +@dataclass +class TrainStage1FlowConfig: + experiment: ExperimentConfig = field(default_factory=ExperimentConfig) + data: DataConfig = field(default_factory=DataConfig) + model: ModelConfig = field(default_factory=ModelConfig) + train: TrainConfig = field(default_factory=TrainConfig) + loss: LossConfig = field(default_factory=LossConfig) + distributed: DistributedConfig = field(default_factory=DistributedConfig) + wandb: WandbConfig = field(default_factory=WandbConfig) + + + +def _resolve_path(path: str, cwd: Path) -> str: + p = Path(path) + if p.is_absolute(): + return str(p) + return str((cwd / p).resolve()) + + + +def _validate_config(cfg: TrainStage1FlowConfig) -> None: + # if cfg.data.pairs_per_epoch <= 0: + # raise ValueError("data.pairs_per_epoch must be > 0") + if cfg.train.max_steps <= 0: + raise ValueError("train.max_steps must be > 0") + if cfg.train.batch_size_per_gpu <= 0: + raise ValueError("train.batch_size_per_gpu must be > 0") + if cfg.train.grad_accum_steps <= 0: + raise ValueError("train.grad_accum_steps must be > 0") + if cfg.train.save_every <= 0: + raise ValueError("train.save_every must be > 0") + is_dexycb = cfg.data.dataset_name == "dexycb" + if is_dexycb: + if not cfg.data.dexycb_root: + raise ValueError("data.dexycb_root must be set when dataset_name='dexycb'") + if not cfg.data.val_subjects and not cfg.data.train_subjects: + raise ValueError( + "At least one of data.val_subjects or data.train_subjects must be set " + "when dataset_name='dexycb'" + ) + else: + if cfg.data.mask_mode not in { + "green_channel", "blue_channel", "auto_smallest_nonzero_region" + }: + raise ValueError( + "data.mask_mode must be 'green_channel', 'blue_channel', or " + "'auto_smallest_nonzero_region'" + ) + if cfg.model.pointmap_embedder_index < 0: + raise ValueError("model.pointmap_embedder_index must be >= 0") + if cfg.distributed.backend.lower().strip() not in {"ddp", "deepspeed"}: + raise ValueError("distributed.backend must be 'ddp' or 'deepspeed'") + if cfg.train.timestep_sampler not in {"uniform", "logit_normal"}: + raise ValueError("train.timestep_sampler must be one of: uniform, logit_normal") + if cfg.train.cfg_prob_uncond < 0.0 or cfg.train.cfg_prob_uncond > 1.0: + raise ValueError("train.cfg_prob_uncond must be in [0, 1]") + if cfg.train.cfg_prob_pm_drop < 0.0 or cfg.train.cfg_prob_pm_drop > 1.0: + raise ValueError("train.cfg_prob_pm_drop must be in [0, 1]") + if cfg.train.cfg_prob_uncond + cfg.train.cfg_prob_pm_drop > 1.0: + raise ValueError("train.cfg_prob_uncond + train.cfg_prob_pm_drop must be <= 1") + if cfg.wandb.mode.lower().strip() not in {"online", "offline", "disabled"}: + raise ValueError("wandb.mode must be one of: online, offline, disabled") + if cfg.wandb.enable_wandb and not cfg.wandb.project.strip(): + raise ValueError("wandb.project must be non-empty when wandb.enable_wandb=true") + + + +def load_config(config_path: str, overrides: list[str] | None = None) -> TrainStage1FlowConfig: + config_path = str(Path(config_path).resolve()) + cfg = OmegaConf.load(config_path) + if overrides: + cfg = OmegaConf.merge(cfg, OmegaConf.from_dotlist(overrides)) + + base = OmegaConf.structured(TrainStage1FlowConfig) + merged = OmegaConf.merge(base, cfg) + obj = OmegaConf.to_object(merged) + if not isinstance(obj, TrainStage1FlowConfig): + obj = TrainStage1FlowConfig(**obj) + + cwd = Path.cwd() + obj.experiment.output_dir = _resolve_path(obj.experiment.output_dir, cwd) + obj.data.ho3d_root = _resolve_path(obj.data.ho3d_root, cwd) + if obj.data.dexycb_root: + obj.data.dexycb_root = _resolve_path(obj.data.dexycb_root, cwd) + if obj.data.pointmap_dir: + obj.data.pointmap_dir = _resolve_path(obj.data.pointmap_dir, cwd) + if obj.data.initial_pose_file: + obj.data.initial_pose_file = _resolve_path(obj.data.initial_pose_file, cwd) + obj.model.pipeline_yaml = _resolve_path(obj.model.pipeline_yaml, cwd) + obj.distributed.deepspeed_config = _resolve_path(obj.distributed.deepspeed_config, cwd) + if obj.wandb.enable_wanb: + obj.wandb.enable_wandb = True + + _validate_config(obj) + return obj + + + +def config_to_dict(cfg: TrainStage1FlowConfig) -> dict: + return asdict(cfg) diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/dataset_unified.py b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/dataset_unified.py new file mode 100644 index 0000000000000000000000000000000000000000..8da80910411abd7440f1e1bb2d41363d3037efd2 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/dataset_unified.py @@ -0,0 +1,537 @@ +"""Unified single-frame 6D-pose dataset (DexYCB + RoboCasa). + +Replaces the old frame-PAIR datasets. Each sample is ONE input frame: + inputs : RGB + object mask + image-shaped (H,W,3) PyTorch3D pointmap + targets : r (rot6d), t (metric, m), s (metric longest-axis extent, m) + +All quantities are in the SAME convention for both datasets (audited): + * pointmap : PyTorch3D camera frame (x-left,y-up,z-fwd), meters + * r, t : PyTorch3D, row-vector cam = s*(canon @ R) + t + * s : scalar metric longest-axis extent (meters) + +Compat: the model wrapper is a first_frame(shape)+query(pose) two-pass; for a +single frame we alias first_frame == query == the input frame. The batch keys +match the old collate plus `gt_scale` and `dataset`. + +Per-dataset saved layout consumed here: + RoboCasa: //ep/pointmap_.npz (pointmap, frame_indices), + .../pseudo_gt_targets.npz (rotation_matrix,translation,s_metric), + //ep/views/.npz (obj_mask), + /videos/chunk-000/observation.images./episode_*.mp4 (RGB) + DexYCB: color_.jpg + labels_.npz (seg) under //, + ///raw/.pt (N_cams,H,W,3), + initial_pose_ycb.pt (rotation_matrix,translation,scale) +""" +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch.utils.data import Dataset + +from .geometry import matrix_to_rot6d, normalize_rot6d + + +# ----------------------------- sample descriptor ------------------------------ +@dataclass +class Sample: + dataset: str # "robocasa" | "dexycb" + key: str # human-readable uid + object_label: int + object_name: str + # robocasa + pm_npz: str = "" + pgt_npz: str = "" + views_npz: str = "" + video: str = "" + cam: str = "" + frame_abs: int = -1 + # dexycb + color_path: str = "" + labels_path: str = "" + pointmap_pt: str = "" + cam_index: int = 0 + rot: tuple = () # (3,3) row-vector P3D + trans: tuple = () # (3,) + scale: float = 0.0 + ycb_id: int = -1 + # dexycb per-frame GT (manifest-driven path): GT read from labels_.npz:pose_y + # at load time. obj_index = row in pose_y; align_R composes the per-frame relative + # motion onto the SAM-3D-canonical anchor (see build_manifest_dexycb_jsonl). + perframe_pose: bool = False + obj_index: int = -1 + align_R: tuple = () # (3,3) constant per sequence: R0 @ Rp_anchor.T + # shape-K/V is FORCED to the video's first frame (GT is anchored to R0), so the + # first_frame slot is set to frame0 while query stays the moving frame. + anchor_color_path: str = "" + anchor_labels_path: str = "" + anchor_pointmap_pt: str = "" + + +# --------------------------------- helpers ------------------------------------ +# OpenCV (x-right,y-down,z-fwd) <-> PyTorch3D (x-left,y-up,z-fwd): M = diag(-1,-1,1) +_CV_TO_P3D = np.diag([-1.0, -1.0, 1.0]).astype(np.float64) + + +def _load_dexycb_pointmap_view(path: str, cam_index: int) -> torch.Tensor: + """Return the (H,W,3) PyTorch3D pointmap for camera ``cam_index``. + + Two on-disk formats are supported transparently (same coordinate values): + * dense ``.pt`` : (N_cams,H,W,3) float32 -> slice [cam_index] + * compact ``.npz`` : sparse per-view (scheme B, see + tools/compact_dexycb_pointmaps.py). Reconstructed bit-exactly here by + scattering the stored float32 xyz into a zero (H,W,3) grid. No transform, + no sign flip -- the stored values are already in this camera's P3D frame. + """ + if path.endswith(".npz"): + d = np.load(path, allow_pickle=False) + n, h, w = (int(x) for x in d["shape"]) + counts = np.asarray(d["counts"], np.int64) + off = int(counts[:cam_index].sum()) + k = int(counts[cam_index]) + view = np.zeros((h * w, 3), np.float32) + if k: + pix = np.asarray(d["pix"][off:off + k], np.int64) + view[pix] = np.asarray(d["xyz"][off:off + k], np.float32) + pm = torch.from_numpy(view.reshape(h, w, 3)) + else: + pm = torch.load(path, map_location="cpu", weights_only=False) + pm = pm[cam_index] if pm.dim() == 4 else pm + # DexYCB stores invalid/background as exact (0,0,0) (finite). The SS pointmap + # embedder flags invalid by isfinite() only (dit/embedder/pointmap.py) and the + # cond builder passes no valid_mask, so a (0,0,0) background would be embedded as + # real geometry at the origin. Convert exact-zero pixels -> NaN to match MoGe's + # NaN-background convention, which the whole pipeline (resize -> SSI -> embedder) + # already handles. (A camera-frame point at exactly (0,0,0) is z=0, i.e. invalid.) + pm = pm.float().clone() + pm[pm.abs().sum(dim=-1) < 1e-6] = float("nan") + return pm + + +def _to_rot6d(R: np.ndarray) -> np.ndarray: + """row-vector (3,3) -> normalized 6D, matching the old datasets.""" + r = torch.from_numpy(np.asarray(R, np.float32)).reshape(1, 3, 3) + return normalize_rot6d(matrix_to_rot6d(r).squeeze(0)).cpu().numpy().astype(np.float32) + + +class _VideoCache: + """Lazy per-path cv2.VideoCapture cache (worker-local).""" + def __init__(self): + self._caps: dict[str, Any] = {} + + def frame(self, path: str, idx: int) -> np.ndarray: + import cv2 + cap = self._caps.get(path) + if cap is None: + cap = cv2.VideoCapture(path) + self._caps[path] = cap + cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx)) + ok, bgr = cap.read() + if not ok: + raise RuntimeError(f"failed to read frame {idx} from {path}") + return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) + + +# --------------------------------- dataset ------------------------------------ +class UnifiedPoseDataset(Dataset): + def __init__(self, samples: list[Sample]): + if not samples: + raise RuntimeError("UnifiedPoseDataset: empty sample list") + self.samples = samples + self._vid = _VideoCache() + + def __len__(self) -> int: + return len(self.samples) + + # -- robocasa ---------------------------------------------------------- + def _load_robocasa(self, s: Sample) -> dict[str, Any]: + pm_d = np.load(s.pm_npz, allow_pickle=False) + pm_fidx = list(pm_d["frame_indices"]) + pj = pm_fidx.index(s.frame_abs) + pointmap = torch.from_numpy(pm_d["pointmap"][pj].astype(np.float32)) # (H,W,3) P3D + + pgt = np.load(s.pgt_npz, allow_pickle=True) + cams = [str(c) for c in pgt["cams"]]; pg_fidx = list(pgt["frame_indices"]) + ci = cams.index(s.cam); fi = pg_fidx.index(s.frame_abs) + R = pgt["rotation_matrix"][ci, fi]; t = pgt["translation"][ci, fi] + scale = float(pgt["s_metric"]) + + vd = np.load(s.views_npz, allow_pickle=False) + vj = list(vd["frame_indices"]).index(s.frame_abs) + mask = vd["obj_mask"][vj].astype(np.uint8) + H, W = mask.shape + rgb = self._vid.frame(s.video, s.frame_abs) + if rgb.shape[:2] != (H, W): + import cv2 + rgb = cv2.resize(rgb, (W, H), interpolation=cv2.INTER_AREA) + return self._pack(rgb, mask, pointmap, R, t, scale, s) + + # -- dexycb ------------------------------------------------------------ + def _load_dexycb(self, s: Sample) -> dict[str, Any]: + from PIL import Image + rgb = np.asarray(Image.open(s.color_path).convert("RGB"), np.uint8) + lab = np.load(s.labels_path) + mask = (lab["seg"] == s.ycb_id).astype(np.uint8) + # (H,W,3) P3D camera-frame map; dense .pt or compact .npz, same values + pointmap = _load_dexycb_pointmap_view(s.pointmap_pt, s.cam_index).float() + if s.perframe_pose: + # per-frame GT from pose_y (OpenCV CAD->cam) -> training frame: + # t = t_cv @ M (matches initial_pose translation, M=diag(-1,-1,1)) + # R = align_R @ (M @ R_cv).T (SAM-3D-canonical frame; relative motion preserved) + py = lab["pose_y"][s.obj_index] + R_cv = np.asarray(py[:, :3], np.float64); t_cv = np.asarray(py[:, 3], np.float64) + Rp = (_CV_TO_P3D @ R_cv).T + R = (np.asarray(s.align_R, np.float64).reshape(3, 3) @ Rp).astype(np.float32) + t = (t_cv @ _CV_TO_P3D).astype(np.float32) + else: + R = np.asarray(s.rot, np.float32).reshape(3, 3) + t = np.asarray(s.trans, np.float32).reshape(3) + # FORCE shape K/V to the video's first frame (frame0); query stays this frame + if s.anchor_color_path and s.anchor_pointmap_pt: + a_rgb = np.asarray(Image.open(s.anchor_color_path).convert("RGB"), np.uint8) + a_mask = (np.load(s.anchor_labels_path)["seg"] == s.ycb_id).astype(np.uint8) + a_pm = _load_dexycb_pointmap_view(s.anchor_pointmap_pt, s.cam_index).float() + else: + a_rgb, a_mask, a_pm = rgb, mask, pointmap + return self._pack(rgb, mask, pointmap, R, t, float(s.scale), s, + ff_rgb=a_rgb, ff_mask=a_mask, ff_pm=a_pm) + + # -- common record ----------------------------------------------------- + def _pack(self, rgb, mask, pointmap, R, t, scale, s: Sample, + ff_rgb=None, ff_mask=None, ff_pm=None) -> dict[str, Any]: + rgb = np.ascontiguousarray(rgb, np.uint8) + mask = np.ascontiguousarray(mask, np.uint8) + ff_rgb = rgb if ff_rgb is None else np.ascontiguousarray(ff_rgb, np.uint8) + ff_mask = mask if ff_mask is None else np.ascontiguousarray(ff_mask, np.uint8) + ff_pm = pointmap if ff_pm is None else ff_pm + rec = { + # first_frame = video frame0 (shape-K/V source); query = this frame + "first_frame_rgb": ff_rgb, "first_frame_mask": ff_mask, + "query_rgb": rgb, "query_mask": mask, + "first_frame_pointmap": ff_pm, "query_pointmap": pointmap, + # targets + "gt_rot6d": _to_rot6d(R), + "gt_trans_metric": np.asarray(t, np.float32).reshape(3), + "gt_scale": np.asarray([scale], np.float32), + # meta + "dataset": s.dataset, "sequence": s.key, "query_frame": s.frame_abs, + "object_name": s.object_name, "object_label": int(s.object_label), + } + return rec + + def __getitem__(self, i: int) -> dict[str, Any]: + s = self.samples[i] + return self._load_robocasa(s) if s.dataset == "robocasa" else self._load_dexycb(s) + + +# --------------------------------- collate ------------------------------------ +def collate_unified(batch: list[dict[str, Any]]) -> dict[str, Any]: + out = { + "first_frame_rgb": [b["first_frame_rgb"] for b in batch], + "first_frame_mask": [b["first_frame_mask"] for b in batch], + "query_rgb": [b["query_rgb"] for b in batch], + "query_mask": [b["query_mask"] for b in batch], + "first_frame_pointmap": [b["first_frame_pointmap"] for b in batch], + "query_pointmap": [b["query_pointmap"] for b in batch], + "gt_rot6d": torch.from_numpy(np.stack([b["gt_rot6d"] for b in batch])).float(), + "gt_trans_metric": torch.from_numpy(np.stack([b["gt_trans_metric"] for b in batch])).float(), + "gt_scale": torch.from_numpy(np.stack([b["gt_scale"] for b in batch])).float(), + "dataset": [b["dataset"] for b in batch], + "sequence": [b["sequence"] for b in batch], + "query_frame": [b["query_frame"] for b in batch], + "object_name": [b["object_name"] for b in batch], + "object_label": torch.tensor([b["object_label"] for b in batch], dtype=torch.long), + } + return out + + +# ------------------------------ manifest builders ----------------------------- +def build_manifest_robocasa(pm_root: str, mv_root: str, dataset_root: str, + task: str, label_map: dict[str, int] | None = None, + stride: int = 1) -> list[Sample]: + """Scan saved robocasa episodes -> per (ep,cam,valid-frame) samples.""" + pm_root, mv_root, dataset_root = Path(pm_root), Path(mv_root), Path(dataset_root) + label_map = label_map or {} + samples: list[Sample] = [] + for ep_dir in sorted(pm_root.glob("ep*")): + pgt_fp = ep_dir / "pseudo_gt_targets.npz" + if not pgt_fp.is_file(): + continue + meta_fp = (mv_root / ep_dir.name / "object_meta.json") + obj_name = "unknown" + if meta_fp.is_file(): + obj_name = json.loads(meta_fp.read_text()).get("mesh_files", ["unknown"])[0] + obj_name = Path(obj_name).parent.parent.name if "/" in obj_name else obj_name + label = label_map.setdefault(obj_name, len(label_map)) + for pm_fp in sorted(ep_dir.glob("pointmap_*.npz")): + cam = pm_fp.stem.replace("pointmap_", "") + fidx = list(np.load(pm_fp, allow_pickle=False)["frame_indices"]) + video = str(dataset_root / "videos" / "chunk-000" / + f"observation.images.{cam}" / f"episode_{int(ep_dir.name[2:]):06d}.mp4") + views = str(mv_root / ep_dir.name / "views" / f"{cam}.npz") + for fa in fidx[::max(1, stride)]: + samples.append(Sample( + dataset="robocasa", key=f"{task}/{ep_dir.name}/{cam}", + object_label=label, object_name=obj_name, + pm_npz=str(pm_fp), pgt_npz=str(pgt_fp), views_npz=views, + video=video, cam=cam, frame_abs=int(fa))) + return samples + + +def build_manifest_dexycb(dexycb_root: str, initial_pose_tmpl: str, pointmap_dir: str, + subjects: list[str] | None = None, ycb_ids: list[int] | None = None, + label_map: dict[str, int] | None = None, stride: int = 1, + pointmap_format: str = "dense") -> list[Sample]: + """Scan stored DexYCB artifacts -> per (sequence, frame) samples. + + initial_pose_tmpl: e.g. ".../initial_pose_postprocessed/{subject}/initial_pose_ycb{ycb_id}.pt" + -> sequences[subject/session/serial] = {rotation_matrix(3,3) P3D row-vec, translation(3), scale(1)} + pointmap_dir: /{subject}/{session}/raw/{frame:06d}.pt (N_cams,H,W,3), per-cam slice + rgb/mask: /{subject}/{session}/{serial}/color_/labels_{frame:06d} + Paths are all arguments; nothing is inferred or computed on GPU. + """ + import glob + import yaml + label_map = label_map if label_map is not None else {} + dexycb_root = Path(dexycb_root); pointmap_dir = Path(pointmap_dir) + _serials_cache: dict[str, list[str]] = {} + + def serials_of(subject, session): + key = f"{subject}/{session}" + if key not in _serials_cache: + meta_fp = dexycb_root / subject / session / "meta.yml" + if not meta_fp.is_file(): + # subject/session in initial_pose but raw data absent locally → skip + _serials_cache[key] = [] + else: + m = yaml.safe_load(meta_fp.read_text()) + _serials_cache[key] = list(m.get("serials", [])) + return _serials_cache[key] + + # resolve which initial_pose files exist (glob over {subject},{ycb_id}) + glob_pat = initial_pose_tmpl.replace("{subject}", "*").replace("{ycb_id}", "*") + samples: list[Sample] = [] + for fp in sorted(glob.glob(glob_pat)): + ycb_id = int(Path(fp).stem.split("ycb")[-1]) + if ycb_ids and ycb_id not in ycb_ids: + continue + d = torch.load(fp, map_location="cpu", weights_only=False) + for seq, pose in d.get("sequences", {}).items(): + subject, session, serial = seq.split("/", 2) + if subjects and subject not in subjects: + continue + ser_list = serials_of(subject, session) + if serial not in ser_list: + continue + cam_index = ser_list.index(serial) + label = label_map.setdefault(f"ycb{ycb_id}", len(label_map)) + R = np.asarray(pose["rotation_matrix"], np.float32).reshape(3, 3) + t = np.asarray(pose["translation"], np.float32).reshape(3) + s = float(np.asarray(pose["scale"]).reshape(-1)[0]) + # full-scene per-view maps are object-independent -> per-session layout + # produced by dexycb_utils/batch.py (aggregate_pointmaps): //raw/ + pm_dir = pointmap_dir / subject / session / "raw" + pm_ext = "*.npz" if pointmap_format == "compact" else "*.pt" + for pm_fp in sorted(pm_dir.glob(pm_ext))[::max(1, stride)]: + frame = int(pm_fp.stem) + ftok = f"{frame:06d}" + color = dexycb_root / subject / session / serial / f"color_{ftok}.jpg" + labels = dexycb_root / subject / session / serial / f"labels_{ftok}.npz" + if not (color.is_file() and labels.is_file()): + continue + samples.append(Sample( + dataset="dexycb", key=seq, object_label=label, object_name=f"ycb{ycb_id}", + color_path=str(color), labels_path=str(labels), pointmap_pt=str(pm_fp), + cam_index=cam_index, rot=tuple(R.ravel().tolist()), trans=tuple(t.tolist()), + scale=s, ycb_id=ycb_id, frame_abs=frame)) + return samples + + +def build_manifest_dexycb_jsonl(manifest_path: str, dexycb_root: str, pointmap_dir: str, + initial_pose_tmpl: str, + label_map: dict[str, int] | None = None, + pointmap_format: str = "dense") -> list[Sample]: + """Manifest-driven DexYCB samples with PER-FRAME GT (for moving-object training). + + Reads a JSONL manifest (rows: subject, session, serial, ycb_id, obj_index, frame, ...) + and emits one Sample per row. GT pose is read per-frame from labels_.npz:pose_y + at load time; ``align_R`` (= R0 @ Rp_anchor.T, constant per sequence) maps the per-frame + relative motion onto the SAM-3D-canonical anchor pose ``R0`` from initial_pose, so targets + stay in the same frame the pretrained pose head uses. Scale is per-object-constant (s0). + Sequences without an initial_pose entry are skipped (no canonical anchor). + """ + import glob + import yaml + label_map = label_map if label_map is not None else {} + dexycb_root = Path(dexycb_root); pointmap_dir = Path(pointmap_dir) + + rows = [json.loads(l) for l in open(manifest_path)] + by_seq: dict[tuple, list[dict]] = {} + for r in rows: + by_seq.setdefault((r["subject"], r["session"], r["serial"], int(r["ycb_id"])), []).append(r) + + # initial_pose lookup: (subject/session/serial, ycb_id) -> {rotation_matrix, scale} + ip: dict[tuple[str, int], dict] = {} + glob_pat = initial_pose_tmpl.replace("{subject}", "*").replace("{ycb_id}", "*") + for fp in sorted(glob.glob(glob_pat)): + yid = int(Path(fp).stem.split("ycb")[-1]) + d = torch.load(fp, map_location="cpu", weights_only=False) + for seq, pose in d.get("sequences", {}).items(): + ip[(seq, yid)] = pose + + _serials: dict[str, list[str]] = {} + def serials_of(subject, session): + key = f"{subject}/{session}" + if key not in _serials: + mp = dexycb_root / subject / session / "meta.yml" + _serials[key] = list(yaml.safe_load(mp.read_text()).get("serials", [])) if mp.is_file() else [] + return _serials[key] + + samples: list[Sample] = [] + for (subject, session, serial, ycb_id), seq_rows in by_seq.items(): + pose0 = ip.get((f"{subject}/{session}/{serial}", ycb_id)) + if pose0 is None: + continue # no canonical anchor -> skip sequence + R0 = np.asarray(pose0["rotation_matrix"], np.float64).reshape(3, 3) + s0 = float(np.asarray(pose0["scale"]).reshape(-1)[0]) + ser_list = serials_of(subject, session) + if serial not in ser_list: + continue + cam_index = ser_list.index(serial) + seq_rows.sort(key=lambda r: r["frame"]) + anchor = seq_rows[0] + # anchor pose_y -> P3D rotation; align maps relative motion onto R0 + alab = dexycb_root / subject / session / serial / f"labels_{anchor['frame']:06d}.npz" + if not alab.is_file(): + continue + py_a = np.load(alab)["pose_y"][int(anchor["obj_index"])] + Rp_a = (_CV_TO_P3D @ np.asarray(py_a[:, :3], np.float64)).T + align = (R0 @ Rp_a.T) + label = label_map.setdefault(f"ycb{ycb_id}", len(label_map)) + # frame0 = the sequence anchor; its frame fills the first_frame (shape-K/V) slot. + pm_ext = "npz" if pointmap_format == "compact" else "pt" + a_ftok = f"{int(anchor['frame']):06d}" + a_color = dexycb_root / subject / session / serial / f"color_{a_ftok}.jpg" + a_pm = pointmap_dir / subject / session / "raw" / f"{a_ftok}.{pm_ext}" + anchor_ok = a_color.is_file() and a_pm.is_file() # alab already verified above + for r in seq_rows: + ftok = f"{int(r['frame']):06d}" + color = dexycb_root / subject / session / serial / f"color_{ftok}.jpg" + labels = dexycb_root / subject / session / serial / f"labels_{ftok}.npz" + pm_pt = pointmap_dir / subject / session / "raw" / f"{ftok}.{pm_ext}" + if not (color.is_file() and labels.is_file() and pm_pt.is_file()): + continue + samples.append(Sample( + dataset="dexycb", key=f"{subject}/{session}/{serial}", + object_label=label, object_name=f"ycb{ycb_id}", + color_path=str(color), labels_path=str(labels), pointmap_pt=str(pm_pt), + cam_index=cam_index, scale=s0, ycb_id=ycb_id, frame_abs=int(r["frame"]), + perframe_pose=True, obj_index=int(r["obj_index"]), + align_R=tuple(align.ravel().tolist()), + anchor_color_path=(str(a_color) if anchor_ok else ""), + anchor_labels_path=(str(alab) if anchor_ok else ""), + anchor_pointmap_pt=(str(a_pm) if anchor_ok else ""))) + return samples + + +def _split_sequences(samples: list[Sample], seed: int, val_frac: float, + val_num_sequences: int, val_stride: int): + """SEQUENCE-LEVEL holdout: keep WHOLE sequences in val so no sequence leaks + across the train/val boundary (the relative-pose validation forms intra- + sequence frame pairs, so a sequence split across both sides would be leakage). + + Grouping key is the sample ``key`` (DexYCB: subject/session/serial; RoboCasa: + task/ep/cam) -> the per-(sequence,object) uid. ``val_fraction`` is now + interpreted as a fraction of SEQUENCES (not frames); ``val_num_sequences`` (if + > 0) overrides it with an absolute sequence count. Only sequences with enough + frames to form at least one relative pair at ``val_stride`` are eligible for + val, so every held-out sequence yields >=1 pair. + + Returns (train_samples, val_samples) where val_samples may be []. + """ + from collections import defaultdict + by_key: dict[str, list[Sample]] = defaultdict(list) + for s in samples: + by_key[s.key].append(s) + keys = sorted(by_key.keys()) + + # eligible val sequences = those that can form a relative pair (span >= stride + # in absolute frame index). group_by_sequence in eval sorts by frame; mirror it. + def _has_pair(seq_samples: list[Sample]) -> bool: + fr = sorted(s.frame_abs for s in seq_samples) + fset = set(fr) + return any((f + val_stride) in fset for f in fr) + + eligible = [k for k in keys if _has_pair(by_key[k])] + + rng = np.random.default_rng(seed) + perm = rng.permutation(len(eligible)) + eligible = [eligible[i] for i in perm] + + if val_num_sequences and val_num_sequences > 0: + n_val_seq = min(int(val_num_sequences), len(eligible)) + elif val_frac > 0: + n_val_seq = int(round(len(keys) * val_frac)) + n_val_seq = min(max(n_val_seq, 1) if len(keys) > 1 else 0, len(eligible)) + else: + n_val_seq = 0 + + val_keys = set(eligible[:n_val_seq]) + train_samples = [s for s in samples if s.key not in val_keys] + val_samples = [s for s in samples if s.key in val_keys] + return train_samples, val_samples + + +def build_train_val_datasets(data_cfg, seed: int): + """Entry used by train_stage1_flow_lora.py. Dispatches on data_cfg.dataset_name. + + Validation split is SEQUENCE-LEVEL (see ``_split_sequences``): whole sequences + are held out so the relative-pose validation can form intra-sequence frame + pairs with no train/val leakage. Controlled by ``data_cfg.val_fraction`` + (fraction of SEQUENCES), optional ``data_cfg.val_num_sequences`` (absolute + override), and ``data_cfg.val_stride`` (relative-pair stride, default 10, + matching inference/eval.py). + """ + name = getattr(data_cfg, "dataset_name", "robocasa") + # "dense" (default, .pt) | "compact" (.npz, see tools/compact_dexycb_pointmaps.py) + pm_fmt = str(getattr(data_cfg, "pointmap_format", "dense")) + samples: list[Sample] = [] + label_map: dict[str, int] = {} # shared object-label map across datasets + stride = int(getattr(data_cfg, "frame_stride", 1)) + val_num_seq = int(getattr(data_cfg, "val_num_sequences", 0) or 0) + val_stride = int(getattr(data_cfg, "val_stride", 10)) + if name == "dexycb_jsonl": + # manifest-driven DexYCB with per-frame GT (moving-object training) + samples += build_manifest_dexycb_jsonl( + manifest_path=data_cfg.dexycb_manifest_jsonl, dexycb_root=data_cfg.dexycb_root, + pointmap_dir=data_cfg.pointmap_dir, initial_pose_tmpl=data_cfg.initial_pose_file, + label_map=label_map, pointmap_format=pm_fmt) + if not samples: + raise RuntimeError(f"no samples from manifest {data_cfg.dexycb_manifest_jsonl!r}") + val_frac = float(getattr(data_cfg, "val_fraction", 0.0)) + train_s, val_s = _split_sequences(samples, seed, val_frac, val_num_seq, val_stride) + val = UnifiedPoseDataset(val_s) if val_s else None + return UnifiedPoseDataset(train_s or samples), val + if name in ("robocasa", "unified"): + samples += build_manifest_robocasa( + pm_root=data_cfg.robocasa_pm_root, mv_root=data_cfg.robocasa_mv_root, + dataset_root=data_cfg.robocasa_dataset_root, task=getattr(data_cfg, "robocasa_task", "task"), + label_map=label_map, stride=stride) + if name in ("dexycb", "unified"): + samples += build_manifest_dexycb( + dexycb_root=data_cfg.dexycb_root, initial_pose_tmpl=data_cfg.initial_pose_file, + pointmap_dir=data_cfg.pointmap_dir, + subjects=list(getattr(data_cfg, "subjects", []) or []) or None, + ycb_ids=list(getattr(data_cfg, "target_ycb_ids", []) or []) or None, + label_map=label_map, stride=stride, pointmap_format=pm_fmt) + if not samples: + raise RuntimeError(f"no samples for dataset_name={name!r}") + val_frac = float(getattr(data_cfg, "val_fraction", 0.05)) + train_s, val_s = _split_sequences(samples, seed, val_frac, val_num_seq, val_stride) + val = UnifiedPoseDataset(val_s) if val_s else None + train = UnifiedPoseDataset(train_s or samples) + return train, val diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/distributed.py b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/distributed.py new file mode 100644 index 0000000000000000000000000000000000000000..21e26ebf3b90e26e3c0cdf08edc4f1ef614c3bc9 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/distributed.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +from dataclasses import dataclass +import json +import os + +import torch +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP + + +@dataclass +class DistributedState: + rank: int + local_rank: int + world_size: int + distributed: bool + device: torch.device + backend: str + + @property + def is_main(self) -> bool: + return self.rank == 0 + + +@dataclass +class RuntimeHandles: + model: torch.nn.Module + optimizer: torch.optim.Optimizer + scheduler: object + use_deepspeed: bool + + + +def init_distributed(backend: str = "ddp") -> DistributedState: + rank = int(os.environ.get("RANK", "0")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + else: + device = torch.device("cpu") + + distributed = world_size > 1 + if distributed and not dist.is_initialized(): + dist_backend = "nccl" if device.type == "cuda" else "gloo" + dist.init_process_group(backend=dist_backend, init_method="env://") + + return DistributedState( + rank=rank, + local_rank=local_rank, + world_size=world_size, + distributed=distributed, + device=device, + backend=backend.lower().strip(), + ) + + + +def setup_runtime( + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + scheduler, + state: DistributedState, + deepspeed_config_path: str, +) -> RuntimeHandles: + backend = state.backend + + if backend == "deepspeed": + try: + import deepspeed + except Exception as exc: + if state.is_main: + print( + f"[distributed] DeepSpeed unavailable ({exc}); falling back to DDP.", + flush=True, + ) + backend = "ddp" + else: + cfg_path = os.path.abspath(deepspeed_config_path) + if not os.path.isfile(cfg_path): + raise FileNotFoundError(f"DeepSpeed config not found: {cfg_path}") + with open(cfg_path, "r", encoding="utf-8") as f: + _ = json.load(f) + + engine, ds_optimizer, _, ds_scheduler = deepspeed.initialize( + model=model, + model_parameters=[p for p in model.parameters() if p.requires_grad], + optimizer=optimizer, + lr_scheduler=scheduler, + config=cfg_path, + ) + if state.is_main: + print(f"[distributed] DeepSpeed initialized with {cfg_path}", flush=True) + return RuntimeHandles( + model=engine, + optimizer=ds_optimizer, + scheduler=ds_scheduler, + use_deepspeed=True, + ) + + wrapped_model = model + if state.distributed: + wrapped_model = DDP( + model, + device_ids=[state.local_rank] if state.device.type == "cuda" else None, + output_device=state.local_rank if state.device.type == "cuda" else None, + find_unused_parameters=False, + ) + return RuntimeHandles( + model=wrapped_model, + optimizer=optimizer, + scheduler=scheduler, + use_deepspeed=False, + ) + + + +def barrier() -> None: + if dist.is_available() and dist.is_initialized(): + if torch.cuda.is_available(): + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + dist.barrier(device_ids=[local_rank]) + else: + dist.barrier() + + + +def cleanup_distributed() -> None: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + + +def reduce_scalar_dict(metrics: dict[str, torch.Tensor], state: DistributedState) -> dict[str, float]: + out: dict[str, float] = {} + for key, value in metrics.items(): + tensor = value.detach().to(state.device, non_blocking=True).float() + if state.distributed: + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + tensor = tensor / float(state.world_size) + out[key] = float(tensor.item()) + return out diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/geometry.py b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..70fc75cfcb37cf8694b23c104518a09ac58d97e1 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/geometry.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import math + +import torch +from pytorch3d.transforms import axis_angle_to_matrix + +from .rotation_constants import ROTATION_6D_MEAN, ROTATION_6D_STD + + +# ── 6D rotation: COLUMN convention ────────────────────────────────── +# The pretrained SAM-3D pose_decoder uses column-based Gram-Schmidt: +# encode: 6D = [R[:,0], R[:,1]] (first two COLUMNS of R) +# decode: torch.stack([b1, b2, b3], dim=-1) → columns +# +# pytorch3d uses ROW convention (dim=-2) — INCOMPATIBLE with pose_decoder. +# Do NOT use pytorch3d matrix_to_rotation_6d / rotation_6d_to_matrix here. + + +def matrix_to_rot6d(R: torch.Tensor) -> torch.Tensor: + """Rotation matrix → 6D (column convention, matching pose_decoder).""" + return torch.cat([R[..., :, 0], R[..., :, 1]], dim=-1) + + +def rot6d_to_matrix(rot6d: torch.Tensor) -> torch.Tensor: + """6D → rotation matrix (column Gram-Schmidt, matching pose_decoder).""" + a1 = rot6d[..., 0:3] + a2 = rot6d[..., 3:6] + b1 = torch.nn.functional.normalize(a1, dim=-1) + b2 = a2 - torch.sum(b1 * a2, dim=-1, keepdim=True) * b1 + b2 = torch.nn.functional.normalize(b2, dim=-1) + b3 = torch.cross(b1, b2, dim=-1) + return torch.stack([b1, b2, b3], dim=-1) + + +def normalize_rot6d(raw_rot6d: torch.Tensor) -> torch.Tensor: + """raw 6D → normalized 6D (matching pretrained backbone convention).""" + mean = ROTATION_6D_MEAN.to(device=raw_rot6d.device, dtype=raw_rot6d.dtype) + std = ROTATION_6D_STD.to(device=raw_rot6d.device, dtype=raw_rot6d.dtype) + return (raw_rot6d - mean) / std + + +def denormalize_rot6d(norm_rot6d: torch.Tensor) -> torch.Tensor: + """normalized 6D → raw 6D (inverse of normalize_rot6d).""" + mean = ROTATION_6D_MEAN.to(device=norm_rot6d.device, dtype=norm_rot6d.dtype) + std = ROTATION_6D_STD.to(device=norm_rot6d.device, dtype=norm_rot6d.dtype) + return norm_rot6d * std + mean + + +def ssi_to_metric( + translation_ssi: torch.Tensor, + pointmap_scale: torch.Tensor, + pointmap_shift: torch.Tensor, +) -> torch.Tensor: + return pointmap_scale * translation_ssi + pointmap_shift + + + +def geodesic_distance_deg(pred_rotation: torch.Tensor, gt_rotation: torch.Tensor) -> torch.Tensor: + rot = pred_rotation @ gt_rotation.transpose(-1, -2) + trace = torch.diagonal(rot, dim1=-2, dim2=-1).sum(dim=-1) + cos = torch.clamp((trace - 1.0) * 0.5, -1.0, 1.0) + return torch.acos(cos) * (180.0 / math.pi) diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/losses.py b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..1cc97199440ebba255ab468db7e1ea390c79c35f --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/losses.py @@ -0,0 +1,204 @@ +from __future__ import annotations + +import torch + +from .geometry import denormalize_rot6d, geodesic_distance_deg, rot6d_to_matrix + + +def _get_loss_weights(loss_cfg) -> tuple[float, float, float]: + rot_weight = float(getattr(loss_cfg, "rotation_weight", 1.0)) + trans_weight = float(getattr(loss_cfg, "translation_weight", 1.0)) + scale_weight = float(getattr(loss_cfg, "scale_weight", 1.0)) + return rot_weight, trans_weight, scale_weight + + +def _get_min_valid_pointmap_scale(loss_cfg) -> float: + return float(getattr(loss_cfg, "min_valid_pointmap_scale", 1e-6)) + + +def _get_max_valid_gt_translation_norm_m(loss_cfg) -> float: + return float(getattr(loss_cfg, "max_valid_gt_translation_norm_m", 10.0)) + + +def _skip_invalid_translation_supervision(loss_cfg) -> bool: + return bool(getattr(loss_cfg, "skip_invalid_translation_supervision", True)) + + +def _get_translation_valid_mask( + pred: dict, + batch: dict, + loss_cfg, + device: torch.device, +) -> torch.Tensor | None: + if not _skip_invalid_translation_supervision(loss_cfg): + return None + + mask = pred.get("translation_valid_mask") + if mask is not None: + base_mask = mask.to(device=device, dtype=torch.bool) + else: + pm_scale = pred.get("pm_scale") + if pm_scale is None: + return None + + scale = pm_scale.to(device=device).float() + if scale.dim() == 0: + scale = scale.reshape(1, 1) + elif scale.dim() == 1: + scale = scale.unsqueeze(-1) + else: + scale = scale.reshape(scale.shape[0], -1) + + min_valid_scale = _get_min_valid_pointmap_scale(loss_cfg) + finite = torch.isfinite(scale).all(dim=-1) + above_floor = (scale.abs() > min_valid_scale).all(dim=-1) + base_mask = finite & above_floor + + gt_trans_metric = batch.get("gt_trans_metric") + if gt_trans_metric is None: + return base_mask + + gt_trans_metric = gt_trans_metric.to(device=device).float().reshape(base_mask.shape[0], -1) + gt_finite = torch.isfinite(gt_trans_metric).all(dim=-1) + gt_norm = torch.linalg.norm(gt_trans_metric, dim=-1) + max_valid_norm = _get_max_valid_gt_translation_norm_m(loss_cfg) + if max_valid_norm > 0.0: + gt_reasonable = gt_norm <= max_valid_norm + else: + gt_reasonable = torch.ones_like(gt_finite, dtype=torch.bool) + return base_mask & gt_finite & gt_reasonable + + +def _masked_mean(values: torch.Tensor, valid_mask: torch.Tensor | None) -> torch.Tensor: + if valid_mask is None: + return values.mean() + if bool(valid_mask.any().item()): + return values[valid_mask].mean() + return values.new_zeros(()) + + +def _masked_finite_mean(values: torch.Tensor, valid_mask: torch.Tensor | None) -> torch.Tensor: + finite_mask = torch.isfinite(values) + if valid_mask is not None: + finite_mask = finite_mask & valid_mask.to(device=values.device, dtype=torch.bool) + if bool(finite_mask.any().item()): + return values[finite_mask].mean() + return values.new_zeros(()) + + + +#### Where to balanche translation and rotation + +def compute_flow_loss(pred: dict, batch: dict, loss_cfg, device: torch.device) -> dict[str, torch.Tensor]: + """Standard flow matching velocity MSE loss.""" + v_pred = pred["v_pred"] + v_target = pred["v_target"] + + loss_rot = torch.nn.functional.mse_loss( + v_pred["6drotation_normalized"], v_target["6drotation_normalized"], + ) + trans_diff = (v_pred["translation"].float() - v_target["translation"].float()).pow(2) + trans_diff = trans_diff.reshape(trans_diff.shape[0], -1).mean(dim=-1) + translation_valid_mask = _get_translation_valid_mask(pred, batch, loss_cfg, device) + loss_trans = _masked_mean(trans_diff, translation_valid_mask) + + rot_weight, trans_weight, scale_weight = _get_loss_weights(loss_cfg) + + # Scale velocity MSE (SSI log-scale latent). Supervised only when both v_pred + # and v_target carry the "scale" key and a scale_valid_mask is present. + loss_scale = loss_rot.new_zeros(()) + scale_valid_count = 0 + if "scale" in v_pred and "scale" in v_target: + scale_diff = (v_pred["scale"].float() - v_target["scale"].float()).pow(2) + scale_diff = scale_diff.reshape(scale_diff.shape[0], -1).mean(dim=-1) + scale_valid_mask = pred.get("scale_valid_mask") + if scale_valid_mask is not None: + scale_valid_mask = scale_valid_mask.to(device=device, dtype=torch.bool) + loss_scale = _masked_mean(scale_diff, scale_valid_mask) + scale_valid_count = ( + int(scale_valid_mask.sum().item()) + if scale_valid_mask is not None + else int(scale_diff.shape[0]) + ) + + total = rot_weight * loss_rot + trans_weight * loss_trans + scale_weight * loss_scale + + valid_count = ( + int(translation_valid_mask.sum().item()) + if translation_valid_mask is not None + else int(trans_diff.shape[0]) + ) + skipped_count = int(trans_diff.shape[0]) - valid_count + + return { + "loss": total, + "loss_rot": loss_rot.detach(), + "loss_trans": loss_trans.detach(), + "loss_scale": loss_scale.detach(), + "translation_valid_samples": torch.tensor(float(valid_count), device=device), + "translation_skipped_samples": torch.tensor(float(skipped_count), device=device), + "scale_valid_samples": torch.tensor(float(scale_valid_count), device=device), + } + + +### currently loss is L2 for both GT 6d and GT translation. +def compute_pose_loss(pred: dict, batch: dict, loss_cfg, device: torch.device) -> dict[str, torch.Tensor]: + """Pose loss on decoded predictions (used for inference/validation).""" + pred_rot_norm = pred["rotation"] + pred_t_ssi = pred["translation_ssi"] + pred_t_metric = pred["translation"] + pm_scale = pred["pm_scale"] + pm_shift = pred["pm_shift"] + + gt_rot_norm = batch["gt_rot6d"].to(device, non_blocking=True) + gt_t_metric = batch["gt_trans_metric"].to(device, non_blocking=True) + + # Rotation loss: L2 norm (not squared) in normalized 6D space. + rot_diff = pred_rot_norm.float() - gt_rot_norm.float() + loss_rot = torch.linalg.norm(rot_diff, dim=-1).mean() + + # Translation loss: L2 norm (not squared) in SSI-normalized camera space. + # t_ssi = (t_metric - shift) / scale per-axis. + pm_scale_safe = pm_scale.float().clamp_min(_get_min_valid_pointmap_scale(loss_cfg)) + gt_t_ssi = (gt_t_metric.float() - pm_shift.float()) / pm_scale_safe + trans_ssi_diff = pred_t_ssi.float() - gt_t_ssi.float() + translation_valid_mask = _get_translation_valid_mask(pred, batch, loss_cfg, device) + loss_trans = _masked_mean( + torch.linalg.norm(trans_ssi_diff, dim=-1), + translation_valid_mask, + ) + + rot_weight, trans_weight, _scale_weight = _get_loss_weights(loss_cfg) + total = rot_weight * loss_rot + trans_weight * loss_trans + + # Metrics for logging (no grad) + with torch.no_grad(): + pred_rotation = rot6d_to_matrix(denormalize_rot6d(pred_rot_norm)) + gt_rotation = rot6d_to_matrix(denormalize_rot6d(gt_rot_norm)) + rot_deg_per_sample = geodesic_distance_deg(pred_rotation.float(), gt_rotation.float()) + rot_deg = _masked_finite_mean(rot_deg_per_sample, valid_mask=None) + trans_diff = pred_t_metric.float() - gt_t_metric.float() + trans_l2_per_sample = torch.linalg.norm(trans_diff, dim=-1) + trans_l2 = _masked_mean( + trans_l2_per_sample, + torch.isfinite(trans_l2_per_sample) + if translation_valid_mask is None + else translation_valid_mask & torch.isfinite(trans_l2_per_sample), + ) + + valid_count = ( + int(translation_valid_mask.sum().item()) + if translation_valid_mask is not None + else int(pred_t_metric.shape[0]) + ) + skipped_count = int(pred_t_metric.shape[0]) - valid_count + + return { + "loss": total, + "loss_rot": loss_rot.detach(), + "loss_trans": loss_trans.detach(), + "metric_rot_deg": rot_deg.detach(), + "metric_trans_l2_m": trans_l2.detach(), + "translation_valid_samples": torch.tensor(float(valid_count), device=device), + "translation_skipped_samples": torch.tensor(float(skipped_count), device=device), + } diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/rotation_constants.py b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/rotation_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..caa1d6bc2a901b27d0cd890b99052f0cf190d85e --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/rotation_constants.py @@ -0,0 +1,32 @@ +"""Rotation 6D normalization constants. + +These match the ROTATION_6D_MEAN / ROTATION_6D_STD in +sam3d_objects.pipeline.inference_utils (pretrained SAM-3D checkpoint). +Duplicated here to avoid importing the heavy inference module (open3d, trimesh, etc.). +""" + +from __future__ import annotations + +import torch + +ROTATION_6D_MEAN = torch.tensor( + [ + -0.06366084883674913, + 0.008438224692279752, + 0.00017084786438302483, + 0.0007126610473540038, + -0.0030916726538816417, + 0.5166093753457688, + ] +) + +ROTATION_6D_STD = torch.tensor( + [ + 0.6656971967514863, + 0.6787012271867754, + 0.30345010594844524, + 0.4394504420678794, + 0.39817973931717104, + 0.6176286868761914, + ] +) diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/test_dataset_build.py b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/test_dataset_build.py new file mode 100644 index 0000000000000000000000000000000000000000..c61f6a255cde99d0760149afa2b1eb010bbd7e6d --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/test_dataset_build.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +"""Smoke-test the unified dataset: build manifest, exercise __getitem__ + collate, +and assert the record shapes/conventions. CPU only. + + /home/nvidia/miniconda3/envs/minha/bin/python \ + sam3d_objects/training/stage1_flow_lora/test_dataset_build.py \ + --pm-root /data/minha/output/unified_6d/robocasa/PickPlaceCounterToCabinet \ + --mv-root /data/minha/output/robocasa_mv_samples/PickPlaceCounterToCabinet \ + --dataset-root /lp-dev/jonghoon/robocasa_full/pickplace_target_human/PickPlaceCounterToCabinet +""" +import argparse +import os +import sys +from pathlib import Path + +os.environ.setdefault("LIDRA_SKIP_INIT", "1") # lightweight: skip heavy package init + +import numpy as np +import torch + +_REPO = Path(__file__).resolve().parents[3] +if str(_REPO) not in sys.path: + sys.path.insert(0, str(_REPO)) + +from sam3d_objects.training.stage1_flow_lora.dataset_unified import ( # noqa: E402 + UnifiedPoseDataset, build_manifest_robocasa, collate_unified) +from sam3d_objects.training.stage1_flow_lora.geometry import ( # noqa: E402 + denormalize_rot6d, rot6d_to_matrix) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--pm-root", required=True) + ap.add_argument("--mv-root", required=True) + ap.add_argument("--dataset-root", required=True) + ap.add_argument("--task", default="PickPlaceCounterToCabinet") + ap.add_argument("--stride", type=int, default=20) + args = ap.parse_args() + + samples = build_manifest_robocasa(args.pm_root, args.mv_root, args.dataset_root, + args.task, stride=args.stride) + print(f"[test] manifest: {len(samples)} samples; objects={sorted({s.object_name for s in samples})}") + assert samples, "empty manifest" + ds = UnifiedPoseDataset(samples) + + # exercise a few items + idxs = [0, len(ds) // 2, len(ds) - 1] + recs = [ds[i] for i in idxs] + r0 = recs[0] + print("[test] record keys:", sorted(r0.keys())) + H, W = r0["query_rgb"].shape[:2] + assert r0["query_rgb"].dtype == np.uint8 and r0["query_rgb"].ndim == 3 + assert r0["query_mask"].shape == (H, W) + assert tuple(r0["query_pointmap"].shape) == (H, W, 3), r0["query_pointmap"].shape + assert r0["gt_rot6d"].shape == (6,) + assert r0["gt_trans_metric"].shape == (3,) + assert r0["gt_scale"].shape == (1,) + assert r0["query_rgb"] is r0["first_frame_rgb"], "single-frame alias broken" + + # rot6d must round-trip to a valid rotation matrix + R = rot6d_to_matrix(denormalize_rot6d(torch.from_numpy(r0["gt_rot6d"])[None]))[0] + det = float(torch.det(R)); orth = float((R.T @ R - torch.eye(3)).abs().max()) + print(f"[test] rot6d->R det={det:.4f} orth_err={orth:.2e} t={np.round(r0['gt_trans_metric'],3)} s={float(r0['gt_scale'][0]):.4f}") + assert abs(det - 1.0) < 1e-3 and orth < 1e-3, "rot6d does not decode to SO(3)" + + # pointmap metric sanity: object-mask points distance ~ |t| + pm = r0["query_pointmap"].numpy(); m = r0["query_mask"].astype(bool) + if m.sum() > 0: + d = np.linalg.norm(pm[m & np.isfinite(pm).all(-1) & (pm[..., 2] != 0)], axis=1) + if len(d): + print(f"[test] obj-pointmap dist med={np.median(d):.3f}m vs |t|={np.linalg.norm(r0['gt_trans_metric']):.3f}m") + + # collate a batch + batch = collate_unified(recs) + assert batch["gt_rot6d"].shape == (len(recs), 6) + assert batch["gt_trans_metric"].shape == (len(recs), 3) + assert batch["gt_scale"].shape == (len(recs), 1) + assert len(batch["query_rgb"]) == len(recs) and len(batch["query_pointmap"]) == len(recs) + assert batch["object_label"].shape == (len(recs),) + print(f"[test] collate OK: gt_rot6d{tuple(batch['gt_rot6d'].shape)} " + f"gt_trans{tuple(batch['gt_trans_metric'].shape)} gt_scale{tuple(batch['gt_scale'].shape)} " + f"rgb_list={len(batch['query_rgb'])} datasets={set(batch['dataset'])}") + print("[test] PASS") + + +if __name__ == "__main__": + main() diff --git a/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/validation_infer.py b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/validation_infer.py new file mode 100644 index 0000000000000000000000000000000000000000..eb6ff1d9908e8534b1785dd72767cde61ee3d778 --- /dev/null +++ b/migrator/code/sam3d-for-6d/sam3d_objects/training/stage1_flow_lora/validation_infer.py @@ -0,0 +1,168 @@ +"""SAM-3D-inference validation for stage-1 training. + +Runs the *stock multi-step* SAM-3D inference (the same path as ``inference/``: +``prime_shape_cache`` → deterministic shape K/V from the first frame, then batched +``decode_query``) on a held-out s0 TEST subsample, and reports the **relative R/t +error** (frame-to-frame, eval.py convention). Per DDP rank shards whole sequences; +metrics are all-reduced. The first-frame (shape-K/V) noise seed is fixed so the +cached shape is deterministic across validation calls. +""" +from __future__ import annotations + +import contextlib +from collections import defaultdict + +import numpy as np +import torch + + +def _autocast(cfg, device): + """Match the training autocast so the backbone's FlashAttention sees bf16/fp16 + (it rejects fp32). Mirrors get_autocast_context used by the training forward.""" + amp = str(getattr(cfg.train, "amp_dtype", "bf16")).lower() + dt = {"bf16": torch.bfloat16, "fp16": torch.float16, "float16": torch.float16}.get(amp) + if dt is not None and device.type == "cuda": + return torch.autocast(device_type="cuda", dtype=dt) + return contextlib.nullcontext() + +from .dataset_unified import build_manifest_dexycb_jsonl, UnifiedPoseDataset +from .geometry import denormalize_rot6d, rot6d_to_matrix + +# reuse the inference package's relative-pose metric helpers +from inference.eval import geodesic_deg, relative_pose +from inference.pose_inference import _SeededNoise + +EVAL_BATCH = 32 + +_CACHE: dict[str, object] = {} + + +@contextlib.contextmanager +def _fixed_noise(ss_generator, seed: int): + """Force deterministic z_0 (fixed seed) for the prime + decode passes so the + shape K/V is reproducible. B is read from the latent shape at call time.""" + orig = ss_generator._generate_noise + + def _gen(x_shape, x_device): + B = int(next(iter(x_shape.values()))[0]) + return _SeededNoise(int(seed), list(range(B)))(x_shape, x_device) + + ss_generator._generate_noise = _gen + try: + yield + finally: + ss_generator._generate_noise = orig + + +def _load_sequences(cfg): + """Build + group the test-eval samples by sequence (cached across val calls).""" + key = cfg.data.dexycb_val_manifest_jsonl + if key not in _CACHE: + samples = build_manifest_dexycb_jsonl( + manifest_path=cfg.data.dexycb_val_manifest_jsonl, + dexycb_root=cfg.data.dexycb_root, pointmap_dir=cfg.data.pointmap_dir, + initial_pose_tmpl=cfg.data.initial_pose_file) + ds = UnifiedPoseDataset(samples) + seqs: dict[str, list[int]] = defaultdict(list) + for i, s in enumerate(samples): + seqs[f"{s.key}/ycb{s.ycb_id}"].append(i) + for k in seqs: + seqs[k].sort(key=lambda i: samples[i].frame_abs) + _CACHE[key] = (ds, sorted(seqs.keys()), seqs) + return _CACHE[key] + + +@torch.no_grad() +def run_sam3d_validation(model, cfg, dist_state) -> dict[str, float]: + """Batched SAM-3D inference on the test subset; returns mean relative R/t error.""" + inner = model.module if hasattr(model, "module") else model + was_training = inner.training + inner.eval() + if torch.cuda.is_available(): + torch.cuda.empty_cache() # free training reserve before the inference eval + + # Match PRODUCTION inference (inference/pose_inference.py): use the distilled + # multi-step shortcut sampler. The training default is no_shortcut=True (single + # Euler step) and leaves CFG strength set; the shortcut path uses NO CFG (only the + # cond slot), so this both makes val faithful to inference/eval.py AND avoids the + # never-trained unconditional branch. Restored after validation. + ss = inner.ss_generator + _orig_ns, _orig_steps = ss.no_shortcut, ss.inference_steps + steps = int(getattr(cfg.train, "eval_shortcut_steps", 2)) + inner.configure_shortcut_inference(no_shortcut=False, steps=steps) + + ds, seq_keys, seqs = _load_sequences(cfg) + my_keys = seq_keys[dist_state.rank::dist_state.world_size] + seed = int(cfg.train.eval_seed) + stride = int(getattr(cfg.train, "eval_rel_stride", 5)) + eval_batch = int(getattr(cfg.train, "eval_infer_batch", EVAL_BATCH)) + + rel_rot: list[float] = [] + rel_trans_cm: list[float] = [] + + for k in my_keys: + recs = [ds[i] for i in seqs[k]] + if len(recs) < 2: + continue + first = recs[0] + with _autocast(cfg, dist_state.device): + cond_first, _, _ = inner.embed_frame( + first["first_frame_rgb"], first["first_frame_mask"], first["first_frame_pointmap"]) + preds_R: list[np.ndarray] = [] + preds_t: list[np.ndarray] = [] + for start in range(0, len(recs), eval_batch): + chunk = recs[start:start + eval_batch] + with _autocast(cfg, dist_state.device): + conds, pms, pmsh = [], [], [] + for r in chunk: + c, ps, psh = inner.embed_frame(r["query_rgb"], r["query_mask"], r["query_pointmap"]) + conds.append(c); pms.append(ps); pmsh.append(psh) + cond_q = torch.cat(conds, 0) + pm_s = torch.cat(pms, 0); pm_sh = torch.cat(pmsh, 0) + B = cond_q.shape[0] + with _fixed_noise(inner.ss_generator, seed): + # Prime the shape K/V ONCE at batch 1 (all B query frames share the same + # first frame); the inject path broadcasts the batch-1 cache to B. This + # keeps the persistent 25-step cache 1x instead of Bx -> batch 32 fits. + inner.prime_shape_cache(cond_first) # deterministic shape K/V, batch 1 + out = inner.decode_query(cond_q, pm_s, pm_sh) # batched B, reuse shape K/V + inner.clear_shape_cache() + # out["rotation"] is the NORMALIZED 6D latent -> denormalize before Gram-Schmidt + # (matches inference/pose_inference.py; the affine doesn't survive the nonlinearity). + R = rot6d_to_matrix(denormalize_rot6d(out["rotation"])).float().cpu().numpy() # [B,3,3] + t = out["translation"].float().cpu().numpy() # [B,3] + for b in range(B): + preds_R.append(R[b]); preds_t.append(t[b]) + # GT (SAM-3D-canonical frame, same as preds) per frame; gt_rot6d is normalized too. + gt_R = [rot6d_to_matrix(denormalize_rot6d(torch.from_numpy(r["gt_rot6d"]).reshape(1, 6))).numpy()[0] + for r in recs] + gt_t = [np.asarray(r["gt_trans_metric"], np.float64) for r in recs] + # relative R/t error at stride + for i in range(0, len(recs) - stride, stride): + j = i + stride + Rp_rel, tp_rel = relative_pose(preds_R[i], preds_t[i], preds_R[j], preds_t[j]) + Rg_rel, tg_rel = relative_pose(gt_R[i], gt_t[i], gt_R[j], gt_t[j]) + rel_rot.append(geodesic_deg(Rp_rel, Rg_rel)) + rel_trans_cm.append(float(np.linalg.norm(tp_rel - tg_rel)) * 100.0) + + # all-reduce sums + counts across ranks + dev = dist_state.device + local = torch.tensor([sum(rel_rot), float(len(rel_rot)), + sum(rel_trans_cm), float(len(rel_trans_cm))], device=dev, dtype=torch.float64) + if getattr(dist_state, "distributed", False): + import torch.distributed as dist + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(local, op=dist.ReduceOp.SUM) + rr_sum, rr_n, tt_sum, tt_n = local.tolist() + + # restore training-default sampler (training forward() doesn't use these, but keep clean) + ss.no_shortcut, ss.inference_steps = _orig_ns, _orig_steps + if was_training: + inner.train() + if rr_n == 0: + return {} + return { + "val_rel_rot_deg": rr_sum / rr_n, + "val_rel_trans_cm": tt_sum / tt_n, + "val_pairs": rr_n, + }