diff --git a/models/sam3/assets/bpe_simple_vocab_16e6.txt.gz b/models/sam3/assets/bpe_simple_vocab_16e6.txt.gz new file mode 100644 index 0000000..7b5088a Binary files /dev/null and b/models/sam3/assets/bpe_simple_vocab_16e6.txt.gz differ diff --git a/models/sam3/logger.py b/models/sam3/logger.py new file mode 100644 index 0000000..db9c0a6 --- /dev/null +++ b/models/sam3/logger.py @@ -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 diff --git a/models/sam3/model_builder.py b/models/sam3/model_builder.py new file mode 100644 index 0000000..058bbec --- /dev/null +++ b/models/sam3/model_builder.py @@ -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 + ) diff --git a/models/sam3/sam/__init__.py b/models/sam3/sam/__init__.py new file mode 100644 index 0000000..8b35da6 --- /dev/null +++ b/models/sam3/sam/__init__.py @@ -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 diff --git a/models/sam3/sam/common.py b/models/sam3/sam/common.py new file mode 100644 index 0000000..b6d1587 --- /dev/null +++ b/models/sam3/sam/common.py @@ -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 diff --git a/models/sam3/sam/mask_decoder.py b/models/sam3/sam/mask_decoder.py new file mode 100644 index 0000000..b4ac397 --- /dev/null +++ b/models/sam3/sam/mask_decoder.py @@ -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 diff --git a/models/sam3/sam/prompt_encoder.py b/models/sam3/sam/prompt_encoder.py new file mode 100644 index 0000000..145ea9f --- /dev/null +++ b/models/sam3/sam/prompt_encoder.py @@ -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 diff --git a/models/sam3/sam/rope.py b/models/sam3/sam/rope.py new file mode 100644 index 0000000..2db01b6 --- /dev/null +++ b/models/sam3/sam/rope.py @@ -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 diff --git a/models/sam3/sam/transformer.py b/models/sam3/sam/transformer.py new file mode 100644 index 0000000..dc5bf2d --- /dev/null +++ b/models/sam3/sam/transformer.py @@ -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 diff --git a/models/sam3/sam/transformer.py.bak b/models/sam3/sam/transformer.py.bak new file mode 100644 index 0000000..3e96c28 --- /dev/null +++ b/models/sam3/sam/transformer.py.bak @@ -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 diff --git a/models/sam3/visualization_utils.py b/models/sam3/visualization_utils.py new file mode 100644 index 0000000..090f086 --- /dev/null +++ b/models/sam3/visualization_utils.py @@ -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