From a3dfbb31784b94c774f339a4a6ab02f1d89fafd7 Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Mon, 24 Nov 2025 20:56:51 -0800 Subject: [PATCH] Add files via upload --- sam3/__init__.py | 7 + sam3/logger.py | 54 +++ sam3/model_builder.py | 793 ++++++++++++++++++++++++++++++ sam3/visualization_utils.py | 941 ++++++++++++++++++++++++++++++++++++ 4 files changed, 1795 insertions(+) create mode 100644 sam3/__init__.py create mode 100644 sam3/logger.py create mode 100644 sam3/model_builder.py create mode 100644 sam3/visualization_utils.py diff --git a/sam3/__init__.py b/sam3/__init__.py new file mode 100644 index 0000000..14270a6 --- /dev/null +++ b/sam3/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from .model_builder import build_sam3_image_model + +__version__ = "0.1.0" + +__all__ = ["build_sam3_image_model"] diff --git a/sam3/logger.py b/sam3/logger.py new file mode 100644 index 0000000..db9c0a6 --- /dev/null +++ b/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/sam3/model_builder.py b/sam3/model_builder.py new file mode 100644 index 0000000..058bbec --- /dev/null +++ b/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/sam3/visualization_utils.py b/sam3/visualization_utils.py new file mode 100644 index 0000000..090f086 --- /dev/null +++ b/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