340 lines
13 KiB
Python
Executable File
340 lines
13 KiB
Python
Executable File
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
|
|
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
|
|
|
from __future__ import annotations
|
|
|
|
from copy import deepcopy
|
|
|
|
import torch
|
|
|
|
from ultralytics.nn.modules.utils import inverse_sigmoid
|
|
from ultralytics.utils.ops import xywh2xyxy
|
|
|
|
from ..modules.sam import SAM2Model
|
|
from .geometry_encoders import Prompt
|
|
from .vl_combiner import SAM3VLBackbone
|
|
|
|
|
|
def _update_out(out, out_name, out_value, auxiliary=True, update_aux=True):
|
|
"""Helper function to update output dictionary with main and auxiliary outputs."""
|
|
out[out_name] = out_value[-1] if auxiliary else out_value
|
|
if auxiliary and update_aux:
|
|
if "aux_outputs" not in out:
|
|
out["aux_outputs"] = [{} for _ in range(len(out_value) - 1)]
|
|
assert len(out["aux_outputs"]) == len(out_value) - 1
|
|
for aux_output, aux_value in zip(out["aux_outputs"], out_value[:-1]):
|
|
aux_output[out_name] = aux_value
|
|
|
|
|
|
class SAM3SemanticModel(torch.nn.Module):
|
|
"""SAM3 model for semantic segmentation with vision-language backbone."""
|
|
|
|
def __init__(
|
|
self,
|
|
backbone: SAM3VLBackbone,
|
|
transformer,
|
|
input_geometry_encoder,
|
|
segmentation_head=None,
|
|
num_feature_levels=1,
|
|
o2m_mask_predict=True,
|
|
dot_prod_scoring=None,
|
|
use_instance_query: bool = True,
|
|
multimask_output: bool = True,
|
|
use_act_checkpoint_seg_head: bool = True,
|
|
matcher=None,
|
|
use_dot_prod_scoring=True,
|
|
supervise_joint_box_scores: bool = False, # only relevant if using presence token/score
|
|
detach_presence_in_joint_score: bool = False, # only relevant if using presence token/score
|
|
separate_scorer_for_instance: bool = False,
|
|
num_interactive_steps_val: int = 0,
|
|
):
|
|
"""Initialize the SAM3SemanticModel."""
|
|
super().__init__()
|
|
self.backbone = backbone
|
|
self.geometry_encoder = input_geometry_encoder
|
|
self.transformer = transformer
|
|
self.hidden_dim = transformer.d_model
|
|
self.num_feature_levels = num_feature_levels
|
|
self.segmentation_head = segmentation_head
|
|
|
|
self.o2m_mask_predict = o2m_mask_predict
|
|
|
|
self.dot_prod_scoring = dot_prod_scoring
|
|
self.use_act_checkpoint_seg_head = use_act_checkpoint_seg_head
|
|
self.matcher = matcher
|
|
|
|
self.num_interactive_steps_val = num_interactive_steps_val
|
|
self.use_dot_prod_scoring = use_dot_prod_scoring
|
|
|
|
if self.use_dot_prod_scoring:
|
|
assert dot_prod_scoring is not None
|
|
self.dot_prod_scoring = dot_prod_scoring
|
|
self.instance_dot_prod_scoring = None
|
|
if separate_scorer_for_instance:
|
|
self.instance_dot_prod_scoring = deepcopy(dot_prod_scoring)
|
|
else:
|
|
self.class_embed = torch.nn.Linear(self.hidden_dim, 1)
|
|
self.instance_class_embed = None
|
|
if separate_scorer_for_instance:
|
|
self.instance_class_embed = deepcopy(self.class_embed)
|
|
|
|
self.supervise_joint_box_scores = supervise_joint_box_scores
|
|
self.detach_presence_in_joint_score = detach_presence_in_joint_score
|
|
|
|
# verify the number of queries for O2O and O2M
|
|
num_o2o_static = self.transformer.decoder.num_queries
|
|
num_o2m_static = self.transformer.decoder.num_o2m_queries
|
|
assert num_o2m_static == (num_o2o_static if self.transformer.decoder.dac else 0)
|
|
self.dac = self.transformer.decoder.dac
|
|
|
|
self.use_instance_query = use_instance_query
|
|
self.multimask_output = multimask_output
|
|
|
|
self.text_embeddings = {}
|
|
self.names = []
|
|
|
|
def _encode_prompt(
|
|
self,
|
|
img_feats,
|
|
img_pos_embeds,
|
|
vis_feat_sizes,
|
|
geometric_prompt,
|
|
visual_prompt_embed=None,
|
|
visual_prompt_mask=None,
|
|
prev_mask_pred=None,
|
|
):
|
|
"""Encode the geometric and visual prompts."""
|
|
if prev_mask_pred is not None:
|
|
img_feats = [img_feats[-1] + prev_mask_pred]
|
|
# Encode geometry
|
|
geo_feats, geo_masks = self.geometry_encoder(
|
|
geo_prompt=geometric_prompt,
|
|
img_feats=img_feats,
|
|
img_sizes=vis_feat_sizes,
|
|
img_pos_embeds=img_pos_embeds,
|
|
)
|
|
if visual_prompt_embed is None:
|
|
visual_prompt_embed = torch.zeros((0, *geo_feats.shape[1:]), device=geo_feats.device)
|
|
visual_prompt_mask = torch.zeros(
|
|
(*geo_masks.shape[:-1], 0),
|
|
device=geo_masks.device,
|
|
dtype=geo_masks.dtype,
|
|
)
|
|
prompt = torch.cat([geo_feats, visual_prompt_embed], dim=0)
|
|
prompt_mask = torch.cat([geo_masks, visual_prompt_mask], dim=1)
|
|
return prompt, prompt_mask
|
|
|
|
def _run_encoder(
|
|
self,
|
|
img_feats,
|
|
img_pos_embeds,
|
|
vis_feat_sizes,
|
|
prompt,
|
|
prompt_mask,
|
|
encoder_extra_kwargs: dict | None = None,
|
|
):
|
|
"""Run the transformer encoder."""
|
|
# Run the encoder
|
|
# make a copy of the image feature lists since the encoder may modify these lists in-place
|
|
memory = self.transformer.encoder(
|
|
src=img_feats.copy(),
|
|
src_key_padding_mask=None,
|
|
src_pos=img_pos_embeds.copy(),
|
|
prompt=prompt,
|
|
prompt_key_padding_mask=prompt_mask,
|
|
feat_sizes=vis_feat_sizes,
|
|
encoder_extra_kwargs=encoder_extra_kwargs,
|
|
)
|
|
encoder_out = {
|
|
# encoded image features
|
|
"encoder_hidden_states": memory["memory"],
|
|
"pos_embed": memory["pos_embed"],
|
|
"padding_mask": memory["padding_mask"],
|
|
"spatial_shapes": memory["spatial_shapes"],
|
|
"valid_ratios": memory["valid_ratios"],
|
|
"vis_feat_sizes": vis_feat_sizes,
|
|
# encoded text features (or other prompts)
|
|
"prompt_before_enc": prompt,
|
|
"prompt_after_enc": memory.get("memory_text", prompt),
|
|
"prompt_mask": prompt_mask,
|
|
}
|
|
return encoder_out
|
|
|
|
def _run_decoder(
|
|
self,
|
|
pos_embed,
|
|
memory,
|
|
src_mask,
|
|
out,
|
|
prompt,
|
|
prompt_mask,
|
|
encoder_out,
|
|
):
|
|
"""Run the transformer decoder."""
|
|
bs = memory.shape[1]
|
|
query_embed = self.transformer.decoder.query_embed.weight
|
|
tgt = query_embed.unsqueeze(1).repeat(1, bs, 1)
|
|
|
|
hs, reference_boxes, dec_presence_out, _ = self.transformer.decoder(
|
|
tgt=tgt,
|
|
memory=memory,
|
|
memory_key_padding_mask=src_mask,
|
|
pos=pos_embed,
|
|
reference_boxes=None,
|
|
spatial_shapes=encoder_out["spatial_shapes"],
|
|
valid_ratios=encoder_out["valid_ratios"],
|
|
tgt_mask=None,
|
|
memory_text=prompt,
|
|
text_attention_mask=prompt_mask,
|
|
apply_dac=False,
|
|
)
|
|
hs = hs.transpose(1, 2) # seq-first to batch-first
|
|
reference_boxes = reference_boxes.transpose(1, 2) # seq-first to batch-first
|
|
if dec_presence_out is not None:
|
|
# seq-first to batch-first
|
|
dec_presence_out = dec_presence_out.transpose(1, 2)
|
|
self._update_scores_and_boxes(
|
|
out,
|
|
hs,
|
|
reference_boxes,
|
|
prompt,
|
|
prompt_mask,
|
|
dec_presence_out=dec_presence_out,
|
|
)
|
|
return out, hs
|
|
|
|
def _update_scores_and_boxes(
|
|
self,
|
|
out,
|
|
hs,
|
|
reference_boxes,
|
|
prompt,
|
|
prompt_mask,
|
|
dec_presence_out=None,
|
|
is_instance_prompt=False,
|
|
):
|
|
"""Update output dict with class scores and box predictions."""
|
|
num_o2o = hs.size(2)
|
|
# score prediction
|
|
if self.use_dot_prod_scoring:
|
|
dot_prod_scoring_head = self.dot_prod_scoring
|
|
if is_instance_prompt and self.instance_dot_prod_scoring is not None:
|
|
dot_prod_scoring_head = self.instance_dot_prod_scoring
|
|
outputs_class = dot_prod_scoring_head(hs, prompt, prompt_mask)
|
|
else:
|
|
class_embed_head = self.class_embed
|
|
if is_instance_prompt and self.instance_class_embed is not None:
|
|
class_embed_head = self.instance_class_embed
|
|
outputs_class = class_embed_head(hs)
|
|
|
|
# box prediction
|
|
box_head = self.transformer.decoder.bbox_embed
|
|
if is_instance_prompt and self.transformer.decoder.instance_bbox_embed is not None:
|
|
box_head = self.transformer.decoder.instance_bbox_embed
|
|
anchor_box_offsets = box_head(hs)
|
|
reference_boxes_inv_sig = inverse_sigmoid(reference_boxes)
|
|
outputs_coord = (reference_boxes_inv_sig + anchor_box_offsets).sigmoid()
|
|
outputs_boxes_xyxy = xywh2xyxy(outputs_coord)
|
|
|
|
if dec_presence_out is not None:
|
|
_update_out(out, "presence_logit_dec", dec_presence_out, update_aux=False)
|
|
|
|
if self.supervise_joint_box_scores:
|
|
assert dec_presence_out is not None
|
|
prob_dec_presence_out = dec_presence_out.clone().sigmoid()
|
|
if self.detach_presence_in_joint_score:
|
|
prob_dec_presence_out = prob_dec_presence_out.detach()
|
|
|
|
outputs_class = inverse_sigmoid(outputs_class.sigmoid() * prob_dec_presence_out.unsqueeze(2)).clamp(
|
|
min=-10.0, max=10.0
|
|
)
|
|
|
|
_update_out(out, "pred_logits", outputs_class[:, :, :num_o2o], update_aux=False)
|
|
_update_out(out, "pred_boxes", outputs_coord[:, :, :num_o2o], update_aux=False)
|
|
_update_out(out, "pred_boxes_xyxy", outputs_boxes_xyxy[:, :, :num_o2o], update_aux=False)
|
|
|
|
def _run_segmentation_heads(
|
|
self,
|
|
out,
|
|
backbone_out,
|
|
encoder_hidden_states,
|
|
prompt,
|
|
prompt_mask,
|
|
hs,
|
|
):
|
|
"""Run segmentation heads and get masks."""
|
|
if self.segmentation_head is not None:
|
|
num_o2o = hs.size(2)
|
|
obj_queries = hs if self.o2m_mask_predict else hs[:, :, :num_o2o]
|
|
seg_head_outputs = self.segmentation_head(
|
|
backbone_feats=backbone_out["backbone_fpn"],
|
|
obj_queries=obj_queries,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
prompt=prompt,
|
|
prompt_mask=prompt_mask,
|
|
)
|
|
for k, v in seg_head_outputs.items():
|
|
if k in self.segmentation_head.instance_keys:
|
|
_update_out(out, k, v[:, :num_o2o], auxiliary=False)
|
|
else:
|
|
out[k] = v
|
|
else:
|
|
backbone_out.pop("backbone_fpn", None)
|
|
|
|
def forward_grounding(
|
|
self, backbone_out: dict[str, torch.Tensor], text_ids: torch.Tensor, geometric_prompt: Prompt = None
|
|
):
|
|
"""Forward pass for grounding (detection + segmentation) given input images and text."""
|
|
backbone_out, img_feats, img_pos_embeds, vis_feat_sizes = SAM2Model._prepare_backbone_features(
|
|
self, backbone_out, batch=len(text_ids)
|
|
)
|
|
backbone_out.update({k: v for k, v in self.text_embeddings.items()})
|
|
with torch.profiler.record_function("SAM3Image._encode_prompt"):
|
|
prompt, prompt_mask = self._encode_prompt(img_feats, img_pos_embeds, vis_feat_sizes, geometric_prompt)
|
|
# index text features (note that regardless of early or late fusion, the batch size of
|
|
# `txt_feats` is always the number of *prompts* in the encoder)
|
|
txt_feats = backbone_out["language_features"][:, text_ids]
|
|
txt_masks = backbone_out["language_mask"][text_ids]
|
|
# encode text
|
|
prompt = torch.cat([txt_feats, prompt], dim=0)
|
|
prompt_mask = torch.cat([txt_masks, prompt_mask], dim=1)
|
|
|
|
# Run the encoder
|
|
with torch.profiler.record_function("SAM3Image._run_encoder"):
|
|
encoder_out = self._run_encoder(img_feats, img_pos_embeds, vis_feat_sizes, prompt, prompt_mask)
|
|
out = {"backbone_out": backbone_out}
|
|
|
|
# Run the decoder
|
|
with torch.profiler.record_function("SAM3Image._run_decoder"):
|
|
out, hs = self._run_decoder(
|
|
memory=encoder_out["encoder_hidden_states"],
|
|
pos_embed=encoder_out["pos_embed"],
|
|
src_mask=encoder_out["padding_mask"],
|
|
out=out,
|
|
prompt=prompt,
|
|
prompt_mask=prompt_mask,
|
|
encoder_out=encoder_out,
|
|
)
|
|
|
|
# Run segmentation heads
|
|
with torch.profiler.record_function("SAM3Image._run_segmentation_heads"):
|
|
self._run_segmentation_heads(
|
|
out=out,
|
|
backbone_out=backbone_out,
|
|
encoder_hidden_states=encoder_out["encoder_hidden_states"],
|
|
prompt=prompt,
|
|
prompt_mask=prompt_mask,
|
|
hs=hs,
|
|
)
|
|
return out
|
|
|
|
def set_classes(self, text: list[str]):
|
|
"""Set the text embeddings for the given class names."""
|
|
self.text_embeddings = self.backbone.forward_text(text)
|
|
self.names = text
|
|
|
|
def set_imgsz(self, imgsz: tuple[int, int]):
|
|
"""Set the image size for the model."""
|
|
self.backbone.set_imgsz(imgsz)
|