| """VoxTell 3D grounding encoder for standalone RepConTUR inference. |
| |
| Adapted from the training-time VoxTellTrainRunner: the architecture is built |
| from an explicit ``arch_kwargs`` dict (stored in the HF config) instead of a |
| ``plans.json`` file, and no checkpoint is loaded at construction — weights come |
| from the RepConTUR safetensors. |
| """ |
|
|
| import pydoc |
| from typing import Tuple |
|
|
| import torch |
| from torch import nn |
| from einops import rearrange |
|
|
| from .voxtell_model import VoxTellModel |
|
|
|
|
| class VoxTellGroundingEncoder(nn.Module): |
| """VoxTell ResidualEncoder + text-prompted transformer decoder. |
| |
| The LLM ``[SEG]`` embeddings (projected to ``query_dim`` by the parent's |
| ``text_hidden_fcs``) are used directly as decoder queries, bypassing |
| VoxTell's original text embedding model. |
| """ |
|
|
| def __init__( |
| self, |
| arch_kwargs: dict, |
| kw_requires_import=('conv_op', 'norm_op', 'dropout_op', 'nonlin'), |
| decoder_layer: int = 4, |
| text_embedding_dim: int = 2560, |
| num_maskformer_stages: int = 5, |
| num_heads: int = 32, |
| query_dim: int = 2048, |
| ): |
| super().__init__() |
| self.decoder_layer = decoder_layer |
| self.hidden_dim = query_dim |
|
|
| resolved = dict(arch_kwargs) |
| for key in kw_requires_import: |
| if resolved.get(key) is not None: |
| resolved[key] = pydoc.locate(resolved[key]) |
|
|
| self.voxtell_model = VoxTellModel( |
| input_channels=1, |
| **resolved, |
| decoder_layer=decoder_layer, |
| text_embedding_dim=text_embedding_dim, |
| num_maskformer_stages=num_maskformer_stages, |
| num_heads=num_heads, |
| query_dim=query_dim, |
| project_to_decoder_hidden_dim=query_dim, |
| deep_supervision=False, |
| ) |
|
|
| self.voxtell_model.encoder.requires_grad_(False) |
| self.voxtell_model.project_text_embed.requires_grad_(False) |
|
|
| def forward_with_language( |
| self, |
| volume: torch.Tensor, |
| language_embd: torch.Tensor, |
| num_objs: int = 1, |
| ) -> torch.Tensor: |
| """Encode a 3D volume and decode with language queries. |
| |
| Args: |
| volume: (B, 1, D, H, W) preprocessed volume. |
| language_embd: (B * num_objs, 1, query_dim) projected LLM embeddings. |
| num_objs: objects per sample. |
| Returns: |
| (B * num_objs, D, H, W) segmentation logits. |
| """ |
| batch_size = volume.shape[0] |
|
|
| with torch.autocast(device_type='cuda', dtype=torch.bfloat16): |
| skips = self.voxtell_model.encoder(volume) |
|
|
| selected_feature = skips[self.decoder_layer] |
| bottleneck_embed = rearrange(selected_feature, 'b c d h w -> b h w d c') |
| bottleneck_embed = self.voxtell_model.project_bottleneck_embed(bottleneck_embed) |
| bottleneck_embed = rearrange(bottleneck_embed, 'b h w d c -> (h w d) b c') |
|
|
| text_embed = language_embd.squeeze(1) |
| text_embed = text_embed.reshape(batch_size, num_objs, -1) |
| text_embed = rearrange(text_embed, 'b n d -> n b d') |
|
|
| mask_embedding, _ = self.voxtell_model.transformer_decoder( |
| tgt=text_embed, |
| memory=bottleneck_embed, |
| pos=self.voxtell_model.pos_embed, |
| memory_key_padding_mask=None, |
| ) |
| mask_embedding = rearrange(mask_embedding, 'n b d -> b n d') |
|
|
| mask_embeddings = [ |
| proj(mask_embedding) |
| for proj in self.voxtell_model.project_to_decoder_channels |
| ] |
|
|
| outs = [] |
| for prompt_idx in range(num_objs): |
| prompt_embeds = [m[:, prompt_idx:prompt_idx + 1] for m in mask_embeddings] |
| outs.append(self.voxtell_model.decoder(skips, prompt_embeds)) |
|
|
| outs = [torch.cat(scale_outs, dim=1) for scale_outs in zip(*outs)] |
| pred_masks = outs[0] |
|
|
| pred_masks = pred_masks.reshape(batch_size * num_objs, *pred_masks.shape[2:]) |
| return pred_masks |
|
|
| @torch.inference_mode() |
| def predict_sliding_window( |
| self, |
| volume: torch.Tensor, |
| language_embd: torch.Tensor, |
| num_objs: int = 1, |
| tile_step_size: float = 0.5, |
| patch_size: Tuple[int, int, int] = (192, 192, 192), |
| ) -> torch.Tensor: |
| """Sliding-window inference over a full volume with Gaussian stitching. |
| |
| Args: |
| volume: (1, 1, D, H, W) full preprocessed volume. |
| language_embd: (num_objs, 1, query_dim) shared language queries. |
| num_objs: number of objects. |
| tile_step_size: step fraction of patch size. |
| patch_size: inference patch size. |
| Returns: |
| (num_objs, D, H, W) logits at full resolution. |
| """ |
| from acvl_utils.cropping_and_padding.padding import pad_nd_image |
| from nnunetv2.inference.sliding_window_prediction import ( |
| compute_gaussian, |
| compute_steps_for_sliding_window, |
| ) |
|
|
| device = volume.device |
| data = volume.squeeze(0) |
|
|
| data, slicer_revert = pad_nd_image( |
| data, patch_size, 'constant', {'value': 0}, True, None, |
| ) |
| spatial_shape = data.shape[1:] |
|
|
| steps = compute_steps_for_sliding_window(spatial_shape, patch_size, tile_step_size) |
| slicers = [] |
| for sx in steps[0]: |
| for sy in steps[1]: |
| for sz in steps[2]: |
| slicers.append(( |
| slice(None), |
| slice(sx, sx + patch_size[0]), |
| slice(sy, sy + patch_size[1]), |
| slice(sz, sz + patch_size[2]), |
| )) |
|
|
| predicted_logits = torch.zeros( |
| (num_objs, *spatial_shape), dtype=torch.float32, device=device, |
| ) |
| n_predictions = torch.zeros(spatial_shape, dtype=torch.float32, device=device) |
|
|
| gaussian = compute_gaussian( |
| tuple(patch_size), sigma_scale=1.0 / 8, |
| value_scaling_factor=10, device=device, |
| ).float() |
|
|
| for slicer in slicers: |
| patch = data[slicer].unsqueeze(0).contiguous() |
| pred = self.forward_with_language(patch, language_embd, num_objs) |
| pred = pred.float() * gaussian |
|
|
| spatial_slicer = slicer[1:] |
| out_slicer = (slice(None), *spatial_slicer) |
| predicted_logits[out_slicer] += pred |
| n_predictions[spatial_slicer] += gaussian |
|
|
| predicted_logits /= n_predictions |
| predicted_logits = predicted_logits[(slice(None), *slicer_revert[1:])] |
| return predicted_logits |
|
|
