Add files via upload
This commit is contained in:
AI Lab
Binary file not shown.
@@ -0,0 +1,54 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
import logging
|
||||
import os
|
||||
|
||||
LOG_LEVELS = {
|
||||
"DEBUG": logging.DEBUG,
|
||||
"INFO": logging.INFO,
|
||||
"WARNING": logging.WARNING,
|
||||
"ERROR": logging.ERROR,
|
||||
"CRITICAL": logging.CRITICAL,
|
||||
}
|
||||
|
||||
|
||||
class ColoredFormatter(logging.Formatter):
|
||||
"""A command line formatter with different colors for each level."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
reset = "\033[0m"
|
||||
colors = {
|
||||
logging.DEBUG: f"{reset}\033[36m", # cyan,
|
||||
logging.INFO: f"{reset}\033[32m", # green
|
||||
logging.WARNING: f"{reset}\033[33m", # yellow
|
||||
logging.ERROR: f"{reset}\033[31m", # red
|
||||
logging.CRITICAL: f"{reset}\033[35m", # magenta
|
||||
}
|
||||
fmt_str = "{color}%(levelname)s %(asctime)s %(process)d %(filename)s:%(lineno)4d:{reset} %(message)s"
|
||||
self.formatters = {
|
||||
level: logging.Formatter(fmt_str.format(color=color, reset=reset))
|
||||
for level, color in colors.items()
|
||||
}
|
||||
self.default_formatter = self.formatters[logging.INFO]
|
||||
|
||||
def format(self, record):
|
||||
formatter = self.formatters.get(record.levelno, self.default_formatter)
|
||||
return formatter.format(record)
|
||||
|
||||
|
||||
def get_logger(name, level=logging.INFO):
|
||||
"""A command line logger."""
|
||||
if "LOG_LEVEL" in os.environ:
|
||||
level = os.environ["LOG_LEVEL"].upper()
|
||||
assert (
|
||||
level in LOG_LEVELS
|
||||
), f"Invalid LOG_LEVEL: {level}, must be one of {list(LOG_LEVELS.keys())}"
|
||||
level = LOG_LEVELS[level]
|
||||
logger = logging.getLogger(name)
|
||||
logger.setLevel(level)
|
||||
logger.propagate = False
|
||||
ch = logging.StreamHandler()
|
||||
ch.setLevel(level)
|
||||
ch.setFormatter(ColoredFormatter())
|
||||
logger.addHandler(ch)
|
||||
return logger
|
||||
@@ -0,0 +1,793 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from huggingface_hub import hf_hub_download
|
||||
from iopath.common.file_io import g_pathmgr
|
||||
from sam3.model.decoder import (
|
||||
TransformerDecoder,
|
||||
TransformerDecoderLayer,
|
||||
TransformerDecoderLayerv2,
|
||||
TransformerEncoderCrossAttention,
|
||||
)
|
||||
from sam3.model.encoder import TransformerEncoderFusion, TransformerEncoderLayer
|
||||
from sam3.model.geometry_encoders import SequenceGeometryEncoder
|
||||
from sam3.model.maskformer_segmentation import PixelDecoder, UniversalSegmentationHead
|
||||
from sam3.model.memory import (
|
||||
CXBlock,
|
||||
SimpleFuser,
|
||||
SimpleMaskDownSampler,
|
||||
SimpleMaskEncoder,
|
||||
)
|
||||
from sam3.model.model_misc import (
|
||||
DotProductScoring,
|
||||
MLP,
|
||||
MultiheadAttentionWrapper as MultiheadAttention,
|
||||
TransformerWrapper,
|
||||
)
|
||||
from sam3.model.necks import Sam3DualViTDetNeck
|
||||
from sam3.model.position_encoding import PositionEmbeddingSine
|
||||
from sam3.model.sam1_task_predictor import SAM3InteractiveImagePredictor
|
||||
from sam3.model.sam3_image import Sam3Image, Sam3ImageOnVideoMultiGPU
|
||||
from sam3.model.sam3_tracking_predictor import Sam3TrackerPredictor
|
||||
from sam3.model.sam3_video_inference import Sam3VideoInferenceWithInstanceInteractivity
|
||||
from sam3.model.sam3_video_predictor import Sam3VideoPredictorMultiGPU
|
||||
from sam3.model.text_encoder_ve import VETextEncoder
|
||||
from sam3.model.tokenizer_ve import SimpleTokenizer
|
||||
from sam3.model.vitdet import ViT
|
||||
from sam3.model.vl_combiner import SAM3VLBackbone
|
||||
from sam3.sam.transformer import RoPEAttention
|
||||
|
||||
|
||||
# Setup TensorFloat-32 for Ampere GPUs if available
|
||||
def _setup_tf32() -> None:
|
||||
"""Enable TensorFloat-32 for Ampere GPUs if available."""
|
||||
if torch.cuda.is_available():
|
||||
device_props = torch.cuda.get_device_properties(0)
|
||||
if device_props.major >= 8:
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
|
||||
|
||||
_setup_tf32()
|
||||
|
||||
|
||||
def _create_position_encoding(precompute_resolution=None):
|
||||
"""Create position encoding for visual backbone."""
|
||||
return PositionEmbeddingSine(
|
||||
num_pos_feats=256,
|
||||
normalize=True,
|
||||
scale=None,
|
||||
temperature=10000,
|
||||
precompute_resolution=precompute_resolution,
|
||||
)
|
||||
|
||||
|
||||
def _create_vit_backbone(compile_mode=None):
|
||||
"""Create ViT backbone for visual feature extraction."""
|
||||
return ViT(
|
||||
img_size=1008,
|
||||
pretrain_img_size=336,
|
||||
patch_size=14,
|
||||
embed_dim=1024,
|
||||
depth=32,
|
||||
num_heads=16,
|
||||
mlp_ratio=4.625,
|
||||
norm_layer="LayerNorm",
|
||||
drop_path_rate=0.1,
|
||||
qkv_bias=True,
|
||||
use_abs_pos=True,
|
||||
tile_abs_pos=True,
|
||||
global_att_blocks=(7, 15, 23, 31),
|
||||
rel_pos_blocks=(),
|
||||
use_rope=True,
|
||||
use_interp_rope=True,
|
||||
window_size=24,
|
||||
pretrain_use_cls_token=True,
|
||||
retain_cls_token=False,
|
||||
ln_pre=True,
|
||||
ln_post=False,
|
||||
return_interm_layers=False,
|
||||
bias_patch_embed=False,
|
||||
compile_mode=compile_mode,
|
||||
)
|
||||
|
||||
|
||||
def _create_vit_neck(position_encoding, vit_backbone, enable_inst_interactivity=False):
|
||||
"""Create ViT neck for feature pyramid."""
|
||||
return Sam3DualViTDetNeck(
|
||||
position_encoding=position_encoding,
|
||||
d_model=256,
|
||||
scale_factors=[4.0, 2.0, 1.0, 0.5],
|
||||
trunk=vit_backbone,
|
||||
add_sam2_neck=enable_inst_interactivity,
|
||||
)
|
||||
|
||||
|
||||
def _create_vl_backbone(vit_neck, text_encoder):
|
||||
"""Create visual-language backbone."""
|
||||
return SAM3VLBackbone(visual=vit_neck, text=text_encoder, scalp=1)
|
||||
|
||||
|
||||
def _create_transformer_encoder() -> TransformerEncoderFusion:
|
||||
"""Create transformer encoder with its layer."""
|
||||
encoder_layer = TransformerEncoderLayer(
|
||||
activation="relu",
|
||||
d_model=256,
|
||||
dim_feedforward=2048,
|
||||
dropout=0.1,
|
||||
pos_enc_at_attn=True,
|
||||
pos_enc_at_cross_attn_keys=False,
|
||||
pos_enc_at_cross_attn_queries=False,
|
||||
pre_norm=True,
|
||||
self_attention=MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0.1,
|
||||
embed_dim=256,
|
||||
batch_first=True,
|
||||
),
|
||||
cross_attention=MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0.1,
|
||||
embed_dim=256,
|
||||
batch_first=True,
|
||||
),
|
||||
)
|
||||
|
||||
encoder = TransformerEncoderFusion(
|
||||
layer=encoder_layer,
|
||||
num_layers=6,
|
||||
d_model=256,
|
||||
num_feature_levels=1,
|
||||
frozen=False,
|
||||
use_act_checkpoint=True,
|
||||
add_pooled_text_to_img_feat=False,
|
||||
pool_text_with_mask=True,
|
||||
)
|
||||
return encoder
|
||||
|
||||
|
||||
def _create_transformer_decoder() -> TransformerDecoder:
|
||||
"""Create transformer decoder with its layer."""
|
||||
decoder_layer = TransformerDecoderLayer(
|
||||
activation="relu",
|
||||
d_model=256,
|
||||
dim_feedforward=2048,
|
||||
dropout=0.1,
|
||||
cross_attention=MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0.1,
|
||||
embed_dim=256,
|
||||
),
|
||||
n_heads=8,
|
||||
use_text_cross_attention=True,
|
||||
)
|
||||
|
||||
decoder = TransformerDecoder(
|
||||
layer=decoder_layer,
|
||||
num_layers=6,
|
||||
num_queries=200,
|
||||
return_intermediate=True,
|
||||
box_refine=True,
|
||||
num_o2m_queries=0,
|
||||
dac=True,
|
||||
boxRPB="log",
|
||||
d_model=256,
|
||||
frozen=False,
|
||||
interaction_layer=None,
|
||||
dac_use_selfatt_ln=True,
|
||||
resolution=1008,
|
||||
stride=14,
|
||||
use_act_checkpoint=True,
|
||||
presence_token=True,
|
||||
)
|
||||
return decoder
|
||||
|
||||
|
||||
def _create_dot_product_scoring():
|
||||
"""Create dot product scoring module."""
|
||||
prompt_mlp = MLP(
|
||||
input_dim=256,
|
||||
hidden_dim=2048,
|
||||
output_dim=256,
|
||||
num_layers=2,
|
||||
dropout=0.1,
|
||||
residual=True,
|
||||
out_norm=nn.LayerNorm(256),
|
||||
)
|
||||
return DotProductScoring(d_model=256, d_proj=256, prompt_mlp=prompt_mlp)
|
||||
|
||||
|
||||
def _create_segmentation_head(compile_mode=None):
|
||||
"""Create segmentation head with pixel decoder."""
|
||||
pixel_decoder = PixelDecoder(
|
||||
num_upsampling_stages=3,
|
||||
interpolation_mode="nearest",
|
||||
hidden_dim=256,
|
||||
compile_mode=compile_mode,
|
||||
)
|
||||
|
||||
cross_attend_prompt = MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0,
|
||||
embed_dim=256,
|
||||
)
|
||||
|
||||
segmentation_head = UniversalSegmentationHead(
|
||||
hidden_dim=256,
|
||||
upsampling_stages=3,
|
||||
aux_masks=False,
|
||||
presence_head=False,
|
||||
dot_product_scorer=None,
|
||||
act_ckpt=True,
|
||||
cross_attend_prompt=cross_attend_prompt,
|
||||
pixel_decoder=pixel_decoder,
|
||||
)
|
||||
return segmentation_head
|
||||
|
||||
|
||||
def _create_geometry_encoder():
|
||||
"""Create geometry encoder with all its components."""
|
||||
# Create position encoding for geometry encoder
|
||||
geo_pos_enc = _create_position_encoding()
|
||||
# Create CX block for fuser
|
||||
cx_block = CXBlock(
|
||||
dim=256,
|
||||
kernel_size=7,
|
||||
padding=3,
|
||||
layer_scale_init_value=1.0e-06,
|
||||
use_dwconv=True,
|
||||
)
|
||||
# Create geometry encoder layer
|
||||
geo_layer = TransformerEncoderLayer(
|
||||
activation="relu",
|
||||
d_model=256,
|
||||
dim_feedforward=2048,
|
||||
dropout=0.1,
|
||||
pos_enc_at_attn=False,
|
||||
pre_norm=True,
|
||||
self_attention=MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0.1,
|
||||
embed_dim=256,
|
||||
batch_first=False,
|
||||
),
|
||||
pos_enc_at_cross_attn_queries=False,
|
||||
pos_enc_at_cross_attn_keys=True,
|
||||
cross_attention=MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0.1,
|
||||
embed_dim=256,
|
||||
batch_first=False,
|
||||
),
|
||||
)
|
||||
|
||||
# Create geometry encoder
|
||||
input_geometry_encoder = SequenceGeometryEncoder(
|
||||
pos_enc=geo_pos_enc,
|
||||
encode_boxes_as_points=False,
|
||||
points_direct_project=True,
|
||||
points_pool=True,
|
||||
points_pos_enc=True,
|
||||
boxes_direct_project=True,
|
||||
boxes_pool=True,
|
||||
boxes_pos_enc=True,
|
||||
d_model=256,
|
||||
num_layers=3,
|
||||
layer=geo_layer,
|
||||
use_act_ckpt=True,
|
||||
add_cls=True,
|
||||
add_post_encode_proj=True,
|
||||
)
|
||||
return input_geometry_encoder
|
||||
|
||||
|
||||
def _create_sam3_model(
|
||||
backbone,
|
||||
transformer,
|
||||
input_geometry_encoder,
|
||||
segmentation_head,
|
||||
dot_prod_scoring,
|
||||
inst_interactive_predictor,
|
||||
eval_mode,
|
||||
):
|
||||
"""Create the SAM3 image model."""
|
||||
common_params = {
|
||||
"backbone": backbone,
|
||||
"transformer": transformer,
|
||||
"input_geometry_encoder": input_geometry_encoder,
|
||||
"segmentation_head": segmentation_head,
|
||||
"num_feature_levels": 1,
|
||||
"o2m_mask_predict": True,
|
||||
"dot_prod_scoring": dot_prod_scoring,
|
||||
"use_instance_query": False,
|
||||
"multimask_output": True,
|
||||
"inst_interactive_predictor": inst_interactive_predictor,
|
||||
}
|
||||
|
||||
matcher = None
|
||||
if not eval_mode:
|
||||
from sam3.train.matcher import BinaryHungarianMatcherV2
|
||||
|
||||
matcher = BinaryHungarianMatcherV2(
|
||||
focal=True,
|
||||
cost_class=2.0,
|
||||
cost_bbox=5.0,
|
||||
cost_giou=2.0,
|
||||
alpha=0.25,
|
||||
gamma=2,
|
||||
stable=False,
|
||||
)
|
||||
common_params["matcher"] = matcher
|
||||
model = Sam3Image(**common_params)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def _create_tracker_maskmem_backbone():
|
||||
"""Create the SAM3 Tracker memory encoder."""
|
||||
# Position encoding for mask memory backbone
|
||||
position_encoding = PositionEmbeddingSine(
|
||||
num_pos_feats=64,
|
||||
normalize=True,
|
||||
scale=None,
|
||||
temperature=10000,
|
||||
precompute_resolution=1008,
|
||||
)
|
||||
|
||||
# Mask processing components
|
||||
mask_downsampler = SimpleMaskDownSampler(
|
||||
kernel_size=3, stride=2, padding=1, interpol_size=[1152, 1152]
|
||||
)
|
||||
|
||||
cx_block_layer = CXBlock(
|
||||
dim=256,
|
||||
kernel_size=7,
|
||||
padding=3,
|
||||
layer_scale_init_value=1.0e-06,
|
||||
use_dwconv=True,
|
||||
)
|
||||
|
||||
fuser = SimpleFuser(layer=cx_block_layer, num_layers=2)
|
||||
|
||||
maskmem_backbone = SimpleMaskEncoder(
|
||||
out_dim=64,
|
||||
position_encoding=position_encoding,
|
||||
mask_downsampler=mask_downsampler,
|
||||
fuser=fuser,
|
||||
)
|
||||
|
||||
return maskmem_backbone
|
||||
|
||||
|
||||
def _create_tracker_transformer():
|
||||
"""Create the SAM3 Tracker transformer components."""
|
||||
# Self attention
|
||||
self_attention = RoPEAttention(
|
||||
embedding_dim=256,
|
||||
num_heads=1,
|
||||
downsample_rate=1,
|
||||
dropout=0.1,
|
||||
rope_theta=10000.0,
|
||||
feat_sizes=[72, 72],
|
||||
use_fa3=False,
|
||||
use_rope_real=False,
|
||||
)
|
||||
|
||||
# Cross attention
|
||||
cross_attention = RoPEAttention(
|
||||
embedding_dim=256,
|
||||
num_heads=1,
|
||||
downsample_rate=1,
|
||||
dropout=0.1,
|
||||
kv_in_dim=64,
|
||||
rope_theta=10000.0,
|
||||
feat_sizes=[72, 72],
|
||||
rope_k_repeat=True,
|
||||
use_fa3=False,
|
||||
use_rope_real=False,
|
||||
)
|
||||
|
||||
# Encoder layer
|
||||
encoder_layer = TransformerDecoderLayerv2(
|
||||
cross_attention_first=False,
|
||||
activation="relu",
|
||||
dim_feedforward=2048,
|
||||
dropout=0.1,
|
||||
pos_enc_at_attn=False,
|
||||
pre_norm=True,
|
||||
self_attention=self_attention,
|
||||
d_model=256,
|
||||
pos_enc_at_cross_attn_keys=True,
|
||||
pos_enc_at_cross_attn_queries=False,
|
||||
cross_attention=cross_attention,
|
||||
)
|
||||
|
||||
# Encoder
|
||||
encoder = TransformerEncoderCrossAttention(
|
||||
remove_cross_attention_layers=[],
|
||||
batch_first=True,
|
||||
d_model=256,
|
||||
frozen=False,
|
||||
pos_enc_at_input=True,
|
||||
layer=encoder_layer,
|
||||
num_layers=4,
|
||||
use_act_checkpoint=False,
|
||||
)
|
||||
|
||||
# Transformer wrapper
|
||||
transformer = TransformerWrapper(
|
||||
encoder=encoder,
|
||||
decoder=None,
|
||||
d_model=256,
|
||||
)
|
||||
|
||||
return transformer
|
||||
|
||||
|
||||
def build_tracker(
|
||||
apply_temporal_disambiguation: bool, with_backbone: bool = False, compile_mode=None
|
||||
) -> Sam3TrackerPredictor:
|
||||
"""
|
||||
Build the SAM3 Tracker module for video tracking.
|
||||
|
||||
Returns:
|
||||
Sam3TrackerPredictor: Wrapped SAM3 Tracker module
|
||||
"""
|
||||
|
||||
# Create model components
|
||||
maskmem_backbone = _create_tracker_maskmem_backbone()
|
||||
transformer = _create_tracker_transformer()
|
||||
backbone = None
|
||||
if with_backbone:
|
||||
vision_backbone = _create_vision_backbone(compile_mode=compile_mode)
|
||||
backbone = SAM3VLBackbone(scalp=1, visual=vision_backbone, text=None)
|
||||
# Create the Tracker module
|
||||
model = Sam3TrackerPredictor(
|
||||
image_size=1008,
|
||||
num_maskmem=7,
|
||||
backbone=backbone,
|
||||
backbone_stride=14,
|
||||
transformer=transformer,
|
||||
maskmem_backbone=maskmem_backbone,
|
||||
# SAM parameters
|
||||
multimask_output_in_sam=True,
|
||||
# Evaluation
|
||||
forward_backbone_per_frame_for_eval=True,
|
||||
trim_past_non_cond_mem_for_eval=False,
|
||||
# Multimask
|
||||
multimask_output_for_tracking=True,
|
||||
multimask_min_pt_num=0,
|
||||
multimask_max_pt_num=1,
|
||||
# Additional settings
|
||||
always_start_from_first_ann_frame=False,
|
||||
# Mask overlap
|
||||
non_overlap_masks_for_mem_enc=False,
|
||||
non_overlap_masks_for_output=False,
|
||||
max_cond_frames_in_attn=4,
|
||||
offload_output_to_cpu_for_eval=False,
|
||||
# SAM decoder settings
|
||||
sam_mask_decoder_extra_args={
|
||||
"dynamic_multimask_via_stability": True,
|
||||
"dynamic_multimask_stability_delta": 0.05,
|
||||
"dynamic_multimask_stability_thresh": 0.98,
|
||||
},
|
||||
clear_non_cond_mem_around_input=True,
|
||||
fill_hole_area=0,
|
||||
use_memory_selection=apply_temporal_disambiguation,
|
||||
)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def _create_text_encoder(bpe_path: str) -> VETextEncoder:
|
||||
"""Create SAM3 text encoder."""
|
||||
tokenizer = SimpleTokenizer(bpe_path=bpe_path)
|
||||
return VETextEncoder(
|
||||
tokenizer=tokenizer,
|
||||
d_model=256,
|
||||
width=1024,
|
||||
heads=16,
|
||||
layers=24,
|
||||
)
|
||||
|
||||
|
||||
def _create_vision_backbone(
|
||||
compile_mode=None, enable_inst_interactivity=True
|
||||
) -> Sam3DualViTDetNeck:
|
||||
"""Create SAM3 visual backbone with ViT and neck."""
|
||||
# Position encoding
|
||||
position_encoding = _create_position_encoding(precompute_resolution=1008)
|
||||
# ViT backbone
|
||||
vit_backbone: ViT = _create_vit_backbone(compile_mode=compile_mode)
|
||||
vit_neck: Sam3DualViTDetNeck = _create_vit_neck(
|
||||
position_encoding,
|
||||
vit_backbone,
|
||||
enable_inst_interactivity=enable_inst_interactivity,
|
||||
)
|
||||
# Visual neck
|
||||
return vit_neck
|
||||
|
||||
|
||||
def _create_sam3_transformer(has_presence_token: bool = True) -> TransformerWrapper:
|
||||
"""Create SAM3 transformer encoder and decoder."""
|
||||
encoder: TransformerEncoderFusion = _create_transformer_encoder()
|
||||
decoder: TransformerDecoder = _create_transformer_decoder()
|
||||
|
||||
return TransformerWrapper(encoder=encoder, decoder=decoder, d_model=256)
|
||||
|
||||
|
||||
def _load_checkpoint(model, checkpoint_path):
|
||||
"""Load model checkpoint from file."""
|
||||
with g_pathmgr.open(checkpoint_path, "rb") as f:
|
||||
ckpt = torch.load(f, map_location="cpu", weights_only=True)
|
||||
if "model" in ckpt and isinstance(ckpt["model"], dict):
|
||||
ckpt = ckpt["model"]
|
||||
sam3_image_ckpt = {
|
||||
k.replace("detector.", ""): v for k, v in ckpt.items() if "detector" in k
|
||||
}
|
||||
if model.inst_interactive_predictor is not None:
|
||||
sam3_image_ckpt.update(
|
||||
{
|
||||
k.replace("tracker.", "inst_interactive_predictor.model."): v
|
||||
for k, v in ckpt.items()
|
||||
if "tracker" in k
|
||||
}
|
||||
)
|
||||
missing_keys, _ = model.load_state_dict(sam3_image_ckpt, strict=False)
|
||||
if len(missing_keys) > 0:
|
||||
print(
|
||||
f"loaded {checkpoint_path} and found "
|
||||
f"missing and/or unexpected keys:\n{missing_keys=}"
|
||||
)
|
||||
|
||||
|
||||
def _setup_device_and_mode(model, device, eval_mode):
|
||||
"""Setup model device and evaluation mode."""
|
||||
if device == "cuda":
|
||||
model = model.cuda()
|
||||
if eval_mode:
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
|
||||
def build_sam3_image_model(
|
||||
bpe_path=None,
|
||||
device="cuda" if torch.cuda.is_available() else "cpu",
|
||||
eval_mode=True,
|
||||
checkpoint_path=None,
|
||||
load_from_HF=True,
|
||||
enable_segmentation=True,
|
||||
enable_inst_interactivity=False,
|
||||
compile=False,
|
||||
):
|
||||
"""
|
||||
Build SAM3 image model
|
||||
|
||||
Args:
|
||||
bpe_path: Path to the BPE tokenizer vocabulary
|
||||
device: Device to load the model on ('cuda' or 'cpu')
|
||||
eval_mode: Whether to set the model to evaluation mode
|
||||
checkpoint_path: Optional path to model checkpoint
|
||||
enable_segmentation: Whether to enable segmentation head
|
||||
enable_inst_interactivity: Whether to enable instance interactivity (SAM 1 task)
|
||||
compile_mode: To enable compilation, set to "default"
|
||||
|
||||
Returns:
|
||||
A SAM3 image model
|
||||
"""
|
||||
if bpe_path is None:
|
||||
bpe_path = os.path.join(
|
||||
os.path.dirname(__file__), "..", "assets", "bpe_simple_vocab_16e6.txt.gz"
|
||||
)
|
||||
# Create visual components
|
||||
compile_mode = "default" if compile else None
|
||||
vision_encoder = _create_vision_backbone(
|
||||
compile_mode=compile_mode, enable_inst_interactivity=enable_inst_interactivity
|
||||
)
|
||||
|
||||
# Create text components
|
||||
text_encoder = _create_text_encoder(bpe_path)
|
||||
|
||||
# Create visual-language backbone
|
||||
backbone = _create_vl_backbone(vision_encoder, text_encoder)
|
||||
|
||||
# Create transformer components
|
||||
transformer = _create_sam3_transformer()
|
||||
|
||||
# Create dot product scoring
|
||||
dot_prod_scoring = _create_dot_product_scoring()
|
||||
|
||||
# Create segmentation head if enabled
|
||||
segmentation_head = (
|
||||
_create_segmentation_head(compile_mode=compile_mode)
|
||||
if enable_segmentation
|
||||
else None
|
||||
)
|
||||
|
||||
# Create geometry encoder
|
||||
input_geometry_encoder = _create_geometry_encoder()
|
||||
if enable_inst_interactivity:
|
||||
sam3_pvs_base = build_tracker(apply_temporal_disambiguation=False)
|
||||
inst_predictor = SAM3InteractiveImagePredictor(sam3_pvs_base)
|
||||
else:
|
||||
inst_predictor = None
|
||||
# Create the SAM3 model
|
||||
model = _create_sam3_model(
|
||||
backbone,
|
||||
transformer,
|
||||
input_geometry_encoder,
|
||||
segmentation_head,
|
||||
dot_prod_scoring,
|
||||
inst_predictor,
|
||||
eval_mode,
|
||||
)
|
||||
if load_from_HF and checkpoint_path is None:
|
||||
checkpoint_path = download_ckpt_from_hf()
|
||||
# Load checkpoint if provided
|
||||
if checkpoint_path is not None:
|
||||
_load_checkpoint(model, checkpoint_path)
|
||||
|
||||
# Setup device and mode
|
||||
model = _setup_device_and_mode(model, device, eval_mode)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def download_ckpt_from_hf():
|
||||
SAM3_MODEL_ID = "facebook/sam3"
|
||||
SAM3_CKPT_NAME = "sam3.pt"
|
||||
SAM3_CFG_NAME = "config.json"
|
||||
_ = hf_hub_download(repo_id=SAM3_MODEL_ID, filename=SAM3_CFG_NAME)
|
||||
checkpoint_path = hf_hub_download(repo_id=SAM3_MODEL_ID, filename=SAM3_CKPT_NAME)
|
||||
return checkpoint_path
|
||||
|
||||
|
||||
def build_sam3_video_model(
|
||||
checkpoint_path: Optional[str] = None,
|
||||
load_from_HF=True,
|
||||
bpe_path: Optional[str] = None,
|
||||
has_presence_token: bool = True,
|
||||
geo_encoder_use_img_cross_attn: bool = True,
|
||||
strict_state_dict_loading: bool = True,
|
||||
apply_temporal_disambiguation: bool = True,
|
||||
device="cuda" if torch.cuda.is_available() else "cpu",
|
||||
compile=False,
|
||||
) -> Sam3VideoInferenceWithInstanceInteractivity:
|
||||
"""
|
||||
Build SAM3 dense tracking model.
|
||||
|
||||
Args:
|
||||
checkpoint_path: Optional path to checkpoint file
|
||||
bpe_path: Path to the BPE tokenizer file
|
||||
|
||||
Returns:
|
||||
Sam3VideoInferenceWithInstanceInteractivity: The instantiated dense tracking model
|
||||
"""
|
||||
if bpe_path is None:
|
||||
bpe_path = os.path.join(
|
||||
os.path.dirname(__file__), "..", "assets", "bpe_simple_vocab_16e6.txt.gz"
|
||||
)
|
||||
|
||||
# Build Tracker module
|
||||
tracker = build_tracker(apply_temporal_disambiguation=apply_temporal_disambiguation)
|
||||
|
||||
# Build Detector components
|
||||
visual_neck = _create_vision_backbone()
|
||||
text_encoder = _create_text_encoder(bpe_path)
|
||||
backbone = SAM3VLBackbone(scalp=1, visual=visual_neck, text=text_encoder)
|
||||
transformer = _create_sam3_transformer(has_presence_token=has_presence_token)
|
||||
segmentation_head: UniversalSegmentationHead = _create_segmentation_head()
|
||||
input_geometry_encoder = _create_geometry_encoder()
|
||||
|
||||
# Create main dot product scoring
|
||||
main_dot_prod_mlp = MLP(
|
||||
input_dim=256,
|
||||
hidden_dim=2048,
|
||||
output_dim=256,
|
||||
num_layers=2,
|
||||
dropout=0.1,
|
||||
residual=True,
|
||||
out_norm=nn.LayerNorm(256),
|
||||
)
|
||||
main_dot_prod_scoring = DotProductScoring(
|
||||
d_model=256, d_proj=256, prompt_mlp=main_dot_prod_mlp
|
||||
)
|
||||
|
||||
# Build Detector module
|
||||
detector = Sam3ImageOnVideoMultiGPU(
|
||||
num_feature_levels=1,
|
||||
backbone=backbone,
|
||||
transformer=transformer,
|
||||
segmentation_head=segmentation_head,
|
||||
semantic_segmentation_head=None,
|
||||
input_geometry_encoder=input_geometry_encoder,
|
||||
use_early_fusion=True,
|
||||
use_dot_prod_scoring=True,
|
||||
dot_prod_scoring=main_dot_prod_scoring,
|
||||
supervise_joint_box_scores=has_presence_token,
|
||||
)
|
||||
|
||||
# Build the main SAM3 video model
|
||||
if apply_temporal_disambiguation:
|
||||
model = Sam3VideoInferenceWithInstanceInteractivity(
|
||||
detector=detector,
|
||||
tracker=tracker,
|
||||
score_threshold_detection=0.5,
|
||||
assoc_iou_thresh=0.1,
|
||||
det_nms_thresh=0.1,
|
||||
new_det_thresh=0.7,
|
||||
hotstart_delay=15,
|
||||
hotstart_unmatch_thresh=8,
|
||||
hotstart_dup_thresh=8,
|
||||
suppress_unmatched_only_within_hotstart=True,
|
||||
min_trk_keep_alive=-1,
|
||||
max_trk_keep_alive=30,
|
||||
init_trk_keep_alive=30,
|
||||
suppress_overlapping_based_on_recent_occlusion_threshold=0.7,
|
||||
suppress_det_close_to_boundary=False,
|
||||
fill_hole_area=16,
|
||||
recondition_every_nth_frame=16,
|
||||
masklet_confirmation_enable=False,
|
||||
decrease_trk_keep_alive_for_empty_masklets=False,
|
||||
image_size=1008,
|
||||
image_mean=(0.5, 0.5, 0.5),
|
||||
image_std=(0.5, 0.5, 0.5),
|
||||
compile_model=compile,
|
||||
)
|
||||
else:
|
||||
# a version without any heuristics for ablation studies
|
||||
model = Sam3VideoInferenceWithInstanceInteractivity(
|
||||
detector=detector,
|
||||
tracker=tracker,
|
||||
score_threshold_detection=0.5,
|
||||
assoc_iou_thresh=0.1,
|
||||
det_nms_thresh=0.1,
|
||||
new_det_thresh=0.7,
|
||||
hotstart_delay=0,
|
||||
hotstart_unmatch_thresh=0,
|
||||
hotstart_dup_thresh=0,
|
||||
suppress_unmatched_only_within_hotstart=True,
|
||||
min_trk_keep_alive=-1,
|
||||
max_trk_keep_alive=30,
|
||||
init_trk_keep_alive=30,
|
||||
suppress_overlapping_based_on_recent_occlusion_threshold=0.7,
|
||||
suppress_det_close_to_boundary=False,
|
||||
fill_hole_area=16,
|
||||
recondition_every_nth_frame=0,
|
||||
masklet_confirmation_enable=False,
|
||||
decrease_trk_keep_alive_for_empty_masklets=False,
|
||||
image_size=1008,
|
||||
image_mean=(0.5, 0.5, 0.5),
|
||||
image_std=(0.5, 0.5, 0.5),
|
||||
compile_model=compile,
|
||||
)
|
||||
|
||||
# Load checkpoint if provided
|
||||
if load_from_HF and checkpoint_path is None:
|
||||
checkpoint_path = download_ckpt_from_hf()
|
||||
if checkpoint_path is not None:
|
||||
with g_pathmgr.open(checkpoint_path, "rb") as f:
|
||||
ckpt = torch.load(f, map_location="cpu", weights_only=True)
|
||||
if "model" in ckpt and isinstance(ckpt["model"], dict):
|
||||
ckpt = ckpt["model"]
|
||||
|
||||
missing_keys, unexpected_keys = model.load_state_dict(
|
||||
ckpt, strict=strict_state_dict_loading
|
||||
)
|
||||
if missing_keys:
|
||||
print(f"Missing keys: {missing_keys}")
|
||||
if unexpected_keys:
|
||||
print(f"Unexpected keys: {unexpected_keys}")
|
||||
|
||||
model.to(device=device)
|
||||
return model
|
||||
|
||||
|
||||
def build_sam3_video_predictor(*model_args, gpus_to_use=None, **model_kwargs):
|
||||
return Sam3VideoPredictorMultiGPU(
|
||||
*model_args, gpus_to_use=gpus_to_use, **model_kwargs
|
||||
)
|
||||
@@ -0,0 +1,4 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
from .mask_decoder import MaskDecoder
|
||||
from .prompt_encoder import PromptEncoder
|
||||
from .transformer import TwoWayTransformer
|
||||
@@ -0,0 +1,39 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
|
||||
from typing import Type
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class MLPBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
mlp_dim: int,
|
||||
act: Type[nn.Module] = nn.GELU,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.lin1 = nn.Linear(embedding_dim, mlp_dim)
|
||||
self.lin2 = nn.Linear(mlp_dim, embedding_dim)
|
||||
self.act = act()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.lin2(self.act(self.lin1(x)))
|
||||
|
||||
|
||||
# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
|
||||
# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa
|
||||
class LayerNorm2d(nn.Module):
|
||||
def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(num_channels))
|
||||
self.bias = nn.Parameter(torch.zeros(num_channels))
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
u = x.mean(1, keepdim=True)
|
||||
s = (x - u).pow(2).mean(1, keepdim=True)
|
||||
x = (x - u) / torch.sqrt(s + self.eps)
|
||||
x = self.weight[:, None, None] * x + self.bias[:, None, None]
|
||||
return x
|
||||
@@ -0,0 +1,319 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
|
||||
from typing import List, Optional, Tuple, Type
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from .common import LayerNorm2d
|
||||
|
||||
|
||||
class MaskDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
transformer_dim: int,
|
||||
transformer: nn.Module,
|
||||
num_multimask_outputs: int = 3,
|
||||
activation: Type[nn.Module] = nn.GELU,
|
||||
iou_head_depth: int = 3,
|
||||
iou_head_hidden_dim: int = 256,
|
||||
use_high_res_features: bool = False,
|
||||
iou_prediction_use_sigmoid=False,
|
||||
dynamic_multimask_via_stability=False,
|
||||
dynamic_multimask_stability_delta=0.05,
|
||||
dynamic_multimask_stability_thresh=0.98,
|
||||
pred_obj_scores: bool = False,
|
||||
pred_obj_scores_mlp: bool = False,
|
||||
use_multimask_token_for_obj_ptr: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Predicts masks given an image and prompt embeddings, using a
|
||||
transformer architecture.
|
||||
|
||||
Arguments:
|
||||
transformer_dim (int): the channel dimension of the transformer
|
||||
transformer (nn.Module): the transformer used to predict masks
|
||||
num_multimask_outputs (int): the number of masks to predict
|
||||
when disambiguating masks
|
||||
activation (nn.Module): the type of activation to use when
|
||||
upscaling masks
|
||||
iou_head_depth (int): the depth of the MLP used to predict
|
||||
mask quality
|
||||
iou_head_hidden_dim (int): the hidden dimension of the MLP
|
||||
used to predict mask quality
|
||||
"""
|
||||
super().__init__()
|
||||
self.transformer_dim = transformer_dim
|
||||
self.transformer = transformer
|
||||
|
||||
self.num_multimask_outputs = num_multimask_outputs
|
||||
|
||||
self.iou_token = nn.Embedding(1, transformer_dim)
|
||||
self.num_mask_tokens = num_multimask_outputs + 1
|
||||
self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)
|
||||
|
||||
self.pred_obj_scores = pred_obj_scores
|
||||
if self.pred_obj_scores:
|
||||
self.obj_score_token = nn.Embedding(1, transformer_dim)
|
||||
self.use_multimask_token_for_obj_ptr = use_multimask_token_for_obj_ptr
|
||||
|
||||
self.output_upscaling = nn.Sequential(
|
||||
nn.ConvTranspose2d(
|
||||
transformer_dim, transformer_dim // 4, kernel_size=2, stride=2
|
||||
),
|
||||
LayerNorm2d(transformer_dim // 4),
|
||||
activation(),
|
||||
nn.ConvTranspose2d(
|
||||
transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2
|
||||
),
|
||||
activation(),
|
||||
)
|
||||
self.use_high_res_features = use_high_res_features
|
||||
if use_high_res_features:
|
||||
self.conv_s0 = nn.Conv2d(
|
||||
transformer_dim, transformer_dim // 8, kernel_size=1, stride=1
|
||||
)
|
||||
self.conv_s1 = nn.Conv2d(
|
||||
transformer_dim, transformer_dim // 4, kernel_size=1, stride=1
|
||||
)
|
||||
|
||||
self.output_hypernetworks_mlps = nn.ModuleList(
|
||||
[
|
||||
MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3)
|
||||
for i in range(self.num_mask_tokens)
|
||||
]
|
||||
)
|
||||
|
||||
self.iou_prediction_head = MLP(
|
||||
transformer_dim,
|
||||
iou_head_hidden_dim,
|
||||
self.num_mask_tokens,
|
||||
iou_head_depth,
|
||||
sigmoid_output=iou_prediction_use_sigmoid,
|
||||
)
|
||||
if self.pred_obj_scores:
|
||||
self.pred_obj_score_head = nn.Linear(transformer_dim, 1)
|
||||
if pred_obj_scores_mlp:
|
||||
self.pred_obj_score_head = MLP(transformer_dim, transformer_dim, 1, 3)
|
||||
|
||||
# When outputting a single mask, optionally we can dynamically fall back to the best
|
||||
# multimask output token if the single mask output token gives low stability scores.
|
||||
self.dynamic_multimask_via_stability = dynamic_multimask_via_stability
|
||||
self.dynamic_multimask_stability_delta = dynamic_multimask_stability_delta
|
||||
self.dynamic_multimask_stability_thresh = dynamic_multimask_stability_thresh
|
||||
|
||||
def forward(
|
||||
self,
|
||||
image_embeddings: torch.Tensor,
|
||||
image_pe: torch.Tensor,
|
||||
sparse_prompt_embeddings: torch.Tensor,
|
||||
dense_prompt_embeddings: torch.Tensor,
|
||||
multimask_output: bool,
|
||||
repeat_image: bool,
|
||||
high_res_features: Optional[List[torch.Tensor]] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Predict masks given image and prompt embeddings.
|
||||
|
||||
Arguments:
|
||||
image_embeddings (torch.Tensor): the embeddings from the image encoder
|
||||
image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
|
||||
sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
|
||||
dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
|
||||
multimask_output (bool): Whether to return multiple masks or a single
|
||||
mask.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: batched predicted masks
|
||||
torch.Tensor: batched predictions of mask quality
|
||||
torch.Tensor: batched SAM token for mask output
|
||||
"""
|
||||
masks, iou_pred, mask_tokens_out, object_score_logits = self.predict_masks(
|
||||
image_embeddings=image_embeddings,
|
||||
image_pe=image_pe,
|
||||
sparse_prompt_embeddings=sparse_prompt_embeddings,
|
||||
dense_prompt_embeddings=dense_prompt_embeddings,
|
||||
repeat_image=repeat_image,
|
||||
high_res_features=high_res_features,
|
||||
)
|
||||
|
||||
# Select the correct mask or masks for output
|
||||
if multimask_output:
|
||||
masks = masks[:, 1:, :, :]
|
||||
iou_pred = iou_pred[:, 1:]
|
||||
elif self.dynamic_multimask_via_stability and not self.training:
|
||||
masks, iou_pred = self._dynamic_multimask_via_stability(masks, iou_pred)
|
||||
else:
|
||||
masks = masks[:, 0:1, :, :]
|
||||
iou_pred = iou_pred[:, 0:1]
|
||||
|
||||
if multimask_output and self.use_multimask_token_for_obj_ptr:
|
||||
sam_tokens_out = mask_tokens_out[:, 1:] # [b, 3, c] shape
|
||||
else:
|
||||
# Take the mask output token. Here we *always* use the token for single mask output.
|
||||
# At test time, even if we track after 1-click (and using multimask_output=True),
|
||||
# we still take the single mask token here. The rationale is that we always track
|
||||
# after multiple clicks during training, so the past tokens seen during training
|
||||
# are always the single mask token (and we'll let it be the object-memory token).
|
||||
sam_tokens_out = mask_tokens_out[:, 0:1] # [b, 1, c] shape
|
||||
|
||||
# Prepare output
|
||||
return masks, iou_pred, sam_tokens_out, object_score_logits
|
||||
|
||||
def predict_masks(
|
||||
self,
|
||||
image_embeddings: torch.Tensor,
|
||||
image_pe: torch.Tensor,
|
||||
sparse_prompt_embeddings: torch.Tensor,
|
||||
dense_prompt_embeddings: torch.Tensor,
|
||||
repeat_image: bool,
|
||||
high_res_features: Optional[List[torch.Tensor]] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Predicts masks. See 'forward' for more details."""
|
||||
# Concatenate output tokens
|
||||
s = 0
|
||||
if self.pred_obj_scores:
|
||||
output_tokens = torch.cat(
|
||||
[
|
||||
self.obj_score_token.weight,
|
||||
self.iou_token.weight,
|
||||
self.mask_tokens.weight,
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
s = 1
|
||||
else:
|
||||
output_tokens = torch.cat(
|
||||
[self.iou_token.weight, self.mask_tokens.weight], dim=0
|
||||
)
|
||||
output_tokens = output_tokens.unsqueeze(0).expand(
|
||||
sparse_prompt_embeddings.size(0), -1, -1
|
||||
)
|
||||
tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)
|
||||
|
||||
# Expand per-image data in batch direction to be per-mask
|
||||
if repeat_image:
|
||||
src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)
|
||||
else:
|
||||
assert image_embeddings.shape[0] == tokens.shape[0]
|
||||
src = image_embeddings
|
||||
src = src + dense_prompt_embeddings
|
||||
assert (
|
||||
image_pe.size(0) == 1
|
||||
), "image_pe should have size 1 in batch dim (from `get_dense_pe()`)"
|
||||
pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
|
||||
b, c, h, w = src.shape
|
||||
|
||||
# Run the transformer
|
||||
hs, src = self.transformer(src, pos_src, tokens)
|
||||
iou_token_out = hs[:, s, :]
|
||||
mask_tokens_out = hs[:, s + 1 : (s + 1 + self.num_mask_tokens), :]
|
||||
|
||||
# Upscale mask embeddings and predict masks using the mask tokens
|
||||
src = src.transpose(1, 2).view(b, c, h, w)
|
||||
if not self.use_high_res_features:
|
||||
upscaled_embedding = self.output_upscaling(src)
|
||||
else:
|
||||
dc1, ln1, act1, dc2, act2 = self.output_upscaling
|
||||
feat_s0, feat_s1 = high_res_features
|
||||
upscaled_embedding = act1(ln1(dc1(src) + feat_s1))
|
||||
upscaled_embedding = act2(dc2(upscaled_embedding) + feat_s0)
|
||||
|
||||
hyper_in_list: List[torch.Tensor] = []
|
||||
for i in range(self.num_mask_tokens):
|
||||
hyper_in_list.append(
|
||||
self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])
|
||||
)
|
||||
hyper_in = torch.stack(hyper_in_list, dim=1)
|
||||
b, c, h, w = upscaled_embedding.shape
|
||||
masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w)
|
||||
|
||||
# Generate mask quality predictions
|
||||
iou_pred = self.iou_prediction_head(iou_token_out)
|
||||
if self.pred_obj_scores:
|
||||
assert s == 1
|
||||
object_score_logits = self.pred_obj_score_head(hs[:, 0, :])
|
||||
else:
|
||||
# Obj scores logits - default to 10.0, i.e. assuming the object is present, sigmoid(10)=1
|
||||
object_score_logits = 10.0 * iou_pred.new_ones(iou_pred.shape[0], 1)
|
||||
|
||||
return masks, iou_pred, mask_tokens_out, object_score_logits
|
||||
|
||||
def _get_stability_scores(self, mask_logits):
|
||||
"""
|
||||
Compute stability scores of the mask logits based on the IoU between upper and
|
||||
lower thresholds.
|
||||
"""
|
||||
mask_logits = mask_logits.flatten(-2)
|
||||
stability_delta = self.dynamic_multimask_stability_delta
|
||||
area_i = torch.sum(mask_logits > stability_delta, dim=-1).float()
|
||||
area_u = torch.sum(mask_logits > -stability_delta, dim=-1).float()
|
||||
stability_scores = torch.where(area_u > 0, area_i / area_u, 1.0)
|
||||
return stability_scores
|
||||
|
||||
def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores):
|
||||
"""
|
||||
When outputting a single mask, if the stability score from the current single-mask
|
||||
output (based on output token 0) falls below a threshold, we instead select from
|
||||
multi-mask outputs (based on output token 1~3) the mask with the highest predicted
|
||||
IoU score. This is intended to ensure a valid mask for both clicking and tracking.
|
||||
"""
|
||||
# The best mask from multimask output tokens (1~3)
|
||||
multimask_logits = all_mask_logits[:, 1:, :, :]
|
||||
multimask_iou_scores = all_iou_scores[:, 1:]
|
||||
best_scores_inds = torch.argmax(multimask_iou_scores, dim=-1)
|
||||
batch_inds = torch.arange(
|
||||
multimask_iou_scores.size(0), device=all_iou_scores.device
|
||||
)
|
||||
best_multimask_logits = multimask_logits[batch_inds, best_scores_inds]
|
||||
best_multimask_logits = best_multimask_logits.unsqueeze(1)
|
||||
best_multimask_iou_scores = multimask_iou_scores[batch_inds, best_scores_inds]
|
||||
best_multimask_iou_scores = best_multimask_iou_scores.unsqueeze(1)
|
||||
|
||||
# The mask from singlemask output token 0 and its stability score
|
||||
singlemask_logits = all_mask_logits[:, 0:1, :, :]
|
||||
singlemask_iou_scores = all_iou_scores[:, 0:1]
|
||||
stability_scores = self._get_stability_scores(singlemask_logits)
|
||||
is_stable = stability_scores >= self.dynamic_multimask_stability_thresh
|
||||
|
||||
# Dynamically fall back to best multimask output upon low stability scores.
|
||||
mask_logits_out = torch.where(
|
||||
is_stable[..., None, None].expand_as(singlemask_logits),
|
||||
singlemask_logits,
|
||||
best_multimask_logits,
|
||||
)
|
||||
iou_scores_out = torch.where(
|
||||
is_stable.expand_as(singlemask_iou_scores),
|
||||
singlemask_iou_scores,
|
||||
best_multimask_iou_scores,
|
||||
)
|
||||
return mask_logits_out, iou_scores_out
|
||||
|
||||
|
||||
# Lightly adapted from
|
||||
# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa
|
||||
class MLP(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: int,
|
||||
hidden_dim: int,
|
||||
output_dim: int,
|
||||
num_layers: int,
|
||||
sigmoid_output: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_layers = num_layers
|
||||
h = [hidden_dim] * (num_layers - 1)
|
||||
self.layers = nn.ModuleList(
|
||||
nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
|
||||
)
|
||||
self.sigmoid_output = sigmoid_output
|
||||
|
||||
def forward(self, x):
|
||||
for i, layer in enumerate(self.layers):
|
||||
x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
|
||||
if self.sigmoid_output:
|
||||
x = F.sigmoid(x)
|
||||
return x
|
||||
@@ -0,0 +1,243 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
|
||||
from typing import Any, Optional, Tuple, Type
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from .common import LayerNorm2d
|
||||
|
||||
|
||||
class PromptEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embed_dim: int,
|
||||
image_embedding_size: Tuple[int, int],
|
||||
input_image_size: Tuple[int, int],
|
||||
mask_in_chans: int,
|
||||
activation: Type[nn.Module] = nn.GELU,
|
||||
) -> None:
|
||||
"""
|
||||
Encodes prompts for input to SAM's mask decoder.
|
||||
|
||||
Arguments:
|
||||
embed_dim (int): The prompts' embedding dimension
|
||||
image_embedding_size (tuple(int, int)): The spatial size of the
|
||||
image embedding, as (H, W).
|
||||
input_image_size (int): The padded size of the image as input
|
||||
to the image encoder, as (H, W).
|
||||
mask_in_chans (int): The number of hidden channels used for
|
||||
encoding input masks.
|
||||
activation (nn.Module): The activation to use when encoding
|
||||
input masks.
|
||||
"""
|
||||
super().__init__()
|
||||
self.embed_dim = embed_dim
|
||||
self.input_image_size = input_image_size
|
||||
self.image_embedding_size = image_embedding_size
|
||||
self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)
|
||||
|
||||
self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners
|
||||
point_embeddings = [
|
||||
nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)
|
||||
]
|
||||
self.point_embeddings = nn.ModuleList(point_embeddings)
|
||||
self.not_a_point_embed = nn.Embedding(1, embed_dim)
|
||||
|
||||
self.mask_input_size = (
|
||||
4 * image_embedding_size[0],
|
||||
4 * image_embedding_size[1],
|
||||
)
|
||||
self.mask_downscaling = nn.Sequential(
|
||||
nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),
|
||||
LayerNorm2d(mask_in_chans // 4),
|
||||
activation(),
|
||||
nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),
|
||||
LayerNorm2d(mask_in_chans),
|
||||
activation(),
|
||||
nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1),
|
||||
)
|
||||
self.no_mask_embed = nn.Embedding(1, embed_dim)
|
||||
|
||||
def get_dense_pe(self) -> torch.Tensor:
|
||||
"""
|
||||
Returns the positional encoding used to encode point prompts,
|
||||
applied to a dense set of points the shape of the image encoding.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Positional encoding with shape
|
||||
1x(embed_dim)x(embedding_h)x(embedding_w)
|
||||
"""
|
||||
return self.pe_layer(self.image_embedding_size).unsqueeze(0)
|
||||
|
||||
def _embed_points(
|
||||
self,
|
||||
points: torch.Tensor,
|
||||
labels: torch.Tensor,
|
||||
pad: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Embeds point prompts."""
|
||||
points = points + 0.5 # Shift to center of pixel
|
||||
if pad:
|
||||
padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)
|
||||
padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
|
||||
points = torch.cat([points, padding_point], dim=1)
|
||||
labels = torch.cat([labels, padding_label], dim=1)
|
||||
point_embedding = self.pe_layer.forward_with_coords(
|
||||
points, self.input_image_size
|
||||
)
|
||||
|
||||
point_embedding = torch.where(
|
||||
(labels == -1).unsqueeze(-1),
|
||||
torch.zeros_like(point_embedding) + self.not_a_point_embed.weight,
|
||||
point_embedding,
|
||||
)
|
||||
point_embedding = torch.where(
|
||||
(labels == 0).unsqueeze(-1),
|
||||
point_embedding + self.point_embeddings[0].weight,
|
||||
point_embedding,
|
||||
)
|
||||
point_embedding = torch.where(
|
||||
(labels == 1).unsqueeze(-1),
|
||||
point_embedding + self.point_embeddings[1].weight,
|
||||
point_embedding,
|
||||
)
|
||||
point_embedding = torch.where(
|
||||
(labels == 2).unsqueeze(-1),
|
||||
point_embedding + self.point_embeddings[2].weight,
|
||||
point_embedding,
|
||||
)
|
||||
point_embedding = torch.where(
|
||||
(labels == 3).unsqueeze(-1),
|
||||
point_embedding + self.point_embeddings[3].weight,
|
||||
point_embedding,
|
||||
)
|
||||
return point_embedding
|
||||
|
||||
def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
|
||||
"""Embeds box prompts."""
|
||||
boxes = boxes + 0.5 # Shift to center of pixel
|
||||
coords = boxes.reshape(-1, 2, 2)
|
||||
corner_embedding = self.pe_layer.forward_with_coords(
|
||||
coords, self.input_image_size
|
||||
)
|
||||
corner_embedding[:, 0, :] += self.point_embeddings[2].weight
|
||||
corner_embedding[:, 1, :] += self.point_embeddings[3].weight
|
||||
return corner_embedding
|
||||
|
||||
def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:
|
||||
"""Embeds mask inputs."""
|
||||
mask_embedding = self.mask_downscaling(masks)
|
||||
return mask_embedding
|
||||
|
||||
def _get_batch_size(
|
||||
self,
|
||||
points: Optional[Tuple[torch.Tensor, torch.Tensor]],
|
||||
boxes: Optional[torch.Tensor],
|
||||
masks: Optional[torch.Tensor],
|
||||
) -> int:
|
||||
"""
|
||||
Gets the batch size of the output given the batch size of the input prompts.
|
||||
"""
|
||||
if points is not None:
|
||||
return points[0].shape[0]
|
||||
elif boxes is not None:
|
||||
return boxes.shape[0]
|
||||
elif masks is not None:
|
||||
return masks.shape[0]
|
||||
else:
|
||||
return 1
|
||||
|
||||
def _get_device(self) -> torch.device:
|
||||
return self.point_embeddings[0].weight.device
|
||||
|
||||
def forward(
|
||||
self,
|
||||
points: Optional[Tuple[torch.Tensor, torch.Tensor]],
|
||||
boxes: Optional[torch.Tensor],
|
||||
masks: Optional[torch.Tensor],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Embeds different types of prompts, returning both sparse and dense
|
||||
embeddings.
|
||||
|
||||
Arguments:
|
||||
points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates
|
||||
and labels to embed.
|
||||
boxes (torch.Tensor or none): boxes to embed
|
||||
masks (torch.Tensor or none): masks to embed
|
||||
|
||||
Returns:
|
||||
torch.Tensor: sparse embeddings for the points and boxes, with shape
|
||||
BxNx(embed_dim), where N is determined by the number of input points
|
||||
and boxes.
|
||||
torch.Tensor: dense embeddings for the masks, in the shape
|
||||
Bx(embed_dim)x(embed_H)x(embed_W)
|
||||
"""
|
||||
bs = self._get_batch_size(points, boxes, masks)
|
||||
sparse_embeddings = torch.empty(
|
||||
(bs, 0, self.embed_dim), device=self._get_device()
|
||||
)
|
||||
if points is not None:
|
||||
coords, labels = points
|
||||
point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))
|
||||
sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
|
||||
if boxes is not None:
|
||||
box_embeddings = self._embed_boxes(boxes)
|
||||
sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
|
||||
|
||||
if masks is not None:
|
||||
dense_embeddings = self._embed_masks(masks)
|
||||
else:
|
||||
dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
|
||||
bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]
|
||||
)
|
||||
|
||||
return sparse_embeddings, dense_embeddings
|
||||
|
||||
|
||||
class PositionEmbeddingRandom(nn.Module):
|
||||
"""
|
||||
Positional encoding using random spatial frequencies.
|
||||
"""
|
||||
|
||||
def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:
|
||||
super().__init__()
|
||||
if scale is None or scale <= 0.0:
|
||||
scale = 1.0
|
||||
self.register_buffer(
|
||||
"positional_encoding_gaussian_matrix",
|
||||
scale * torch.randn((2, num_pos_feats)),
|
||||
)
|
||||
|
||||
def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
|
||||
"""Positionally encode points that are normalized to [0,1]."""
|
||||
# assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
|
||||
coords = 2 * coords - 1
|
||||
coords = coords @ self.positional_encoding_gaussian_matrix
|
||||
coords = 2 * np.pi * coords
|
||||
# outputs d_1 x ... x d_n x C shape
|
||||
return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
|
||||
|
||||
def forward(self, size: Tuple[int, int]) -> torch.Tensor:
|
||||
"""Generate positional encoding for a grid of the specified size."""
|
||||
h, w = size
|
||||
device: Any = self.positional_encoding_gaussian_matrix.device
|
||||
grid = torch.ones((h, w), device=device, dtype=torch.float32)
|
||||
y_embed = grid.cumsum(dim=0) - 0.5
|
||||
x_embed = grid.cumsum(dim=1) - 0.5
|
||||
y_embed = y_embed / h
|
||||
x_embed = x_embed / w
|
||||
|
||||
pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
|
||||
return pe.permute(2, 0, 1) # C x H x W
|
||||
|
||||
def forward_with_coords(
|
||||
self, coords_input: torch.Tensor, image_size: Tuple[int, int]
|
||||
) -> torch.Tensor:
|
||||
"""Positionally encode points that are not normalized to [0,1]."""
|
||||
coords = coords_input.clone()
|
||||
coords[:, :, 0] = coords[:, :, 0] / image_size[1]
|
||||
coords[:, :, 1] = coords[:, :, 1] / image_size[0]
|
||||
return self._pe_encoding(coords.to(torch.float)) # B x N x C
|
||||
@@ -0,0 +1,161 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
|
||||
"""
|
||||
Adapted from:
|
||||
1. https://github.com/meta-llama/codellama/blob/main/llama/model.py
|
||||
2. https://github.com/naver-ai/rope-vit
|
||||
3. https://github.com/lucidrains/rotary-embedding-torch
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from einops import rearrange, repeat
|
||||
from torch import broadcast_tensors, nn
|
||||
|
||||
|
||||
def init_t_xy(end_x: int, end_y: int, scale: float = 1.0, offset: int = 0, device=None):
|
||||
t = torch.arange(end_x * end_y, dtype=torch.float32, device=device)
|
||||
t_x = (t % end_x).float()
|
||||
t_y = torch.div(t, end_x, rounding_mode="floor").float()
|
||||
return t_x * scale + offset, t_y * scale + offset
|
||||
|
||||
|
||||
def compute_axial_cis(
|
||||
dim: int,
|
||||
end_x: int,
|
||||
end_y: int,
|
||||
theta: float = 10000.0,
|
||||
scale_pos: float = 1.0,
|
||||
offset: int = 0,
|
||||
device=None,
|
||||
):
|
||||
freqs_x = 1.0 / (
|
||||
theta ** (torch.arange(0, dim, 4, device=device)[: (dim // 4)].float() / dim)
|
||||
)
|
||||
freqs_y = 1.0 / (
|
||||
theta ** (torch.arange(0, dim, 4, device=device)[: (dim // 4)].float() / dim)
|
||||
)
|
||||
|
||||
t_x, t_y = init_t_xy(end_x, end_y, scale_pos, offset, device=device)
|
||||
freqs_x = torch.outer(t_x, freqs_x)
|
||||
freqs_y = torch.outer(t_y, freqs_y)
|
||||
freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x)
|
||||
freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y)
|
||||
return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1)
|
||||
|
||||
|
||||
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
|
||||
ndim = x.ndim
|
||||
assert 0 <= 1 < ndim
|
||||
assert freqs_cis.shape == (x.shape[-2], x.shape[-1])
|
||||
shape = [d if i >= ndim - 2 else 1 for i, d in enumerate(x.shape)]
|
||||
return freqs_cis.view(*shape)
|
||||
|
||||
|
||||
def apply_rotary_enc(
|
||||
xq: torch.Tensor,
|
||||
xk: torch.Tensor,
|
||||
freqs_cis: torch.Tensor,
|
||||
repeat_freqs_k: bool = False,
|
||||
):
|
||||
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
|
||||
xk_ = (
|
||||
torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
|
||||
if xk.shape[-2] != 0
|
||||
else None
|
||||
)
|
||||
freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
|
||||
xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
|
||||
if xk_ is None:
|
||||
# no keys to rotate, due to dropout
|
||||
return xq_out.type_as(xq).to(xq.device), xk
|
||||
# repeat freqs along seq_len dim to match k seq_len
|
||||
if repeat_freqs_k:
|
||||
r = xk_.shape[-2] // xq_.shape[-2]
|
||||
freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1)
|
||||
xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
|
||||
return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)
|
||||
|
||||
|
||||
def complex_mult(xq_real, xq_imag, freqs_cis_real, freqs_cis_imag):
|
||||
# Compute the real part of the product
|
||||
real_part = xq_real * freqs_cis_real - xq_imag * freqs_cis_imag
|
||||
# Compute the imaginary part of the product
|
||||
imag_part = xq_real * freqs_cis_imag + xq_imag * freqs_cis_real
|
||||
# Stack the real and imaginary parts along the last dimension
|
||||
return torch.stack([real_part, imag_part], dim=-1)
|
||||
|
||||
|
||||
def apply_rotary_enc_real(
|
||||
xq: torch.Tensor,
|
||||
xk: torch.Tensor,
|
||||
freqs_cis_real: torch.Tensor,
|
||||
freqs_cis_imag: torch.Tensor,
|
||||
repeat_freqs_k: bool = False,
|
||||
):
|
||||
assert xk is not None
|
||||
assert xk.shape[-2] != 0
|
||||
|
||||
xq_real = xq.float().reshape(*xq.shape[:-1], -1, 2)[..., 0]
|
||||
xq_imag = xq.float().reshape(*xq.shape[:-1], -1, 2)[..., 1]
|
||||
xk_real = xk.float().reshape(*xk.shape[:-1], -1, 2)[..., 0]
|
||||
xk_imag = xk.float().reshape(*xk.shape[:-1], -1, 2)[..., 1]
|
||||
freqs_cis_real = reshape_for_broadcast(freqs_cis_real, xq_real)
|
||||
freqs_cis_imag = reshape_for_broadcast(freqs_cis_imag, xq_imag)
|
||||
xq_out = complex_mult(xq_real, xq_imag, freqs_cis_real, freqs_cis_imag).flatten(3)
|
||||
if repeat_freqs_k:
|
||||
r = xk_real.shape[-2] // xq_real.shape[-2]
|
||||
freqs_cis_real = freqs_cis_real.repeat(*([1] * (freqs_cis_real.ndim - 2)), r, 1)
|
||||
freqs_cis_imag = freqs_cis_imag.repeat(*([1] * (freqs_cis_imag.ndim - 2)), r, 1)
|
||||
xk_out = complex_mult(xk_real, xk_imag, freqs_cis_real, freqs_cis_imag).flatten(3)
|
||||
# xq_out = torch.view_as_real(torch.complex(xq_real, xq_imag) * torch.complex(freqs_cis_real, freqs_cis_imag)).flatten(3)
|
||||
# xk_out = torch.view_as_real(torch.compelx(xk_real, xk_imag) * torch.complex(freqs_cis_real, freqs_cis_imag)).flatten(3)
|
||||
return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)
|
||||
|
||||
|
||||
# rotary embedding helper functions
|
||||
def broadcat(tensors, dim=-1):
|
||||
broadcasted_tensors = broadcast_tensors(*tensors)
|
||||
return torch.cat(broadcasted_tensors, dim=dim)
|
||||
|
||||
|
||||
def rotate_half(x: torch.Tensor):
|
||||
x = rearrange(x, "... (d r) -> ... d r", r=2)
|
||||
x1, x2 = x.unbind(dim=-1)
|
||||
x = torch.stack((-x2, x1), dim=-1)
|
||||
return rearrange(x, "... d r -> ... (d r)")
|
||||
|
||||
|
||||
class VisionRotaryEmbeddingVE(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
seq_len: int,
|
||||
pt_seq_len: Optional[int] = None,
|
||||
theta: float = 10000.0,
|
||||
offset: int = 1, # specific to VE
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
|
||||
scale = 1.0
|
||||
if pt_seq_len is not None:
|
||||
scale = pt_seq_len / seq_len
|
||||
|
||||
# offset of +1 following VE - even though for the
|
||||
# attention op only differences matter
|
||||
t = torch.arange(seq_len) * scale + offset
|
||||
|
||||
freqs = torch.einsum("..., f -> ... f", t, freqs)
|
||||
freqs = repeat(freqs, "... n -> ... (n r)", r=2)
|
||||
|
||||
freqs = broadcat((freqs[None, :, :], freqs[:, None, :]), dim=-1)
|
||||
freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
|
||||
freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
|
||||
|
||||
self.register_buffer("freqs_cos", freqs_cos)
|
||||
self.register_buffer("freqs_sin", freqs_sin)
|
||||
|
||||
def forward(self, t: torch.Tensor):
|
||||
return t * self.freqs_cos + rotate_half(t) * self.freqs_sin
|
||||
@@ -0,0 +1,361 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
|
||||
import math
|
||||
from functools import partial
|
||||
from typing import Tuple, Type
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sam3.sam.rope import apply_rotary_enc, apply_rotary_enc_real, compute_axial_cis
|
||||
from torch import nn, Tensor
|
||||
|
||||
from .common import MLPBlock
|
||||
|
||||
|
||||
class TwoWayTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
depth: int,
|
||||
embedding_dim: int,
|
||||
num_heads: int,
|
||||
mlp_dim: int,
|
||||
activation: Type[nn.Module] = nn.ReLU,
|
||||
attention_downsample_rate: int = 2,
|
||||
) -> None:
|
||||
"""
|
||||
A transformer decoder that attends to an input image using
|
||||
queries whose positional embedding is supplied.
|
||||
|
||||
Args:
|
||||
depth (int): number of layers in the transformer
|
||||
embedding_dim (int): the channel dimension for the input embeddings
|
||||
num_heads (int): the number of heads for multihead attention. Must
|
||||
divide embedding_dim
|
||||
mlp_dim (int): the channel dimension internal to the MLP block
|
||||
activation (nn.Module): the activation to use in the MLP block
|
||||
"""
|
||||
super().__init__()
|
||||
self.depth = depth
|
||||
self.embedding_dim = embedding_dim
|
||||
self.num_heads = num_heads
|
||||
self.mlp_dim = mlp_dim
|
||||
self.layers = nn.ModuleList()
|
||||
|
||||
for i in range(depth):
|
||||
self.layers.append(
|
||||
TwoWayAttentionBlock(
|
||||
embedding_dim=embedding_dim,
|
||||
num_heads=num_heads,
|
||||
mlp_dim=mlp_dim,
|
||||
activation=activation,
|
||||
attention_downsample_rate=attention_downsample_rate,
|
||||
skip_first_layer_pe=(i == 0),
|
||||
)
|
||||
)
|
||||
|
||||
self.final_attn_token_to_image = Attention(
|
||||
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
|
||||
)
|
||||
self.norm_final_attn = nn.LayerNorm(embedding_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
image_embedding: Tensor,
|
||||
image_pe: Tensor,
|
||||
point_embedding: Tensor,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
"""
|
||||
Args:
|
||||
image_embedding (torch.Tensor): image to attend to. Should be shape
|
||||
B x embedding_dim x h x w for any h and w.
|
||||
image_pe (torch.Tensor): the positional encoding to add to the image. Must
|
||||
have the same shape as image_embedding.
|
||||
point_embedding (torch.Tensor): the embedding to add to the query points.
|
||||
Must have shape B x N_points x embedding_dim for any N_points.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: the processed point_embedding
|
||||
torch.Tensor: the processed image_embedding
|
||||
"""
|
||||
# BxCxHxW -> BxHWxC == B x N_image_tokens x C
|
||||
bs, c, h, w = image_embedding.shape
|
||||
image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
|
||||
image_pe = image_pe.flatten(2).permute(0, 2, 1)
|
||||
|
||||
# Prepare queries
|
||||
queries = point_embedding
|
||||
keys = image_embedding
|
||||
|
||||
# Apply transformer blocks and final layernorm
|
||||
for layer in self.layers:
|
||||
queries, keys = layer(
|
||||
queries=queries,
|
||||
keys=keys,
|
||||
query_pe=point_embedding,
|
||||
key_pe=image_pe,
|
||||
)
|
||||
|
||||
# Apply the final attention layer from the points to the image
|
||||
q = queries + point_embedding
|
||||
k = keys + image_pe
|
||||
attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
|
||||
queries = queries + attn_out
|
||||
queries = self.norm_final_attn(queries)
|
||||
|
||||
return queries, keys
|
||||
|
||||
|
||||
class TwoWayAttentionBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
num_heads: int,
|
||||
mlp_dim: int = 2048,
|
||||
activation: Type[nn.Module] = nn.ReLU,
|
||||
attention_downsample_rate: int = 2,
|
||||
skip_first_layer_pe: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
A transformer block with four layers: (1) self-attention of sparse
|
||||
inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp
|
||||
block on sparse inputs, and (4) cross attention of dense inputs to sparse
|
||||
inputs.
|
||||
|
||||
Arguments:
|
||||
embedding_dim (int): the channel dimension of the embeddings
|
||||
num_heads (int): the number of heads in the attention layers
|
||||
mlp_dim (int): the hidden dimension of the mlp block
|
||||
activation (nn.Module): the activation of the mlp block
|
||||
skip_first_layer_pe (bool): skip the PE on the first layer
|
||||
"""
|
||||
super().__init__()
|
||||
self.self_attn = Attention(embedding_dim, num_heads)
|
||||
self.norm1 = nn.LayerNorm(embedding_dim)
|
||||
|
||||
self.cross_attn_token_to_image = Attention(
|
||||
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
|
||||
)
|
||||
self.norm2 = nn.LayerNorm(embedding_dim)
|
||||
|
||||
self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
|
||||
self.norm3 = nn.LayerNorm(embedding_dim)
|
||||
|
||||
self.norm4 = nn.LayerNorm(embedding_dim)
|
||||
self.cross_attn_image_to_token = Attention(
|
||||
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
|
||||
)
|
||||
|
||||
self.skip_first_layer_pe = skip_first_layer_pe
|
||||
|
||||
def forward(
|
||||
self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
# Self attention block
|
||||
if self.skip_first_layer_pe:
|
||||
queries = self.self_attn(q=queries, k=queries, v=queries)
|
||||
else:
|
||||
q = queries + query_pe
|
||||
attn_out = self.self_attn(q=q, k=q, v=queries)
|
||||
queries = queries + attn_out
|
||||
queries = self.norm1(queries)
|
||||
|
||||
# Cross attention block, tokens attending to image embedding
|
||||
q = queries + query_pe
|
||||
k = keys + key_pe
|
||||
attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
|
||||
queries = queries + attn_out
|
||||
queries = self.norm2(queries)
|
||||
|
||||
# MLP block
|
||||
mlp_out = self.mlp(queries)
|
||||
queries = queries + mlp_out
|
||||
queries = self.norm3(queries)
|
||||
|
||||
# Cross attention block, image embedding attending to tokens
|
||||
q = queries + query_pe
|
||||
k = keys + key_pe
|
||||
attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)
|
||||
keys = keys + attn_out
|
||||
keys = self.norm4(keys)
|
||||
|
||||
return queries, keys
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
"""
|
||||
An attention layer that allows for downscaling the size of the embedding
|
||||
after projection to queries, keys, and values.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
num_heads: int,
|
||||
downsample_rate: int = 1,
|
||||
dropout: float = 0.0,
|
||||
kv_in_dim: int = None,
|
||||
use_fa3: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.embedding_dim = embedding_dim
|
||||
self.kv_in_dim = kv_in_dim if kv_in_dim is not None else embedding_dim
|
||||
self.internal_dim = embedding_dim // downsample_rate
|
||||
self.num_heads = num_heads
|
||||
self.use_fa3 = use_fa3
|
||||
assert (
|
||||
self.internal_dim % num_heads == 0
|
||||
), "num_heads must divide embedding_dim."
|
||||
|
||||
self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
|
||||
self.k_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
|
||||
self.v_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
|
||||
self.out_proj = nn.Linear(self.internal_dim, embedding_dim)
|
||||
|
||||
self.dropout_p = dropout
|
||||
|
||||
def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:
|
||||
b, n, c = x.shape
|
||||
x = x.reshape(b, n, num_heads, c // num_heads)
|
||||
return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head
|
||||
|
||||
def _recombine_heads(self, x: Tensor) -> Tensor:
|
||||
b, n_heads, n_tokens, c_per_head = x.shape
|
||||
x = x.transpose(1, 2)
|
||||
return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C
|
||||
|
||||
def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
|
||||
# Input projections
|
||||
q = self.q_proj(q)
|
||||
k = self.k_proj(k)
|
||||
v = self.v_proj(v)
|
||||
|
||||
# Separate into heads
|
||||
q = self._separate_heads(q, self.num_heads)
|
||||
k = self._separate_heads(k, self.num_heads)
|
||||
v = self._separate_heads(v, self.num_heads)
|
||||
|
||||
dropout_p = self.dropout_p if self.training else 0.0
|
||||
# Attention
|
||||
# with torch.backends.cuda.sdp_kernel(
|
||||
# enable_flash=USE_FLASH_ATTN,
|
||||
# # if Flash attention kernel is off, then math kernel needs to be enabled
|
||||
# enable_math=(OLD_GPU and dropout_p > 0.0) or MATH_KERNEL_ON,
|
||||
# enable_mem_efficient=OLD_GPU,
|
||||
# ):
|
||||
# Let's trust the dispatcher....
|
||||
if self.use_fa3:
|
||||
from sam3.perflib.fa3 import flash_attn_func
|
||||
|
||||
assert dropout_p == 0.0
|
||||
out = flash_attn_func(
|
||||
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
|
||||
).transpose(1, 2)
|
||||
else:
|
||||
torch.backends.cuda.enable_flash_sdp(True)
|
||||
torch.backends.cuda.enable_math_sdp(True)
|
||||
torch.backends.cuda.enable_mem_efficient_sdp(True)
|
||||
out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
|
||||
|
||||
out = self._recombine_heads(out)
|
||||
out = self.out_proj(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class RoPEAttention(Attention):
|
||||
"""Attention with rotary position encoding."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args,
|
||||
rope_theta=10000.0,
|
||||
# whether to repeat q rope to match k length
|
||||
# this is needed for cross-attention to memories
|
||||
rope_k_repeat=False,
|
||||
feat_sizes=(64, 64), # [w, h] for stride 16 feats at 1024 resolution
|
||||
use_rope_real=False,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.use_rope_real = use_rope_real
|
||||
self.compute_cis = partial(
|
||||
compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta
|
||||
)
|
||||
self.freqs_cis = self.compute_cis(
|
||||
end_x=feat_sizes[0], end_y=feat_sizes[1], device=None
|
||||
)
|
||||
if self.use_rope_real:
|
||||
self.freqs_cis_real = self.freqs_cis.real
|
||||
self.freqs_cis_imag = self.freqs_cis.imag
|
||||
self.rope_k_repeat = rope_k_repeat
|
||||
|
||||
def forward(
|
||||
self, q: Tensor, k: Tensor, v: Tensor, num_k_exclude_rope: int = 0
|
||||
) -> Tensor:
|
||||
# Input projections
|
||||
q = self.q_proj(q)
|
||||
k = self.k_proj(k)
|
||||
v = self.v_proj(v)
|
||||
|
||||
# Separate into heads
|
||||
q = self._separate_heads(q, self.num_heads)
|
||||
k = self._separate_heads(k, self.num_heads)
|
||||
v = self._separate_heads(v, self.num_heads)
|
||||
|
||||
# Apply rotary position encoding
|
||||
w = h = math.sqrt(q.shape[-2])
|
||||
if (
|
||||
self.freqs_cis.device != q.device
|
||||
or self.freqs_cis.shape[0] != q.shape[-2]
|
||||
):
|
||||
self.freqs_cis = self.compute_cis(end_x=w, end_y=h, device=q.device)
|
||||
if self.use_rope_real:
|
||||
self.freqs_cis_real = self.freqs_cis.real
|
||||
self.freqs_cis_imag = self.freqs_cis.imag
|
||||
if q.shape[-2] != k.shape[-2]:
|
||||
assert self.rope_k_repeat
|
||||
|
||||
num_k_rope = k.size(-2) - num_k_exclude_rope
|
||||
if self.use_rope_real:
|
||||
q, k[:, :, :num_k_rope] = apply_rotary_enc_real(
|
||||
q,
|
||||
k[:, :, :num_k_rope],
|
||||
freqs_cis_real=self.freqs_cis_real,
|
||||
freqs_cis_imag=self.freqs_cis_imag,
|
||||
repeat_freqs_k=self.rope_k_repeat,
|
||||
)
|
||||
else:
|
||||
q, k[:, :, :num_k_rope] = apply_rotary_enc(
|
||||
q,
|
||||
k[:, :, :num_k_rope],
|
||||
self.freqs_cis,
|
||||
repeat_freqs_k=self.rope_k_repeat,
|
||||
)
|
||||
|
||||
dropout_p = self.dropout_p if self.training else 0.0
|
||||
# Attention
|
||||
# with torch.backends.cuda.sdp_kernel(
|
||||
# enable_flash=USE_FLASH_ATTN,
|
||||
# # if Flash attention kernel is off, then math kernel needs to be enabled
|
||||
# enable_math=(OLD_GPU and dropout_p > 0.0) or MATH_KERNEL_ON,
|
||||
# enable_mem_efficient=OLD_GPU,
|
||||
# ):
|
||||
# Let's trust the dispatcher....
|
||||
if self.use_fa3:
|
||||
from sam3.perflib.fa3 import flash_attn_func
|
||||
|
||||
assert dropout_p == 0.0
|
||||
out = flash_attn_func(
|
||||
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
|
||||
).transpose(1, 2)
|
||||
else:
|
||||
torch.backends.cuda.enable_flash_sdp(True)
|
||||
torch.backends.cuda.enable_math_sdp(True)
|
||||
torch.backends.cuda.enable_mem_efficient_sdp(True)
|
||||
out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
|
||||
|
||||
out = self._recombine_heads(out)
|
||||
out = self.out_proj(out)
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,358 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
|
||||
import math
|
||||
from functools import partial
|
||||
from typing import Tuple, Type
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sam3.sam.rope import apply_rotary_enc, apply_rotary_enc_real, compute_axial_cis
|
||||
from torch import nn, Tensor
|
||||
|
||||
from .common import MLPBlock
|
||||
|
||||
|
||||
class TwoWayTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
depth: int,
|
||||
embedding_dim: int,
|
||||
num_heads: int,
|
||||
mlp_dim: int,
|
||||
activation: Type[nn.Module] = nn.ReLU,
|
||||
attention_downsample_rate: int = 2,
|
||||
) -> None:
|
||||
"""
|
||||
A transformer decoder that attends to an input image using
|
||||
queries whose positional embedding is supplied.
|
||||
|
||||
Args:
|
||||
depth (int): number of layers in the transformer
|
||||
embedding_dim (int): the channel dimension for the input embeddings
|
||||
num_heads (int): the number of heads for multihead attention. Must
|
||||
divide embedding_dim
|
||||
mlp_dim (int): the channel dimension internal to the MLP block
|
||||
activation (nn.Module): the activation to use in the MLP block
|
||||
"""
|
||||
super().__init__()
|
||||
self.depth = depth
|
||||
self.embedding_dim = embedding_dim
|
||||
self.num_heads = num_heads
|
||||
self.mlp_dim = mlp_dim
|
||||
self.layers = nn.ModuleList()
|
||||
|
||||
for i in range(depth):
|
||||
self.layers.append(
|
||||
TwoWayAttentionBlock(
|
||||
embedding_dim=embedding_dim,
|
||||
num_heads=num_heads,
|
||||
mlp_dim=mlp_dim,
|
||||
activation=activation,
|
||||
attention_downsample_rate=attention_downsample_rate,
|
||||
skip_first_layer_pe=(i == 0),
|
||||
)
|
||||
)
|
||||
|
||||
self.final_attn_token_to_image = Attention(
|
||||
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
|
||||
)
|
||||
self.norm_final_attn = nn.LayerNorm(embedding_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
image_embedding: Tensor,
|
||||
image_pe: Tensor,
|
||||
point_embedding: Tensor,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
"""
|
||||
Args:
|
||||
image_embedding (torch.Tensor): image to attend to. Should be shape
|
||||
B x embedding_dim x h x w for any h and w.
|
||||
image_pe (torch.Tensor): the positional encoding to add to the image. Must
|
||||
have the same shape as image_embedding.
|
||||
point_embedding (torch.Tensor): the embedding to add to the query points.
|
||||
Must have shape B x N_points x embedding_dim for any N_points.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: the processed point_embedding
|
||||
torch.Tensor: the processed image_embedding
|
||||
"""
|
||||
# BxCxHxW -> BxHWxC == B x N_image_tokens x C
|
||||
bs, c, h, w = image_embedding.shape
|
||||
image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
|
||||
image_pe = image_pe.flatten(2).permute(0, 2, 1)
|
||||
|
||||
# Prepare queries
|
||||
queries = point_embedding
|
||||
keys = image_embedding
|
||||
|
||||
# Apply transformer blocks and final layernorm
|
||||
for layer in self.layers:
|
||||
queries, keys = layer(
|
||||
queries=queries,
|
||||
keys=keys,
|
||||
query_pe=point_embedding,
|
||||
key_pe=image_pe,
|
||||
)
|
||||
|
||||
# Apply the final attention layer from the points to the image
|
||||
q = queries + point_embedding
|
||||
k = keys + image_pe
|
||||
attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
|
||||
queries = queries + attn_out
|
||||
queries = self.norm_final_attn(queries)
|
||||
|
||||
return queries, keys
|
||||
|
||||
|
||||
class TwoWayAttentionBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
num_heads: int,
|
||||
mlp_dim: int = 2048,
|
||||
activation: Type[nn.Module] = nn.ReLU,
|
||||
attention_downsample_rate: int = 2,
|
||||
skip_first_layer_pe: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
A transformer block with four layers: (1) self-attention of sparse
|
||||
inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp
|
||||
block on sparse inputs, and (4) cross attention of dense inputs to sparse
|
||||
inputs.
|
||||
|
||||
Arguments:
|
||||
embedding_dim (int): the channel dimension of the embeddings
|
||||
num_heads (int): the number of heads in the attention layers
|
||||
mlp_dim (int): the hidden dimension of the mlp block
|
||||
activation (nn.Module): the activation of the mlp block
|
||||
skip_first_layer_pe (bool): skip the PE on the first layer
|
||||
"""
|
||||
super().__init__()
|
||||
self.self_attn = Attention(embedding_dim, num_heads)
|
||||
self.norm1 = nn.LayerNorm(embedding_dim)
|
||||
|
||||
self.cross_attn_token_to_image = Attention(
|
||||
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
|
||||
)
|
||||
self.norm2 = nn.LayerNorm(embedding_dim)
|
||||
|
||||
self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
|
||||
self.norm3 = nn.LayerNorm(embedding_dim)
|
||||
|
||||
self.norm4 = nn.LayerNorm(embedding_dim)
|
||||
self.cross_attn_image_to_token = Attention(
|
||||
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
|
||||
)
|
||||
|
||||
self.skip_first_layer_pe = skip_first_layer_pe
|
||||
|
||||
def forward(
|
||||
self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
# Self attention block
|
||||
if self.skip_first_layer_pe:
|
||||
queries = self.self_attn(q=queries, k=queries, v=queries)
|
||||
else:
|
||||
q = queries + query_pe
|
||||
attn_out = self.self_attn(q=q, k=q, v=queries)
|
||||
queries = queries + attn_out
|
||||
queries = self.norm1(queries)
|
||||
|
||||
# Cross attention block, tokens attending to image embedding
|
||||
q = queries + query_pe
|
||||
k = keys + key_pe
|
||||
attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
|
||||
queries = queries + attn_out
|
||||
queries = self.norm2(queries)
|
||||
|
||||
# MLP block
|
||||
mlp_out = self.mlp(queries)
|
||||
queries = queries + mlp_out
|
||||
queries = self.norm3(queries)
|
||||
|
||||
# Cross attention block, image embedding attending to tokens
|
||||
q = queries + query_pe
|
||||
k = keys + key_pe
|
||||
attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)
|
||||
keys = keys + attn_out
|
||||
keys = self.norm4(keys)
|
||||
|
||||
return queries, keys
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
"""
|
||||
An attention layer that allows for downscaling the size of the embedding
|
||||
after projection to queries, keys, and values.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
num_heads: int,
|
||||
downsample_rate: int = 1,
|
||||
dropout: float = 0.0,
|
||||
kv_in_dim: int = None,
|
||||
use_fa3: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.embedding_dim = embedding_dim
|
||||
self.kv_in_dim = kv_in_dim if kv_in_dim is not None else embedding_dim
|
||||
self.internal_dim = embedding_dim // downsample_rate
|
||||
self.num_heads = num_heads
|
||||
self.use_fa3 = use_fa3
|
||||
assert (
|
||||
self.internal_dim % num_heads == 0
|
||||
), "num_heads must divide embedding_dim."
|
||||
|
||||
self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
|
||||
self.k_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
|
||||
self.v_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
|
||||
self.out_proj = nn.Linear(self.internal_dim, embedding_dim)
|
||||
|
||||
self.dropout_p = dropout
|
||||
|
||||
def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:
|
||||
b, n, c = x.shape
|
||||
x = x.reshape(b, n, num_heads, c // num_heads)
|
||||
return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head
|
||||
|
||||
def _recombine_heads(self, x: Tensor) -> Tensor:
|
||||
b, n_heads, n_tokens, c_per_head = x.shape
|
||||
x = x.transpose(1, 2)
|
||||
return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C
|
||||
|
||||
def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
|
||||
# Input projections
|
||||
q = self.q_proj(q)
|
||||
k = self.k_proj(k)
|
||||
v = self.v_proj(v)
|
||||
|
||||
# Separate into heads
|
||||
q = self._separate_heads(q, self.num_heads)
|
||||
k = self._separate_heads(k, self.num_heads)
|
||||
v = self._separate_heads(v, self.num_heads)
|
||||
|
||||
dropout_p = self.dropout_p if self.training else 0.0
|
||||
# Attention
|
||||
# with torch.backends.cuda.sdp_kernel(
|
||||
# enable_flash=USE_FLASH_ATTN,
|
||||
# # if Flash attention kernel is off, then math kernel needs to be enabled
|
||||
# enable_math=(OLD_GPU and dropout_p > 0.0) or MATH_KERNEL_ON,
|
||||
# enable_mem_efficient=OLD_GPU,
|
||||
# ):
|
||||
# Let's trust the dispatcher....
|
||||
if self.use_fa3:
|
||||
from sam3.perflib.fa3 import flash_attn_func
|
||||
|
||||
assert dropout_p == 0.0
|
||||
out = flash_attn_func(
|
||||
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
|
||||
).transpose(1, 2)
|
||||
else:
|
||||
torch.backends.cuda.enable_flash_sdp(True)
|
||||
torch.backends.cuda.enable_math_sdp(True)
|
||||
torch.backends.cuda.enable_mem_efficient_sdp(True)
|
||||
out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
|
||||
|
||||
out = self._recombine_heads(out)
|
||||
out = self.out_proj(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class RoPEAttention(Attention):
|
||||
"""Attention with rotary position encoding."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args,
|
||||
rope_theta=10000.0,
|
||||
# whether to repeat q rope to match k length
|
||||
# this is needed for cross-attention to memories
|
||||
rope_k_repeat=False,
|
||||
feat_sizes=(64, 64), # [w, h] for stride 16 feats at 1024 resolution
|
||||
use_rope_real=False,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.use_rope_real = use_rope_real
|
||||
self.compute_cis = partial(
|
||||
compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta
|
||||
)
|
||||
device = torch.device("cuda") if torch.cuda.is_available() else None
|
||||
self.freqs_cis = self.compute_cis(
|
||||
end_x=feat_sizes[0], end_y=feat_sizes[1], device=device
|
||||
)
|
||||
if self.use_rope_real:
|
||||
self.freqs_cis_real = self.freqs_cis.real
|
||||
self.freqs_cis_imag = self.freqs_cis.imag
|
||||
self.rope_k_repeat = rope_k_repeat
|
||||
|
||||
def forward(
|
||||
self, q: Tensor, k: Tensor, v: Tensor, num_k_exclude_rope: int = 0
|
||||
) -> Tensor:
|
||||
# Input projections
|
||||
q = self.q_proj(q)
|
||||
k = self.k_proj(k)
|
||||
v = self.v_proj(v)
|
||||
|
||||
# Separate into heads
|
||||
q = self._separate_heads(q, self.num_heads)
|
||||
k = self._separate_heads(k, self.num_heads)
|
||||
v = self._separate_heads(v, self.num_heads)
|
||||
|
||||
# Apply rotary position encoding
|
||||
w = h = math.sqrt(q.shape[-2])
|
||||
if self.freqs_cis.shape[0] != q.shape[-2]:
|
||||
self.freqs_cis = self.compute_cis(end_x=w, end_y=h, device=q.device)
|
||||
self.freqs_cis_real = self.freqs_cis.real
|
||||
self.freqs_cis_imag = self.freqs_cis.imag
|
||||
if q.shape[-2] != k.shape[-2]:
|
||||
assert self.rope_k_repeat
|
||||
|
||||
num_k_rope = k.size(-2) - num_k_exclude_rope
|
||||
if self.use_rope_real:
|
||||
q, k[:, :, :num_k_rope] = apply_rotary_enc_real(
|
||||
q,
|
||||
k[:, :, :num_k_rope],
|
||||
freqs_cis_real=self.freqs_cis_real,
|
||||
freqs_cis_imag=self.freqs_cis_imag,
|
||||
repeat_freqs_k=self.rope_k_repeat,
|
||||
)
|
||||
else:
|
||||
q, k[:, :, :num_k_rope] = apply_rotary_enc(
|
||||
q,
|
||||
k[:, :, :num_k_rope],
|
||||
self.freqs_cis,
|
||||
repeat_freqs_k=self.rope_k_repeat,
|
||||
)
|
||||
|
||||
dropout_p = self.dropout_p if self.training else 0.0
|
||||
# Attention
|
||||
# with torch.backends.cuda.sdp_kernel(
|
||||
# enable_flash=USE_FLASH_ATTN,
|
||||
# # if Flash attention kernel is off, then math kernel needs to be enabled
|
||||
# enable_math=(OLD_GPU and dropout_p > 0.0) or MATH_KERNEL_ON,
|
||||
# enable_mem_efficient=OLD_GPU,
|
||||
# ):
|
||||
# Let's trust the dispatcher....
|
||||
if self.use_fa3:
|
||||
from sam3.perflib.fa3 import flash_attn_func
|
||||
|
||||
assert dropout_p == 0.0
|
||||
out = flash_attn_func(
|
||||
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
|
||||
).transpose(1, 2)
|
||||
else:
|
||||
torch.backends.cuda.enable_flash_sdp(True)
|
||||
torch.backends.cuda.enable_math_sdp(True)
|
||||
torch.backends.cuda.enable_mem_efficient_sdp(True)
|
||||
out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
|
||||
|
||||
out = self._recombine_heads(out)
|
||||
out = self.out_proj(out)
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,941 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import matplotlib.patches as patches
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pycocotools.mask as mask_utils
|
||||
import torch
|
||||
from matplotlib.colors import to_rgb
|
||||
from PIL import Image
|
||||
from skimage.color import lab2rgb, rgb2lab
|
||||
from sklearn.cluster import KMeans
|
||||
from torchvision.ops import masks_to_boxes
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def generate_colors(n_colors=256, n_samples=5000):
|
||||
# Step 1: Random RGB samples
|
||||
np.random.seed(42)
|
||||
rgb = np.random.rand(n_samples, 3)
|
||||
# Step 2: Convert to LAB for perceptual uniformity
|
||||
# print(f"Converting {n_samples} RGB samples to LAB color space...")
|
||||
lab = rgb2lab(rgb.reshape(1, -1, 3)).reshape(-1, 3)
|
||||
# print("Conversion to LAB complete.")
|
||||
# Step 3: k-means clustering in LAB
|
||||
kmeans = KMeans(n_clusters=n_colors, n_init=10)
|
||||
# print(f"Fitting KMeans with {n_colors} clusters on {n_samples} samples...")
|
||||
kmeans.fit(lab)
|
||||
# print("KMeans fitting complete.")
|
||||
centers_lab = kmeans.cluster_centers_
|
||||
# Step 4: Convert LAB back to RGB
|
||||
colors_rgb = lab2rgb(centers_lab.reshape(1, -1, 3)).reshape(-1, 3)
|
||||
colors_rgb = np.clip(colors_rgb, 0, 1)
|
||||
return colors_rgb
|
||||
|
||||
|
||||
COLORS = generate_colors(n_colors=128, n_samples=5000)
|
||||
|
||||
|
||||
def show_img_tensor(img_batch, vis_img_idx=0):
|
||||
MEAN_IMG = np.array([0.485, 0.456, 0.406])
|
||||
STD_IMG = np.array([0.229, 0.224, 0.225])
|
||||
im_tensor = img_batch[vis_img_idx].detach().cpu()
|
||||
assert im_tensor.dim() == 3
|
||||
im_tensor = im_tensor.numpy().transpose((1, 2, 0))
|
||||
im_tensor = (im_tensor * STD_IMG) + MEAN_IMG
|
||||
im_tensor = np.clip(im_tensor, 0, 1)
|
||||
plt.imshow(im_tensor)
|
||||
|
||||
|
||||
def draw_box_on_image(image, box, color=(0, 255, 0)):
|
||||
"""
|
||||
Draws a rectangle on a given PIL image using the provided box coordinates in xywh format.
|
||||
:param image: PIL.Image - The image on which to draw the rectangle.
|
||||
:param box: tuple - A tuple (x, y, w, h) representing the top-left corner, width, and height of the rectangle.
|
||||
:param color: tuple - A tuple (R, G, B) representing the color of the rectangle. Default is red.
|
||||
:return: PIL.Image - The image with the rectangle drawn on it.
|
||||
"""
|
||||
# Ensure the image is in RGB mode
|
||||
image = image.convert("RGB")
|
||||
# Unpack the box coordinates
|
||||
x, y, w, h = box
|
||||
x, y, w, h = int(x), int(y), int(w), int(h)
|
||||
# Get the pixel data
|
||||
pixels = image.load()
|
||||
# Draw the top and bottom edges
|
||||
for i in range(x, x + w):
|
||||
pixels[i, y] = color
|
||||
pixels[i, y + h - 1] = color
|
||||
pixels[i, y + 1] = color
|
||||
pixels[i, y + h] = color
|
||||
pixels[i, y - 1] = color
|
||||
pixels[i, y + h - 2] = color
|
||||
# Draw the left and right edges
|
||||
for j in range(y, y + h):
|
||||
pixels[x, j] = color
|
||||
pixels[x + 1, j] = color
|
||||
pixels[x - 1, j] = color
|
||||
pixels[x + w - 1, j] = color
|
||||
pixels[x + w, j] = color
|
||||
pixels[x + w - 2, j] = color
|
||||
return image
|
||||
|
||||
|
||||
def plot_bbox(
|
||||
img_height,
|
||||
img_width,
|
||||
box,
|
||||
box_format="XYXY",
|
||||
relative_coords=True,
|
||||
color="r",
|
||||
linestyle="solid",
|
||||
text=None,
|
||||
ax=None,
|
||||
):
|
||||
if box_format == "XYXY":
|
||||
x, y, x2, y2 = box
|
||||
w = x2 - x
|
||||
h = y2 - y
|
||||
elif box_format == "XYWH":
|
||||
x, y, w, h = box
|
||||
elif box_format == "CxCyWH":
|
||||
cx, cy, w, h = box
|
||||
x = cx - w / 2
|
||||
y = cy - h / 2
|
||||
else:
|
||||
raise RuntimeError(f"Invalid box_format {box_format}")
|
||||
|
||||
if relative_coords:
|
||||
x *= img_width
|
||||
w *= img_width
|
||||
y *= img_height
|
||||
h *= img_height
|
||||
|
||||
if ax is None:
|
||||
ax = plt.gca()
|
||||
rect = patches.Rectangle(
|
||||
(x, y),
|
||||
w,
|
||||
h,
|
||||
linewidth=1.5,
|
||||
edgecolor=color,
|
||||
facecolor="none",
|
||||
linestyle=linestyle,
|
||||
)
|
||||
ax.add_patch(rect)
|
||||
if text is not None:
|
||||
facecolor = "w"
|
||||
ax.text(
|
||||
x,
|
||||
y - 5,
|
||||
text,
|
||||
color=color,
|
||||
weight="bold",
|
||||
fontsize=8,
|
||||
bbox={"facecolor": facecolor, "alpha": 0.75, "pad": 2},
|
||||
)
|
||||
|
||||
|
||||
def plot_mask(mask, color="r", ax=None):
|
||||
im_h, im_w = mask.shape
|
||||
mask_img = np.zeros((im_h, im_w, 4), dtype=np.float32)
|
||||
mask_img[..., :3] = to_rgb(color)
|
||||
mask_img[..., 3] = mask * 0.5
|
||||
# Use the provided ax or the current axis
|
||||
if ax is None:
|
||||
ax = plt.gca()
|
||||
ax.imshow(mask_img)
|
||||
|
||||
|
||||
def normalize_bbox(bbox_xywh, img_w, img_h):
|
||||
# Assumes bbox_xywh is in XYWH format
|
||||
if isinstance(bbox_xywh, list):
|
||||
assert (
|
||||
len(bbox_xywh) == 4
|
||||
), "bbox_xywh list must have 4 elements. Batching not support except for torch tensors."
|
||||
normalized_bbox = bbox_xywh.copy()
|
||||
normalized_bbox[0] /= img_w
|
||||
normalized_bbox[1] /= img_h
|
||||
normalized_bbox[2] /= img_w
|
||||
normalized_bbox[3] /= img_h
|
||||
else:
|
||||
assert isinstance(
|
||||
bbox_xywh, torch.Tensor
|
||||
), "Only torch tensors are supported for batching."
|
||||
normalized_bbox = bbox_xywh.clone()
|
||||
assert (
|
||||
normalized_bbox.size(-1) == 4
|
||||
), "bbox_xywh tensor must have last dimension of size 4."
|
||||
normalized_bbox[..., 0] /= img_w
|
||||
normalized_bbox[..., 1] /= img_h
|
||||
normalized_bbox[..., 2] /= img_w
|
||||
normalized_bbox[..., 3] /= img_h
|
||||
return normalized_bbox
|
||||
|
||||
|
||||
def visualize_frame_output(frame_idx, video_frames, outputs, figsize=(12, 8)):
|
||||
plt.figure(figsize=figsize)
|
||||
plt.title(f"frame {frame_idx}")
|
||||
img = load_frame(video_frames[frame_idx])
|
||||
img_H, img_W, _ = img.shape
|
||||
plt.imshow(img)
|
||||
for i in range(len(outputs["out_probs"])):
|
||||
box_xywh = outputs["out_boxes_xywh"][i]
|
||||
prob = outputs["out_probs"][i]
|
||||
obj_id = outputs["out_obj_ids"][i]
|
||||
binary_mask = outputs["out_binary_masks"][i]
|
||||
color = COLORS[obj_id % len(COLORS)]
|
||||
plot_bbox(
|
||||
img_H,
|
||||
img_W,
|
||||
box_xywh,
|
||||
text=f"(id={obj_id}, {prob=:.2f})",
|
||||
box_format="XYWH",
|
||||
color=color,
|
||||
)
|
||||
plot_mask(binary_mask, color=color)
|
||||
|
||||
|
||||
def visualize_formatted_frame_output(
|
||||
frame_idx,
|
||||
video_frames,
|
||||
outputs_list,
|
||||
titles=None,
|
||||
points_list=None,
|
||||
points_labels_list=None,
|
||||
figsize=(12, 8),
|
||||
title_suffix="",
|
||||
prompt_info=None,
|
||||
):
|
||||
"""Visualize up to three sets of segmentation masks on a video frame.
|
||||
|
||||
Args:
|
||||
frame_idx: Frame index to visualize
|
||||
image_files: List of image file paths
|
||||
outputs_list: List of {frame_idx: {obj_id: mask_tensor}} or single dict {obj_id: mask_tensor}
|
||||
titles: List of titles for each set of outputs_list
|
||||
points_list: Optional list of point coordinates
|
||||
points_labels_list: Optional list of point labels
|
||||
figsize: Figure size tuple
|
||||
save: Whether to save the visualization to file
|
||||
output_dir: Base output directory when saving
|
||||
scenario_name: Scenario name for organizing saved files
|
||||
title_suffix: Additional title suffix
|
||||
prompt_info: Dictionary with prompt information (boxes, points, etc.)
|
||||
"""
|
||||
# Handle single output dict case
|
||||
if isinstance(outputs_list, dict) and frame_idx in outputs_list:
|
||||
# This is a single outputs dict with frame indices as keys
|
||||
outputs_list = [outputs_list]
|
||||
elif isinstance(outputs_list, dict) and not any(
|
||||
isinstance(k, int) for k in outputs_list.keys()
|
||||
):
|
||||
# This is a single frame's outputs {obj_id: mask}
|
||||
single_frame_outputs = {frame_idx: outputs_list}
|
||||
outputs_list = [single_frame_outputs]
|
||||
|
||||
num_outputs = len(outputs_list)
|
||||
if titles is None:
|
||||
titles = [f"Set {i+1}" for i in range(num_outputs)]
|
||||
assert (
|
||||
len(titles) == num_outputs
|
||||
), "length of `titles` should match that of `outputs_list` if not None."
|
||||
|
||||
_, axes = plt.subplots(1, num_outputs, figsize=figsize)
|
||||
if num_outputs == 1:
|
||||
axes = [axes] # Make it iterable
|
||||
|
||||
img = load_frame(video_frames[frame_idx])
|
||||
img_H, img_W, _ = img.shape
|
||||
|
||||
for idx in range(num_outputs):
|
||||
ax, outputs_set, ax_title = axes[idx], outputs_list[idx], titles[idx]
|
||||
ax.set_title(f"Frame {frame_idx} - {ax_title}{title_suffix}")
|
||||
ax.imshow(img)
|
||||
|
||||
if frame_idx in outputs_set:
|
||||
_outputs = outputs_set[frame_idx]
|
||||
else:
|
||||
print(f"Warning: Frame {frame_idx} not found in outputs_set")
|
||||
continue
|
||||
|
||||
if prompt_info and frame_idx == 0: # Show prompts on first frame
|
||||
if "boxes" in prompt_info:
|
||||
for box in prompt_info["boxes"]:
|
||||
# box is in [x, y, w, h] normalized format
|
||||
x, y, w, h = box
|
||||
plot_bbox(
|
||||
img_H,
|
||||
img_W,
|
||||
[x, y, x + w, y + h], # Convert to XYXY
|
||||
box_format="XYXY",
|
||||
relative_coords=True,
|
||||
color="yellow",
|
||||
linestyle="dashed",
|
||||
text="PROMPT BOX",
|
||||
ax=ax,
|
||||
)
|
||||
|
||||
if "points" in prompt_info and "point_labels" in prompt_info:
|
||||
points = np.array(prompt_info["points"])
|
||||
labels = np.array(prompt_info["point_labels"])
|
||||
# Convert normalized to pixel coordinates
|
||||
points_pixel = points * np.array([img_W, img_H])
|
||||
|
||||
# Draw positive points (green stars)
|
||||
pos_points = points_pixel[labels == 1]
|
||||
if len(pos_points) > 0:
|
||||
ax.scatter(
|
||||
pos_points[:, 0],
|
||||
pos_points[:, 1],
|
||||
color="lime",
|
||||
marker="*",
|
||||
s=200,
|
||||
edgecolor="white",
|
||||
linewidth=2,
|
||||
label="Positive Points",
|
||||
zorder=10,
|
||||
)
|
||||
|
||||
# Draw negative points (red stars)
|
||||
neg_points = points_pixel[labels == 0]
|
||||
if len(neg_points) > 0:
|
||||
ax.scatter(
|
||||
neg_points[:, 0],
|
||||
neg_points[:, 1],
|
||||
color="red",
|
||||
marker="*",
|
||||
s=200,
|
||||
edgecolor="white",
|
||||
linewidth=2,
|
||||
label="Negative Points",
|
||||
zorder=10,
|
||||
)
|
||||
|
||||
objects_drawn = 0
|
||||
for obj_id, binary_mask in _outputs.items():
|
||||
mask_sum = (
|
||||
binary_mask.sum()
|
||||
if hasattr(binary_mask, "sum")
|
||||
else np.sum(binary_mask)
|
||||
)
|
||||
|
||||
if mask_sum > 0: # Only draw if mask has content
|
||||
# Convert to torch tensor if it's not already
|
||||
if not isinstance(binary_mask, torch.Tensor):
|
||||
binary_mask = torch.tensor(binary_mask)
|
||||
|
||||
# Find bounding box from mask
|
||||
if binary_mask.any():
|
||||
box_xyxy = masks_to_boxes(binary_mask.unsqueeze(0)).squeeze()
|
||||
box_xyxy = normalize_bbox(box_xyxy, img_W, img_H)
|
||||
else:
|
||||
# Fallback: create a small box at center
|
||||
box_xyxy = [0.45, 0.45, 0.55, 0.55]
|
||||
|
||||
color = COLORS[obj_id % len(COLORS)]
|
||||
|
||||
plot_bbox(
|
||||
img_H,
|
||||
img_W,
|
||||
box_xyxy,
|
||||
text=f"(id={obj_id})",
|
||||
box_format="XYXY",
|
||||
color=color,
|
||||
ax=ax,
|
||||
)
|
||||
|
||||
# Convert back to numpy for plotting
|
||||
mask_np = (
|
||||
binary_mask.numpy()
|
||||
if isinstance(binary_mask, torch.Tensor)
|
||||
else binary_mask
|
||||
)
|
||||
plot_mask(mask_np, color=color, ax=ax)
|
||||
objects_drawn += 1
|
||||
|
||||
if objects_drawn == 0:
|
||||
ax.text(
|
||||
0.5,
|
||||
0.5,
|
||||
"No objects detected",
|
||||
transform=ax.transAxes,
|
||||
fontsize=16,
|
||||
ha="center",
|
||||
va="center",
|
||||
color="red",
|
||||
weight="bold",
|
||||
)
|
||||
|
||||
# Draw additional points if provided
|
||||
if points_list is not None and points_list[idx] is not None:
|
||||
show_points(
|
||||
points_list[idx], points_labels_list[idx], ax=ax, marker_size=200
|
||||
)
|
||||
|
||||
ax.axis("off")
|
||||
|
||||
plt.tight_layout()
|
||||
plt.show()
|
||||
|
||||
|
||||
def render_masklet_frame(img, outputs, frame_idx=None, alpha=0.5):
|
||||
"""
|
||||
Overlays masklets and bounding boxes on a single image frame.
|
||||
Args:
|
||||
img: np.ndarray, shape (H, W, 3), uint8 or float32 in [0,255] or [0,1]
|
||||
outputs: dict with keys: out_boxes_xywh, out_probs, out_obj_ids, out_binary_masks
|
||||
frame_idx: int or None, for overlaying frame index text
|
||||
alpha: float, mask overlay alpha
|
||||
Returns:
|
||||
overlay: np.ndarray, shape (H, W, 3), uint8
|
||||
"""
|
||||
if img.dtype == np.float32 or img.max() <= 1.0:
|
||||
img = (img * 255).astype(np.uint8)
|
||||
img = img[..., :3] # drop alpha if present
|
||||
height, width = img.shape[:2]
|
||||
overlay = img.copy()
|
||||
|
||||
for i in range(len(outputs["out_probs"])):
|
||||
obj_id = outputs["out_obj_ids"][i]
|
||||
color = COLORS[obj_id % len(COLORS)]
|
||||
color255 = (color * 255).astype(np.uint8)
|
||||
mask = outputs["out_binary_masks"][i]
|
||||
if mask.shape != img.shape[:2]:
|
||||
mask = cv2.resize(
|
||||
mask.astype(np.float32),
|
||||
(img.shape[1], img.shape[0]),
|
||||
interpolation=cv2.INTER_NEAREST,
|
||||
)
|
||||
mask_bool = mask > 0.5
|
||||
for c in range(3):
|
||||
overlay[..., c][mask_bool] = (
|
||||
alpha * color255[c] + (1 - alpha) * overlay[..., c][mask_bool]
|
||||
).astype(np.uint8)
|
||||
|
||||
# Draw bounding boxes and text
|
||||
for i in range(len(outputs["out_probs"])):
|
||||
box_xywh = outputs["out_boxes_xywh"][i]
|
||||
obj_id = outputs["out_obj_ids"][i]
|
||||
prob = outputs["out_probs"][i]
|
||||
color = COLORS[obj_id % len(COLORS)]
|
||||
color255 = tuple(int(x * 255) for x in color)
|
||||
x, y, w, h = box_xywh
|
||||
x1 = int(x * width)
|
||||
y1 = int(y * height)
|
||||
x2 = int((x + w) * width)
|
||||
y2 = int((y + h) * height)
|
||||
cv2.rectangle(overlay, (x1, y1), (x2, y2), color255, 2)
|
||||
if prob is not None:
|
||||
label = f"id={obj_id}, p={prob:.2f}"
|
||||
else:
|
||||
label = f"id={obj_id}"
|
||||
cv2.putText(
|
||||
overlay,
|
||||
label,
|
||||
(x1, max(y1 - 10, 0)),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.5,
|
||||
color255,
|
||||
1,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
# Overlay frame index at the top-left corner
|
||||
if frame_idx is not None:
|
||||
cv2.putText(
|
||||
overlay,
|
||||
f"Frame {frame_idx}",
|
||||
(10, 30),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
1.0,
|
||||
(255, 255, 255),
|
||||
2,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
return overlay
|
||||
|
||||
|
||||
def save_masklet_video(video_frames, outputs, out_path, alpha=0.5, fps=10):
|
||||
# Each outputs dict has keys: "out_boxes_xywh", "out_probs", "out_obj_ids", "out_binary_masks"
|
||||
# video_frames: list of video frame data, same length as outputs_list
|
||||
|
||||
# Read first frame to get size
|
||||
first_img = load_frame(video_frames[0])
|
||||
height, width = first_img.shape[:2]
|
||||
if first_img.dtype == np.float32 or first_img.max() <= 1.0:
|
||||
first_img = (first_img * 255).astype(np.uint8)
|
||||
# Use 'mp4v' for best compatibility with VSCode playback (.mp4 files)
|
||||
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
writer = cv2.VideoWriter("temp.mp4", fourcc, fps, (width, height))
|
||||
|
||||
outputs_list = [
|
||||
(video_frames[frame_idx], frame_idx, outputs[frame_idx])
|
||||
for frame_idx in sorted(outputs.keys())
|
||||
]
|
||||
|
||||
for frame, frame_idx, frame_outputs in tqdm(outputs_list):
|
||||
img = load_frame(frame)
|
||||
overlay = render_masklet_frame(
|
||||
img, frame_outputs, frame_idx=frame_idx, alpha=alpha
|
||||
)
|
||||
writer.write(cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
|
||||
|
||||
writer.release()
|
||||
|
||||
# Re-encode the video for VSCode compatibility using ffmpeg
|
||||
subprocess.run(["ffmpeg", "-y", "-i", "temp.mp4", out_path])
|
||||
print(f"Re-encoded video saved to {out_path}")
|
||||
|
||||
os.remove("temp.mp4") # Clean up temporary file
|
||||
|
||||
|
||||
def save_masklet_image(frame, outputs, out_path, alpha=0.5, frame_idx=None):
|
||||
"""
|
||||
Save a single image with masklet overlays.
|
||||
"""
|
||||
img = load_frame(frame)
|
||||
overlay = render_masklet_frame(img, outputs, frame_idx=frame_idx, alpha=alpha)
|
||||
Image.fromarray(overlay).save(out_path)
|
||||
print(f"Overlay image saved to {out_path}")
|
||||
|
||||
|
||||
def prepare_masks_for_visualization(frame_to_output):
|
||||
# frame_to_obj_masks --> {frame_idx: {'output_probs': np.array, `out_obj_ids`: np.array, `out_binary_masks`: np.array}}
|
||||
for frame_idx, out in frame_to_output.items():
|
||||
_processed_out = {}
|
||||
for idx, obj_id in enumerate(out["out_obj_ids"].tolist()):
|
||||
if out["out_binary_masks"][idx].any():
|
||||
_processed_out[obj_id] = out["out_binary_masks"][idx]
|
||||
frame_to_output[frame_idx] = _processed_out
|
||||
return frame_to_output
|
||||
|
||||
|
||||
def convert_coco_to_masklet_format(
|
||||
annotations, img_info, is_prediction=False, score_threshold=0.5
|
||||
):
|
||||
"""
|
||||
Convert COCO format annotations to format expected by render_masklet_frame
|
||||
"""
|
||||
outputs = {
|
||||
"out_boxes_xywh": [],
|
||||
"out_probs": [],
|
||||
"out_obj_ids": [],
|
||||
"out_binary_masks": [],
|
||||
}
|
||||
|
||||
img_h, img_w = img_info["height"], img_info["width"]
|
||||
|
||||
for idx, ann in enumerate(annotations):
|
||||
# Get bounding box in relative XYWH format
|
||||
if "bbox" in ann:
|
||||
bbox = ann["bbox"]
|
||||
if max(bbox) > 1.0: # Convert absolute to relative coordinates
|
||||
bbox = [
|
||||
bbox[0] / img_w,
|
||||
bbox[1] / img_h,
|
||||
bbox[2] / img_w,
|
||||
bbox[3] / img_h,
|
||||
]
|
||||
else:
|
||||
mask = mask_utils.decode(ann["segmentation"])
|
||||
rows = np.any(mask, axis=1)
|
||||
cols = np.any(mask, axis=0)
|
||||
if np.any(rows) and np.any(cols):
|
||||
rmin, rmax = np.where(rows)[0][[0, -1]]
|
||||
cmin, cmax = np.where(cols)[0][[0, -1]]
|
||||
# Convert to relative XYWH
|
||||
bbox = [
|
||||
cmin / img_w,
|
||||
rmin / img_h,
|
||||
(cmax - cmin + 1) / img_w,
|
||||
(rmax - rmin + 1) / img_h,
|
||||
]
|
||||
else:
|
||||
bbox = [0, 0, 0, 0]
|
||||
|
||||
outputs["out_boxes_xywh"].append(bbox)
|
||||
|
||||
# Get probability/score
|
||||
if is_prediction:
|
||||
prob = ann["score"]
|
||||
else:
|
||||
prob = 1.0 # GT has no probability
|
||||
outputs["out_probs"].append(prob)
|
||||
|
||||
outputs["out_obj_ids"].append(idx)
|
||||
mask = mask_utils.decode(ann["segmentation"])
|
||||
mask = (mask > score_threshold).astype(np.uint8)
|
||||
|
||||
outputs["out_binary_masks"].append(mask)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
def save_side_by_side_visualization(img, gt_anns, pred_anns, noun_phrase):
|
||||
"""
|
||||
Create side-by-side visualization of GT and predictions
|
||||
"""
|
||||
|
||||
# Create side-by-side visualization
|
||||
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 7))
|
||||
|
||||
main_title = f"Noun phrase: '{noun_phrase}'"
|
||||
fig.suptitle(main_title, fontsize=16, fontweight="bold")
|
||||
|
||||
gt_overlay = render_masklet_frame(img, gt_anns, alpha=0.5)
|
||||
ax1.imshow(gt_overlay)
|
||||
ax1.set_title("Ground Truth", fontsize=14, fontweight="bold")
|
||||
ax1.axis("off")
|
||||
|
||||
pred_overlay = render_masklet_frame(img, pred_anns, alpha=0.5)
|
||||
ax2.imshow(pred_overlay)
|
||||
ax2.set_title("Predictions", fontsize=14, fontweight="bold")
|
||||
ax2.axis("off")
|
||||
|
||||
plt.subplots_adjust(top=0.88)
|
||||
plt.tight_layout()
|
||||
|
||||
|
||||
def bitget(val, idx):
|
||||
return (val >> idx) & 1
|
||||
|
||||
|
||||
def pascal_color_map():
|
||||
colormap = np.zeros((512, 3), dtype=int)
|
||||
ind = np.arange(512, dtype=int)
|
||||
for shift in reversed(list(range(8))):
|
||||
for channel in range(3):
|
||||
colormap[:, channel] |= bitget(ind, channel) << shift
|
||||
ind >>= 3
|
||||
|
||||
return colormap.astype(np.uint8)
|
||||
|
||||
|
||||
def draw_masks_to_frame(
|
||||
frame: np.ndarray, masks: np.ndarray, colors: np.ndarray
|
||||
) -> np.ndarray:
|
||||
masked_frame = frame
|
||||
for mask, color in zip(masks, colors):
|
||||
curr_masked_frame = np.where(mask[..., None], color, masked_frame)
|
||||
masked_frame = cv2.addWeighted(masked_frame, 0.75, curr_masked_frame, 0.25, 0)
|
||||
|
||||
if int(cv2.__version__[0]) > 3:
|
||||
contours, _ = cv2.findContours(
|
||||
np.array(mask, dtype=np.uint8).copy(),
|
||||
cv2.RETR_TREE,
|
||||
cv2.CHAIN_APPROX_NONE,
|
||||
)
|
||||
else:
|
||||
_, contours, _ = cv2.findContours(
|
||||
np.array(mask, dtype=np.uint8).copy(),
|
||||
cv2.RETR_TREE,
|
||||
cv2.CHAIN_APPROX_NONE,
|
||||
)
|
||||
|
||||
cv2.drawContours(
|
||||
masked_frame, contours, -1, (255, 255, 255), 7
|
||||
) # White outer contour
|
||||
cv2.drawContours(
|
||||
masked_frame, contours, -1, (0, 0, 0), 5
|
||||
) # Black middle contour
|
||||
cv2.drawContours(
|
||||
masked_frame, contours, -1, color.tolist(), 3
|
||||
) # Original color inner contour
|
||||
return masked_frame
|
||||
|
||||
|
||||
def get_annot_df(file_path: str):
|
||||
with open(file_path, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
dfs = {}
|
||||
|
||||
for k, v in data.items():
|
||||
if k in ("info", "licenses"):
|
||||
dfs[k] = v
|
||||
continue
|
||||
df = pd.DataFrame(v)
|
||||
dfs[k] = df
|
||||
|
||||
return dfs
|
||||
|
||||
|
||||
def get_annot_dfs(file_list: list[str]):
|
||||
dfs = {}
|
||||
for annot_file in tqdm(file_list):
|
||||
dataset_name = Path(annot_file).stem
|
||||
dfs[dataset_name] = get_annot_df(annot_file)
|
||||
|
||||
return dfs
|
||||
|
||||
|
||||
def get_media_dir(media_dir: str, dataset: str):
|
||||
if dataset in ["saco_veval_sav_test", "saco_veval_sav_val"]:
|
||||
return os.path.join(media_dir, "saco_sav", "JPEGImages_24fps")
|
||||
elif dataset in ["saco_veval_yt1b_test", "saco_veval_yt1b_val"]:
|
||||
return os.path.join(media_dir, "saco_yt1b", "JPEGImages_6fps")
|
||||
elif dataset in ["saco_veval_smartglasses_test", "saco_veval_smartglasses_val"]:
|
||||
return os.path.join(media_dir, "saco_sg", "JPEGImages_6fps")
|
||||
elif dataset == "sa_fari_test":
|
||||
return os.path.join(media_dir, "sa_fari", "JPEGImages_6fps")
|
||||
else:
|
||||
raise ValueError(f"Dataset {dataset} not found")
|
||||
|
||||
|
||||
def get_all_annotations_for_frame(
|
||||
dataset_df: pd.DataFrame, video_id: int, frame_idx: int, data_dir: str, dataset: str
|
||||
):
|
||||
media_dir = os.path.join(data_dir, "media")
|
||||
|
||||
# Load the annotation and video data
|
||||
annot_df = dataset_df["annotations"]
|
||||
video_df = dataset_df["videos"]
|
||||
|
||||
# Get the frame
|
||||
video_df_current = video_df[video_df.id == video_id]
|
||||
assert (
|
||||
len(video_df_current) == 1
|
||||
), f"Expected 1 video row, got {len(video_df_current)}"
|
||||
video_row = video_df_current.iloc[0]
|
||||
file_name = video_row.file_names[frame_idx]
|
||||
file_path = os.path.join(
|
||||
get_media_dir(media_dir=media_dir, dataset=dataset), file_name
|
||||
)
|
||||
frame = cv2.imread(file_path)
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# Get the masks and noun phrases annotated in this video in this frame
|
||||
annot_df_current_video = annot_df[annot_df.video_id == video_id]
|
||||
if len(annot_df_current_video) == 0:
|
||||
print(f"No annotations found for video_id {video_id}")
|
||||
return frame, None, None
|
||||
else:
|
||||
empty_mask = np.zeros(frame.shape[:2], dtype=np.uint8)
|
||||
mask_np_pairs = annot_df_current_video.apply(
|
||||
lambda row: (
|
||||
(
|
||||
mask_utils.decode(row.segmentations[frame_idx])
|
||||
if row.segmentations[frame_idx]
|
||||
else empty_mask
|
||||
),
|
||||
row.noun_phrase,
|
||||
),
|
||||
axis=1,
|
||||
)
|
||||
# sort based on noun_phrases
|
||||
mask_np_pairs = sorted(mask_np_pairs, key=lambda x: x[1])
|
||||
masks, noun_phrases = zip(*mask_np_pairs)
|
||||
|
||||
return frame, masks, noun_phrases
|
||||
|
||||
|
||||
def visualize_prompt_overlay(
|
||||
frame_idx,
|
||||
video_frames,
|
||||
title="Prompt Visualization",
|
||||
text_prompt=None,
|
||||
point_prompts=None,
|
||||
point_labels=None,
|
||||
bounding_boxes=None,
|
||||
box_labels=None,
|
||||
obj_id=None,
|
||||
):
|
||||
"""Simple prompt visualization function"""
|
||||
img = Image.fromarray(load_frame(video_frames[frame_idx]))
|
||||
fig, ax = plt.subplots(1, figsize=(6, 4))
|
||||
ax.imshow(img)
|
||||
|
||||
img_w, img_h = img.size
|
||||
|
||||
if text_prompt:
|
||||
ax.text(
|
||||
0.02,
|
||||
0.98,
|
||||
f'Text: "{text_prompt}"',
|
||||
transform=ax.transAxes,
|
||||
fontsize=12,
|
||||
color="white",
|
||||
weight="bold",
|
||||
bbox=dict(boxstyle="round,pad=0.3", facecolor="red", alpha=0.7),
|
||||
verticalalignment="top",
|
||||
)
|
||||
|
||||
if point_prompts:
|
||||
for i, point in enumerate(point_prompts):
|
||||
x, y = point
|
||||
# Convert relative to absolute coordinates
|
||||
x_img, y_img = x * img_w, y * img_h
|
||||
|
||||
# Use different colors for positive/negative points
|
||||
if point_labels and len(point_labels) > i:
|
||||
color = "green" if point_labels[i] == 1 else "red"
|
||||
marker = "o" if point_labels[i] == 1 else "x"
|
||||
else:
|
||||
color = "green"
|
||||
marker = "o"
|
||||
|
||||
ax.plot(
|
||||
x_img,
|
||||
y_img,
|
||||
marker=marker,
|
||||
color=color,
|
||||
markersize=10,
|
||||
markeredgewidth=2,
|
||||
markeredgecolor="white",
|
||||
)
|
||||
ax.text(
|
||||
x_img + 5,
|
||||
y_img - 5,
|
||||
f"P{i+1}",
|
||||
color=color,
|
||||
fontsize=10,
|
||||
weight="bold",
|
||||
bbox=dict(boxstyle="round,pad=0.2", facecolor="white", alpha=0.8),
|
||||
)
|
||||
|
||||
if bounding_boxes:
|
||||
for i, box in enumerate(bounding_boxes):
|
||||
x, y, w, h = box
|
||||
# Convert relative to absolute coordinates
|
||||
x_img, y_img = x * img_w, y * img_h
|
||||
w_img, h_img = w * img_w, h * img_h
|
||||
|
||||
# Use different colors for positive/negative boxes
|
||||
if box_labels and len(box_labels) > i:
|
||||
color = "green" if box_labels[i] == 1 else "red"
|
||||
else:
|
||||
color = "green"
|
||||
|
||||
rect = patches.Rectangle(
|
||||
(x_img, y_img),
|
||||
w_img,
|
||||
h_img,
|
||||
linewidth=2,
|
||||
edgecolor=color,
|
||||
facecolor="none",
|
||||
)
|
||||
ax.add_patch(rect)
|
||||
ax.text(
|
||||
x_img,
|
||||
y_img - 5,
|
||||
f"B{i+1}",
|
||||
color=color,
|
||||
fontsize=10,
|
||||
weight="bold",
|
||||
bbox=dict(boxstyle="round,pad=0.2", facecolor="white", alpha=0.8),
|
||||
)
|
||||
|
||||
# Add object ID info if provided
|
||||
if obj_id is not None:
|
||||
ax.text(
|
||||
0.02,
|
||||
0.02,
|
||||
f"Object ID: {obj_id}",
|
||||
transform=ax.transAxes,
|
||||
fontsize=10,
|
||||
color="white",
|
||||
weight="bold",
|
||||
bbox=dict(boxstyle="round,pad=0.3", facecolor="blue", alpha=0.7),
|
||||
verticalalignment="bottom",
|
||||
)
|
||||
|
||||
ax.set_title(title)
|
||||
ax.axis("off")
|
||||
plt.tight_layout()
|
||||
plt.show()
|
||||
|
||||
|
||||
def plot_results(img, results):
|
||||
plt.figure(figsize=(12, 8))
|
||||
plt.imshow(img)
|
||||
nb_objects = len(results["scores"])
|
||||
print(f"found {nb_objects} object(s)")
|
||||
for i in range(nb_objects):
|
||||
color = COLORS[i % len(COLORS)]
|
||||
plot_mask(results["masks"][i].squeeze(0).cpu(), color=color)
|
||||
w, h = img.size
|
||||
prob = results["scores"][i].item()
|
||||
plot_bbox(
|
||||
h,
|
||||
w,
|
||||
results["boxes"][i].cpu(),
|
||||
text=f"(id={i}, {prob=:.2f})",
|
||||
box_format="XYXY",
|
||||
color=color,
|
||||
relative_coords=False,
|
||||
)
|
||||
|
||||
|
||||
def single_visualization(img, anns, title):
|
||||
"""
|
||||
Create a single image visualization with overlays.
|
||||
"""
|
||||
fig, ax = plt.subplots(figsize=(7, 7))
|
||||
fig.suptitle(title, fontsize=16, fontweight="bold")
|
||||
overlay = render_masklet_frame(img, anns, alpha=0.5)
|
||||
ax.imshow(overlay)
|
||||
ax.axis("off")
|
||||
plt.tight_layout()
|
||||
|
||||
|
||||
def show_mask(mask, ax, obj_id=None, random_color=False):
|
||||
if random_color:
|
||||
color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
|
||||
else:
|
||||
cmap = plt.get_cmap("tab10")
|
||||
cmap_idx = 0 if obj_id is None else obj_id
|
||||
color = np.array([*cmap(cmap_idx)[:3], 0.6])
|
||||
h, w = mask.shape[-2:]
|
||||
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
|
||||
ax.imshow(mask_image)
|
||||
|
||||
|
||||
def show_box(box, ax):
|
||||
x0, y0 = box[0], box[1]
|
||||
w, h = box[2] - box[0], box[3] - box[1]
|
||||
ax.add_patch(
|
||||
plt.Rectangle((x0, y0), w, h, edgecolor="green", facecolor=(0, 0, 0, 0), lw=2)
|
||||
)
|
||||
|
||||
|
||||
def show_points(coords, labels, ax, marker_size=375):
|
||||
pos_points = coords[labels == 1]
|
||||
neg_points = coords[labels == 0]
|
||||
ax.scatter(
|
||||
pos_points[:, 0],
|
||||
pos_points[:, 1],
|
||||
color="green",
|
||||
marker="*",
|
||||
s=marker_size,
|
||||
edgecolor="white",
|
||||
linewidth=1.25,
|
||||
)
|
||||
ax.scatter(
|
||||
neg_points[:, 0],
|
||||
neg_points[:, 1],
|
||||
color="red",
|
||||
marker="*",
|
||||
s=marker_size,
|
||||
edgecolor="white",
|
||||
linewidth=1.25,
|
||||
)
|
||||
|
||||
|
||||
def load_frame(frame):
|
||||
if isinstance(frame, np.ndarray):
|
||||
img = frame
|
||||
elif isinstance(frame, Image.Image):
|
||||
img = np.array(frame)
|
||||
elif isinstance(frame, str) and os.path.isfile(frame):
|
||||
img = plt.imread(frame)
|
||||
else:
|
||||
raise ValueError(f"Invalid video frame type: {type(frame)=}")
|
||||
return img
|
||||
Reference in New Issue
Block a user