File size: 4,085 Bytes
7eee2d7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """SAM 3's text-prompted segmentation, split into the three graphs the browser runs.
The split follows `Sam3Model.forward` (Hugging Face Transformers, `models/sam3/modeling_sam3.py`)
for a text prompt with no box prompts:
- **vision**: the Perception Encoder ViT and its neck, run once per image. Its outputs are the
first three feature levels (288², 144² and 72², 256 channels); `Sam3Model` drops the fourth.
- **text**: the CLIP text encoder and `text_projection`, run once per phrase.
- **detector**: the DETR encoder and decoder, the box head, the dot-product scoring and the mask
decoder, run once per phrase over the vision graph's features.
The 72² level's sine position encoding depends only on the input size, which is fixed at
1008 × 1008, so the detector carries it as a constant instead of taking it as an input.
"""
import torch
from torch import nn
from transformers.models.sam3.modeling_sam3 import box_cxcywh_to_xyxy, inverse_sigmoid
IMAGE_SIZE = 1008
TEXT_TOKENS = 32
class Sam3Vision(nn.Module):
def __init__(self, model):
super().__init__()
self.vision_encoder = model.vision_encoder
def forward(self, pixel_values: torch.Tensor):
levels = self.vision_encoder(pixel_values).fpn_hidden_states
return levels[0], levels[1], levels[2]
class Sam3Text(nn.Module):
def __init__(self, model):
super().__init__()
self.text_encoder = model.text_encoder
self.text_projection = model.text_projection
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor):
hidden = self.text_encoder(
input_ids=input_ids, attention_mask=attention_mask, return_dict=True
).last_hidden_state
return self.text_projection(hidden)
class Sam3Detector(nn.Module):
def __init__(self, model, position_encoding: torch.Tensor):
super().__init__()
self.detr_encoder = model.detr_encoder
self.detr_decoder = model.detr_decoder
self.dot_product_scoring = model.dot_product_scoring
self.mask_decoder = model.mask_decoder
self.register_buffer("position_encoding", position_encoding)
def forward(
self,
features_288: torch.Tensor,
features_144: torch.Tensor,
features_72: torch.Tensor,
text_features: torch.Tensor,
text_mask: torch.Tensor,
):
mask = text_mask.bool()
encoder = self.detr_encoder(
vision_features=[features_72],
text_features=text_features,
vision_pos_embeds=[self.position_encoding],
text_mask=mask,
)
decoder = self.detr_decoder(
vision_features=encoder.last_hidden_state,
text_features=encoder.text_features,
vision_pos_encoding=encoder.pos_embeds_flattened,
text_mask=mask,
spatial_shapes=encoder.spatial_shapes,
)
hidden = decoder.intermediate_hidden_states[-1:]
offsets = self.detr_decoder.box_head(hidden)
boxes = (inverse_sigmoid(decoder.reference_boxes[-1:]) + offsets).sigmoid()
logits = self.dot_product_scoring(
decoder_hidden_states=hidden,
text_features=encoder.text_features,
text_mask=mask,
).squeeze(-1)
masks = self.mask_decoder(
decoder_queries=hidden[0],
backbone_features=[features_288, features_144, features_72],
encoder_hidden_states=encoder.last_hidden_state,
prompt_features=text_features,
prompt_mask=mask,
)
return (
logits[0],
box_cxcywh_to_xyxy(boxes)[0],
decoder.presence_logits[-1],
masks.pred_masks,
masks.semantic_seg,
)
def position_encoding_72(model) -> torch.Tensor:
"""The 72² level's position encoding for a 1008² image, as the vision model computes it."""
with torch.no_grad():
dummy = torch.zeros(1, 3, IMAGE_SIZE, IMAGE_SIZE)
return model.vision_encoder(dummy).fpn_position_encoding[2].clone()
|