Add files via upload
This commit is contained in:
AI Lab
@@ -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"]
|
||||
@@ -0,0 +1,54 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
import logging
|
||||
import os
|
||||
|
||||
LOG_LEVELS = {
|
||||
"DEBUG": logging.DEBUG,
|
||||
"INFO": logging.INFO,
|
||||
"WARNING": logging.WARNING,
|
||||
"ERROR": logging.ERROR,
|
||||
"CRITICAL": logging.CRITICAL,
|
||||
}
|
||||
|
||||
|
||||
class ColoredFormatter(logging.Formatter):
|
||||
"""A command line formatter with different colors for each level."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
reset = "\033[0m"
|
||||
colors = {
|
||||
logging.DEBUG: f"{reset}\033[36m", # cyan,
|
||||
logging.INFO: f"{reset}\033[32m", # green
|
||||
logging.WARNING: f"{reset}\033[33m", # yellow
|
||||
logging.ERROR: f"{reset}\033[31m", # red
|
||||
logging.CRITICAL: f"{reset}\033[35m", # magenta
|
||||
}
|
||||
fmt_str = "{color}%(levelname)s %(asctime)s %(process)d %(filename)s:%(lineno)4d:{reset} %(message)s"
|
||||
self.formatters = {
|
||||
level: logging.Formatter(fmt_str.format(color=color, reset=reset))
|
||||
for level, color in colors.items()
|
||||
}
|
||||
self.default_formatter = self.formatters[logging.INFO]
|
||||
|
||||
def format(self, record):
|
||||
formatter = self.formatters.get(record.levelno, self.default_formatter)
|
||||
return formatter.format(record)
|
||||
|
||||
|
||||
def get_logger(name, level=logging.INFO):
|
||||
"""A command line logger."""
|
||||
if "LOG_LEVEL" in os.environ:
|
||||
level = os.environ["LOG_LEVEL"].upper()
|
||||
assert (
|
||||
level in LOG_LEVELS
|
||||
), f"Invalid LOG_LEVEL: {level}, must be one of {list(LOG_LEVELS.keys())}"
|
||||
level = LOG_LEVELS[level]
|
||||
logger = logging.getLogger(name)
|
||||
logger.setLevel(level)
|
||||
logger.propagate = False
|
||||
ch = logging.StreamHandler()
|
||||
ch.setLevel(level)
|
||||
ch.setFormatter(ColoredFormatter())
|
||||
logger.addHandler(ch)
|
||||
return logger
|
||||
@@ -0,0 +1,793 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from huggingface_hub import hf_hub_download
|
||||
from iopath.common.file_io import g_pathmgr
|
||||
from sam3.model.decoder import (
|
||||
TransformerDecoder,
|
||||
TransformerDecoderLayer,
|
||||
TransformerDecoderLayerv2,
|
||||
TransformerEncoderCrossAttention,
|
||||
)
|
||||
from sam3.model.encoder import TransformerEncoderFusion, TransformerEncoderLayer
|
||||
from sam3.model.geometry_encoders import SequenceGeometryEncoder
|
||||
from sam3.model.maskformer_segmentation import PixelDecoder, UniversalSegmentationHead
|
||||
from sam3.model.memory import (
|
||||
CXBlock,
|
||||
SimpleFuser,
|
||||
SimpleMaskDownSampler,
|
||||
SimpleMaskEncoder,
|
||||
)
|
||||
from sam3.model.model_misc import (
|
||||
DotProductScoring,
|
||||
MLP,
|
||||
MultiheadAttentionWrapper as MultiheadAttention,
|
||||
TransformerWrapper,
|
||||
)
|
||||
from sam3.model.necks import Sam3DualViTDetNeck
|
||||
from sam3.model.position_encoding import PositionEmbeddingSine
|
||||
from sam3.model.sam1_task_predictor import SAM3InteractiveImagePredictor
|
||||
from sam3.model.sam3_image import Sam3Image, Sam3ImageOnVideoMultiGPU
|
||||
from sam3.model.sam3_tracking_predictor import Sam3TrackerPredictor
|
||||
from sam3.model.sam3_video_inference import Sam3VideoInferenceWithInstanceInteractivity
|
||||
from sam3.model.sam3_video_predictor import Sam3VideoPredictorMultiGPU
|
||||
from sam3.model.text_encoder_ve import VETextEncoder
|
||||
from sam3.model.tokenizer_ve import SimpleTokenizer
|
||||
from sam3.model.vitdet import ViT
|
||||
from sam3.model.vl_combiner import SAM3VLBackbone
|
||||
from sam3.sam.transformer import RoPEAttention
|
||||
|
||||
|
||||
# Setup TensorFloat-32 for Ampere GPUs if available
|
||||
def _setup_tf32() -> None:
|
||||
"""Enable TensorFloat-32 for Ampere GPUs if available."""
|
||||
if torch.cuda.is_available():
|
||||
device_props = torch.cuda.get_device_properties(0)
|
||||
if device_props.major >= 8:
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
|
||||
|
||||
_setup_tf32()
|
||||
|
||||
|
||||
def _create_position_encoding(precompute_resolution=None):
|
||||
"""Create position encoding for visual backbone."""
|
||||
return PositionEmbeddingSine(
|
||||
num_pos_feats=256,
|
||||
normalize=True,
|
||||
scale=None,
|
||||
temperature=10000,
|
||||
precompute_resolution=precompute_resolution,
|
||||
)
|
||||
|
||||
|
||||
def _create_vit_backbone(compile_mode=None):
|
||||
"""Create ViT backbone for visual feature extraction."""
|
||||
return ViT(
|
||||
img_size=1008,
|
||||
pretrain_img_size=336,
|
||||
patch_size=14,
|
||||
embed_dim=1024,
|
||||
depth=32,
|
||||
num_heads=16,
|
||||
mlp_ratio=4.625,
|
||||
norm_layer="LayerNorm",
|
||||
drop_path_rate=0.1,
|
||||
qkv_bias=True,
|
||||
use_abs_pos=True,
|
||||
tile_abs_pos=True,
|
||||
global_att_blocks=(7, 15, 23, 31),
|
||||
rel_pos_blocks=(),
|
||||
use_rope=True,
|
||||
use_interp_rope=True,
|
||||
window_size=24,
|
||||
pretrain_use_cls_token=True,
|
||||
retain_cls_token=False,
|
||||
ln_pre=True,
|
||||
ln_post=False,
|
||||
return_interm_layers=False,
|
||||
bias_patch_embed=False,
|
||||
compile_mode=compile_mode,
|
||||
)
|
||||
|
||||
|
||||
def _create_vit_neck(position_encoding, vit_backbone, enable_inst_interactivity=False):
|
||||
"""Create ViT neck for feature pyramid."""
|
||||
return Sam3DualViTDetNeck(
|
||||
position_encoding=position_encoding,
|
||||
d_model=256,
|
||||
scale_factors=[4.0, 2.0, 1.0, 0.5],
|
||||
trunk=vit_backbone,
|
||||
add_sam2_neck=enable_inst_interactivity,
|
||||
)
|
||||
|
||||
|
||||
def _create_vl_backbone(vit_neck, text_encoder):
|
||||
"""Create visual-language backbone."""
|
||||
return SAM3VLBackbone(visual=vit_neck, text=text_encoder, scalp=1)
|
||||
|
||||
|
||||
def _create_transformer_encoder() -> TransformerEncoderFusion:
|
||||
"""Create transformer encoder with its layer."""
|
||||
encoder_layer = TransformerEncoderLayer(
|
||||
activation="relu",
|
||||
d_model=256,
|
||||
dim_feedforward=2048,
|
||||
dropout=0.1,
|
||||
pos_enc_at_attn=True,
|
||||
pos_enc_at_cross_attn_keys=False,
|
||||
pos_enc_at_cross_attn_queries=False,
|
||||
pre_norm=True,
|
||||
self_attention=MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0.1,
|
||||
embed_dim=256,
|
||||
batch_first=True,
|
||||
),
|
||||
cross_attention=MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0.1,
|
||||
embed_dim=256,
|
||||
batch_first=True,
|
||||
),
|
||||
)
|
||||
|
||||
encoder = TransformerEncoderFusion(
|
||||
layer=encoder_layer,
|
||||
num_layers=6,
|
||||
d_model=256,
|
||||
num_feature_levels=1,
|
||||
frozen=False,
|
||||
use_act_checkpoint=True,
|
||||
add_pooled_text_to_img_feat=False,
|
||||
pool_text_with_mask=True,
|
||||
)
|
||||
return encoder
|
||||
|
||||
|
||||
def _create_transformer_decoder() -> TransformerDecoder:
|
||||
"""Create transformer decoder with its layer."""
|
||||
decoder_layer = TransformerDecoderLayer(
|
||||
activation="relu",
|
||||
d_model=256,
|
||||
dim_feedforward=2048,
|
||||
dropout=0.1,
|
||||
cross_attention=MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0.1,
|
||||
embed_dim=256,
|
||||
),
|
||||
n_heads=8,
|
||||
use_text_cross_attention=True,
|
||||
)
|
||||
|
||||
decoder = TransformerDecoder(
|
||||
layer=decoder_layer,
|
||||
num_layers=6,
|
||||
num_queries=200,
|
||||
return_intermediate=True,
|
||||
box_refine=True,
|
||||
num_o2m_queries=0,
|
||||
dac=True,
|
||||
boxRPB="log",
|
||||
d_model=256,
|
||||
frozen=False,
|
||||
interaction_layer=None,
|
||||
dac_use_selfatt_ln=True,
|
||||
resolution=1008,
|
||||
stride=14,
|
||||
use_act_checkpoint=True,
|
||||
presence_token=True,
|
||||
)
|
||||
return decoder
|
||||
|
||||
|
||||
def _create_dot_product_scoring():
|
||||
"""Create dot product scoring module."""
|
||||
prompt_mlp = MLP(
|
||||
input_dim=256,
|
||||
hidden_dim=2048,
|
||||
output_dim=256,
|
||||
num_layers=2,
|
||||
dropout=0.1,
|
||||
residual=True,
|
||||
out_norm=nn.LayerNorm(256),
|
||||
)
|
||||
return DotProductScoring(d_model=256, d_proj=256, prompt_mlp=prompt_mlp)
|
||||
|
||||
|
||||
def _create_segmentation_head(compile_mode=None):
|
||||
"""Create segmentation head with pixel decoder."""
|
||||
pixel_decoder = PixelDecoder(
|
||||
num_upsampling_stages=3,
|
||||
interpolation_mode="nearest",
|
||||
hidden_dim=256,
|
||||
compile_mode=compile_mode,
|
||||
)
|
||||
|
||||
cross_attend_prompt = MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0,
|
||||
embed_dim=256,
|
||||
)
|
||||
|
||||
segmentation_head = UniversalSegmentationHead(
|
||||
hidden_dim=256,
|
||||
upsampling_stages=3,
|
||||
aux_masks=False,
|
||||
presence_head=False,
|
||||
dot_product_scorer=None,
|
||||
act_ckpt=True,
|
||||
cross_attend_prompt=cross_attend_prompt,
|
||||
pixel_decoder=pixel_decoder,
|
||||
)
|
||||
return segmentation_head
|
||||
|
||||
|
||||
def _create_geometry_encoder():
|
||||
"""Create geometry encoder with all its components."""
|
||||
# Create position encoding for geometry encoder
|
||||
geo_pos_enc = _create_position_encoding()
|
||||
# Create CX block for fuser
|
||||
cx_block = CXBlock(
|
||||
dim=256,
|
||||
kernel_size=7,
|
||||
padding=3,
|
||||
layer_scale_init_value=1.0e-06,
|
||||
use_dwconv=True,
|
||||
)
|
||||
# Create geometry encoder layer
|
||||
geo_layer = TransformerEncoderLayer(
|
||||
activation="relu",
|
||||
d_model=256,
|
||||
dim_feedforward=2048,
|
||||
dropout=0.1,
|
||||
pos_enc_at_attn=False,
|
||||
pre_norm=True,
|
||||
self_attention=MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0.1,
|
||||
embed_dim=256,
|
||||
batch_first=False,
|
||||
),
|
||||
pos_enc_at_cross_attn_queries=False,
|
||||
pos_enc_at_cross_attn_keys=True,
|
||||
cross_attention=MultiheadAttention(
|
||||
num_heads=8,
|
||||
dropout=0.1,
|
||||
embed_dim=256,
|
||||
batch_first=False,
|
||||
),
|
||||
)
|
||||
|
||||
# Create geometry encoder
|
||||
input_geometry_encoder = SequenceGeometryEncoder(
|
||||
pos_enc=geo_pos_enc,
|
||||
encode_boxes_as_points=False,
|
||||
points_direct_project=True,
|
||||
points_pool=True,
|
||||
points_pos_enc=True,
|
||||
boxes_direct_project=True,
|
||||
boxes_pool=True,
|
||||
boxes_pos_enc=True,
|
||||
d_model=256,
|
||||
num_layers=3,
|
||||
layer=geo_layer,
|
||||
use_act_ckpt=True,
|
||||
add_cls=True,
|
||||
add_post_encode_proj=True,
|
||||
)
|
||||
return input_geometry_encoder
|
||||
|
||||
|
||||
def _create_sam3_model(
|
||||
backbone,
|
||||
transformer,
|
||||
input_geometry_encoder,
|
||||
segmentation_head,
|
||||
dot_prod_scoring,
|
||||
inst_interactive_predictor,
|
||||
eval_mode,
|
||||
):
|
||||
"""Create the SAM3 image model."""
|
||||
common_params = {
|
||||
"backbone": backbone,
|
||||
"transformer": transformer,
|
||||
"input_geometry_encoder": input_geometry_encoder,
|
||||
"segmentation_head": segmentation_head,
|
||||
"num_feature_levels": 1,
|
||||
"o2m_mask_predict": True,
|
||||
"dot_prod_scoring": dot_prod_scoring,
|
||||
"use_instance_query": False,
|
||||
"multimask_output": True,
|
||||
"inst_interactive_predictor": inst_interactive_predictor,
|
||||
}
|
||||
|
||||
matcher = None
|
||||
if not eval_mode:
|
||||
from sam3.train.matcher import BinaryHungarianMatcherV2
|
||||
|
||||
matcher = BinaryHungarianMatcherV2(
|
||||
focal=True,
|
||||
cost_class=2.0,
|
||||
cost_bbox=5.0,
|
||||
cost_giou=2.0,
|
||||
alpha=0.25,
|
||||
gamma=2,
|
||||
stable=False,
|
||||
)
|
||||
common_params["matcher"] = matcher
|
||||
model = Sam3Image(**common_params)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def _create_tracker_maskmem_backbone():
|
||||
"""Create the SAM3 Tracker memory encoder."""
|
||||
# Position encoding for mask memory backbone
|
||||
position_encoding = PositionEmbeddingSine(
|
||||
num_pos_feats=64,
|
||||
normalize=True,
|
||||
scale=None,
|
||||
temperature=10000,
|
||||
precompute_resolution=1008,
|
||||
)
|
||||
|
||||
# Mask processing components
|
||||
mask_downsampler = SimpleMaskDownSampler(
|
||||
kernel_size=3, stride=2, padding=1, interpol_size=[1152, 1152]
|
||||
)
|
||||
|
||||
cx_block_layer = CXBlock(
|
||||
dim=256,
|
||||
kernel_size=7,
|
||||
padding=3,
|
||||
layer_scale_init_value=1.0e-06,
|
||||
use_dwconv=True,
|
||||
)
|
||||
|
||||
fuser = SimpleFuser(layer=cx_block_layer, num_layers=2)
|
||||
|
||||
maskmem_backbone = SimpleMaskEncoder(
|
||||
out_dim=64,
|
||||
position_encoding=position_encoding,
|
||||
mask_downsampler=mask_downsampler,
|
||||
fuser=fuser,
|
||||
)
|
||||
|
||||
return maskmem_backbone
|
||||
|
||||
|
||||
def _create_tracker_transformer():
|
||||
"""Create the SAM3 Tracker transformer components."""
|
||||
# Self attention
|
||||
self_attention = RoPEAttention(
|
||||
embedding_dim=256,
|
||||
num_heads=1,
|
||||
downsample_rate=1,
|
||||
dropout=0.1,
|
||||
rope_theta=10000.0,
|
||||
feat_sizes=[72, 72],
|
||||
use_fa3=False,
|
||||
use_rope_real=False,
|
||||
)
|
||||
|
||||
# Cross attention
|
||||
cross_attention = RoPEAttention(
|
||||
embedding_dim=256,
|
||||
num_heads=1,
|
||||
downsample_rate=1,
|
||||
dropout=0.1,
|
||||
kv_in_dim=64,
|
||||
rope_theta=10000.0,
|
||||
feat_sizes=[72, 72],
|
||||
rope_k_repeat=True,
|
||||
use_fa3=False,
|
||||
use_rope_real=False,
|
||||
)
|
||||
|
||||
# Encoder layer
|
||||
encoder_layer = TransformerDecoderLayerv2(
|
||||
cross_attention_first=False,
|
||||
activation="relu",
|
||||
dim_feedforward=2048,
|
||||
dropout=0.1,
|
||||
pos_enc_at_attn=False,
|
||||
pre_norm=True,
|
||||
self_attention=self_attention,
|
||||
d_model=256,
|
||||
pos_enc_at_cross_attn_keys=True,
|
||||
pos_enc_at_cross_attn_queries=False,
|
||||
cross_attention=cross_attention,
|
||||
)
|
||||
|
||||
# Encoder
|
||||
encoder = TransformerEncoderCrossAttention(
|
||||
remove_cross_attention_layers=[],
|
||||
batch_first=True,
|
||||
d_model=256,
|
||||
frozen=False,
|
||||
pos_enc_at_input=True,
|
||||
layer=encoder_layer,
|
||||
num_layers=4,
|
||||
use_act_checkpoint=False,
|
||||
)
|
||||
|
||||
# Transformer wrapper
|
||||
transformer = TransformerWrapper(
|
||||
encoder=encoder,
|
||||
decoder=None,
|
||||
d_model=256,
|
||||
)
|
||||
|
||||
return transformer
|
||||
|
||||
|
||||
def build_tracker(
|
||||
apply_temporal_disambiguation: bool, with_backbone: bool = False, compile_mode=None
|
||||
) -> Sam3TrackerPredictor:
|
||||
"""
|
||||
Build the SAM3 Tracker module for video tracking.
|
||||
|
||||
Returns:
|
||||
Sam3TrackerPredictor: Wrapped SAM3 Tracker module
|
||||
"""
|
||||
|
||||
# Create model components
|
||||
maskmem_backbone = _create_tracker_maskmem_backbone()
|
||||
transformer = _create_tracker_transformer()
|
||||
backbone = None
|
||||
if with_backbone:
|
||||
vision_backbone = _create_vision_backbone(compile_mode=compile_mode)
|
||||
backbone = SAM3VLBackbone(scalp=1, visual=vision_backbone, text=None)
|
||||
# Create the Tracker module
|
||||
model = Sam3TrackerPredictor(
|
||||
image_size=1008,
|
||||
num_maskmem=7,
|
||||
backbone=backbone,
|
||||
backbone_stride=14,
|
||||
transformer=transformer,
|
||||
maskmem_backbone=maskmem_backbone,
|
||||
# SAM parameters
|
||||
multimask_output_in_sam=True,
|
||||
# Evaluation
|
||||
forward_backbone_per_frame_for_eval=True,
|
||||
trim_past_non_cond_mem_for_eval=False,
|
||||
# Multimask
|
||||
multimask_output_for_tracking=True,
|
||||
multimask_min_pt_num=0,
|
||||
multimask_max_pt_num=1,
|
||||
# Additional settings
|
||||
always_start_from_first_ann_frame=False,
|
||||
# Mask overlap
|
||||
non_overlap_masks_for_mem_enc=False,
|
||||
non_overlap_masks_for_output=False,
|
||||
max_cond_frames_in_attn=4,
|
||||
offload_output_to_cpu_for_eval=False,
|
||||
# SAM decoder settings
|
||||
sam_mask_decoder_extra_args={
|
||||
"dynamic_multimask_via_stability": True,
|
||||
"dynamic_multimask_stability_delta": 0.05,
|
||||
"dynamic_multimask_stability_thresh": 0.98,
|
||||
},
|
||||
clear_non_cond_mem_around_input=True,
|
||||
fill_hole_area=0,
|
||||
use_memory_selection=apply_temporal_disambiguation,
|
||||
)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def _create_text_encoder(bpe_path: str) -> VETextEncoder:
|
||||
"""Create SAM3 text encoder."""
|
||||
tokenizer = SimpleTokenizer(bpe_path=bpe_path)
|
||||
return VETextEncoder(
|
||||
tokenizer=tokenizer,
|
||||
d_model=256,
|
||||
width=1024,
|
||||
heads=16,
|
||||
layers=24,
|
||||
)
|
||||
|
||||
|
||||
def _create_vision_backbone(
|
||||
compile_mode=None, enable_inst_interactivity=True
|
||||
) -> Sam3DualViTDetNeck:
|
||||
"""Create SAM3 visual backbone with ViT and neck."""
|
||||
# Position encoding
|
||||
position_encoding = _create_position_encoding(precompute_resolution=1008)
|
||||
# ViT backbone
|
||||
vit_backbone: ViT = _create_vit_backbone(compile_mode=compile_mode)
|
||||
vit_neck: Sam3DualViTDetNeck = _create_vit_neck(
|
||||
position_encoding,
|
||||
vit_backbone,
|
||||
enable_inst_interactivity=enable_inst_interactivity,
|
||||
)
|
||||
# Visual neck
|
||||
return vit_neck
|
||||
|
||||
|
||||
def _create_sam3_transformer(has_presence_token: bool = True) -> TransformerWrapper:
|
||||
"""Create SAM3 transformer encoder and decoder."""
|
||||
encoder: TransformerEncoderFusion = _create_transformer_encoder()
|
||||
decoder: TransformerDecoder = _create_transformer_decoder()
|
||||
|
||||
return TransformerWrapper(encoder=encoder, decoder=decoder, d_model=256)
|
||||
|
||||
|
||||
def _load_checkpoint(model, checkpoint_path):
|
||||
"""Load model checkpoint from file."""
|
||||
with g_pathmgr.open(checkpoint_path, "rb") as f:
|
||||
ckpt = torch.load(f, map_location="cpu", weights_only=True)
|
||||
if "model" in ckpt and isinstance(ckpt["model"], dict):
|
||||
ckpt = ckpt["model"]
|
||||
sam3_image_ckpt = {
|
||||
k.replace("detector.", ""): v for k, v in ckpt.items() if "detector" in k
|
||||
}
|
||||
if model.inst_interactive_predictor is not None:
|
||||
sam3_image_ckpt.update(
|
||||
{
|
||||
k.replace("tracker.", "inst_interactive_predictor.model."): v
|
||||
for k, v in ckpt.items()
|
||||
if "tracker" in k
|
||||
}
|
||||
)
|
||||
missing_keys, _ = model.load_state_dict(sam3_image_ckpt, strict=False)
|
||||
if len(missing_keys) > 0:
|
||||
print(
|
||||
f"loaded {checkpoint_path} and found "
|
||||
f"missing and/or unexpected keys:\n{missing_keys=}"
|
||||
)
|
||||
|
||||
|
||||
def _setup_device_and_mode(model, device, eval_mode):
|
||||
"""Setup model device and evaluation mode."""
|
||||
if device == "cuda":
|
||||
model = model.cuda()
|
||||
if eval_mode:
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
|
||||
def build_sam3_image_model(
|
||||
bpe_path=None,
|
||||
device="cuda" if torch.cuda.is_available() else "cpu",
|
||||
eval_mode=True,
|
||||
checkpoint_path=None,
|
||||
load_from_HF=True,
|
||||
enable_segmentation=True,
|
||||
enable_inst_interactivity=False,
|
||||
compile=False,
|
||||
):
|
||||
"""
|
||||
Build SAM3 image model
|
||||
|
||||
Args:
|
||||
bpe_path: Path to the BPE tokenizer vocabulary
|
||||
device: Device to load the model on ('cuda' or 'cpu')
|
||||
eval_mode: Whether to set the model to evaluation mode
|
||||
checkpoint_path: Optional path to model checkpoint
|
||||
enable_segmentation: Whether to enable segmentation head
|
||||
enable_inst_interactivity: Whether to enable instance interactivity (SAM 1 task)
|
||||
compile_mode: To enable compilation, set to "default"
|
||||
|
||||
Returns:
|
||||
A SAM3 image model
|
||||
"""
|
||||
if bpe_path is None:
|
||||
bpe_path = os.path.join(
|
||||
os.path.dirname(__file__), "..", "assets", "bpe_simple_vocab_16e6.txt.gz"
|
||||
)
|
||||
# Create visual components
|
||||
compile_mode = "default" if compile else None
|
||||
vision_encoder = _create_vision_backbone(
|
||||
compile_mode=compile_mode, enable_inst_interactivity=enable_inst_interactivity
|
||||
)
|
||||
|
||||
# Create text components
|
||||
text_encoder = _create_text_encoder(bpe_path)
|
||||
|
||||
# Create visual-language backbone
|
||||
backbone = _create_vl_backbone(vision_encoder, text_encoder)
|
||||
|
||||
# Create transformer components
|
||||
transformer = _create_sam3_transformer()
|
||||
|
||||
# Create dot product scoring
|
||||
dot_prod_scoring = _create_dot_product_scoring()
|
||||
|
||||
# Create segmentation head if enabled
|
||||
segmentation_head = (
|
||||
_create_segmentation_head(compile_mode=compile_mode)
|
||||
if enable_segmentation
|
||||
else None
|
||||
)
|
||||
|
||||
# Create geometry encoder
|
||||
input_geometry_encoder = _create_geometry_encoder()
|
||||
if enable_inst_interactivity:
|
||||
sam3_pvs_base = build_tracker(apply_temporal_disambiguation=False)
|
||||
inst_predictor = SAM3InteractiveImagePredictor(sam3_pvs_base)
|
||||
else:
|
||||
inst_predictor = None
|
||||
# Create the SAM3 model
|
||||
model = _create_sam3_model(
|
||||
backbone,
|
||||
transformer,
|
||||
input_geometry_encoder,
|
||||
segmentation_head,
|
||||
dot_prod_scoring,
|
||||
inst_predictor,
|
||||
eval_mode,
|
||||
)
|
||||
if load_from_HF and checkpoint_path is None:
|
||||
checkpoint_path = download_ckpt_from_hf()
|
||||
# Load checkpoint if provided
|
||||
if checkpoint_path is not None:
|
||||
_load_checkpoint(model, checkpoint_path)
|
||||
|
||||
# Setup device and mode
|
||||
model = _setup_device_and_mode(model, device, eval_mode)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def download_ckpt_from_hf():
|
||||
SAM3_MODEL_ID = "facebook/sam3"
|
||||
SAM3_CKPT_NAME = "sam3.pt"
|
||||
SAM3_CFG_NAME = "config.json"
|
||||
_ = hf_hub_download(repo_id=SAM3_MODEL_ID, filename=SAM3_CFG_NAME)
|
||||
checkpoint_path = hf_hub_download(repo_id=SAM3_MODEL_ID, filename=SAM3_CKPT_NAME)
|
||||
return checkpoint_path
|
||||
|
||||
|
||||
def build_sam3_video_model(
|
||||
checkpoint_path: Optional[str] = None,
|
||||
load_from_HF=True,
|
||||
bpe_path: Optional[str] = None,
|
||||
has_presence_token: bool = True,
|
||||
geo_encoder_use_img_cross_attn: bool = True,
|
||||
strict_state_dict_loading: bool = True,
|
||||
apply_temporal_disambiguation: bool = True,
|
||||
device="cuda" if torch.cuda.is_available() else "cpu",
|
||||
compile=False,
|
||||
) -> Sam3VideoInferenceWithInstanceInteractivity:
|
||||
"""
|
||||
Build SAM3 dense tracking model.
|
||||
|
||||
Args:
|
||||
checkpoint_path: Optional path to checkpoint file
|
||||
bpe_path: Path to the BPE tokenizer file
|
||||
|
||||
Returns:
|
||||
Sam3VideoInferenceWithInstanceInteractivity: The instantiated dense tracking model
|
||||
"""
|
||||
if bpe_path is None:
|
||||
bpe_path = os.path.join(
|
||||
os.path.dirname(__file__), "..", "assets", "bpe_simple_vocab_16e6.txt.gz"
|
||||
)
|
||||
|
||||
# Build Tracker module
|
||||
tracker = build_tracker(apply_temporal_disambiguation=apply_temporal_disambiguation)
|
||||
|
||||
# Build Detector components
|
||||
visual_neck = _create_vision_backbone()
|
||||
text_encoder = _create_text_encoder(bpe_path)
|
||||
backbone = SAM3VLBackbone(scalp=1, visual=visual_neck, text=text_encoder)
|
||||
transformer = _create_sam3_transformer(has_presence_token=has_presence_token)
|
||||
segmentation_head: UniversalSegmentationHead = _create_segmentation_head()
|
||||
input_geometry_encoder = _create_geometry_encoder()
|
||||
|
||||
# Create main dot product scoring
|
||||
main_dot_prod_mlp = MLP(
|
||||
input_dim=256,
|
||||
hidden_dim=2048,
|
||||
output_dim=256,
|
||||
num_layers=2,
|
||||
dropout=0.1,
|
||||
residual=True,
|
||||
out_norm=nn.LayerNorm(256),
|
||||
)
|
||||
main_dot_prod_scoring = DotProductScoring(
|
||||
d_model=256, d_proj=256, prompt_mlp=main_dot_prod_mlp
|
||||
)
|
||||
|
||||
# Build Detector module
|
||||
detector = Sam3ImageOnVideoMultiGPU(
|
||||
num_feature_levels=1,
|
||||
backbone=backbone,
|
||||
transformer=transformer,
|
||||
segmentation_head=segmentation_head,
|
||||
semantic_segmentation_head=None,
|
||||
input_geometry_encoder=input_geometry_encoder,
|
||||
use_early_fusion=True,
|
||||
use_dot_prod_scoring=True,
|
||||
dot_prod_scoring=main_dot_prod_scoring,
|
||||
supervise_joint_box_scores=has_presence_token,
|
||||
)
|
||||
|
||||
# Build the main SAM3 video model
|
||||
if apply_temporal_disambiguation:
|
||||
model = Sam3VideoInferenceWithInstanceInteractivity(
|
||||
detector=detector,
|
||||
tracker=tracker,
|
||||
score_threshold_detection=0.5,
|
||||
assoc_iou_thresh=0.1,
|
||||
det_nms_thresh=0.1,
|
||||
new_det_thresh=0.7,
|
||||
hotstart_delay=15,
|
||||
hotstart_unmatch_thresh=8,
|
||||
hotstart_dup_thresh=8,
|
||||
suppress_unmatched_only_within_hotstart=True,
|
||||
min_trk_keep_alive=-1,
|
||||
max_trk_keep_alive=30,
|
||||
init_trk_keep_alive=30,
|
||||
suppress_overlapping_based_on_recent_occlusion_threshold=0.7,
|
||||
suppress_det_close_to_boundary=False,
|
||||
fill_hole_area=16,
|
||||
recondition_every_nth_frame=16,
|
||||
masklet_confirmation_enable=False,
|
||||
decrease_trk_keep_alive_for_empty_masklets=False,
|
||||
image_size=1008,
|
||||
image_mean=(0.5, 0.5, 0.5),
|
||||
image_std=(0.5, 0.5, 0.5),
|
||||
compile_model=compile,
|
||||
)
|
||||
else:
|
||||
# a version without any heuristics for ablation studies
|
||||
model = Sam3VideoInferenceWithInstanceInteractivity(
|
||||
detector=detector,
|
||||
tracker=tracker,
|
||||
score_threshold_detection=0.5,
|
||||
assoc_iou_thresh=0.1,
|
||||
det_nms_thresh=0.1,
|
||||
new_det_thresh=0.7,
|
||||
hotstart_delay=0,
|
||||
hotstart_unmatch_thresh=0,
|
||||
hotstart_dup_thresh=0,
|
||||
suppress_unmatched_only_within_hotstart=True,
|
||||
min_trk_keep_alive=-1,
|
||||
max_trk_keep_alive=30,
|
||||
init_trk_keep_alive=30,
|
||||
suppress_overlapping_based_on_recent_occlusion_threshold=0.7,
|
||||
suppress_det_close_to_boundary=False,
|
||||
fill_hole_area=16,
|
||||
recondition_every_nth_frame=0,
|
||||
masklet_confirmation_enable=False,
|
||||
decrease_trk_keep_alive_for_empty_masklets=False,
|
||||
image_size=1008,
|
||||
image_mean=(0.5, 0.5, 0.5),
|
||||
image_std=(0.5, 0.5, 0.5),
|
||||
compile_model=compile,
|
||||
)
|
||||
|
||||
# Load checkpoint if provided
|
||||
if load_from_HF and checkpoint_path is None:
|
||||
checkpoint_path = download_ckpt_from_hf()
|
||||
if checkpoint_path is not None:
|
||||
with g_pathmgr.open(checkpoint_path, "rb") as f:
|
||||
ckpt = torch.load(f, map_location="cpu", weights_only=True)
|
||||
if "model" in ckpt and isinstance(ckpt["model"], dict):
|
||||
ckpt = ckpt["model"]
|
||||
|
||||
missing_keys, unexpected_keys = model.load_state_dict(
|
||||
ckpt, strict=strict_state_dict_loading
|
||||
)
|
||||
if missing_keys:
|
||||
print(f"Missing keys: {missing_keys}")
|
||||
if unexpected_keys:
|
||||
print(f"Unexpected keys: {unexpected_keys}")
|
||||
|
||||
model.to(device=device)
|
||||
return model
|
||||
|
||||
|
||||
def build_sam3_video_predictor(*model_args, gpus_to_use=None, **model_kwargs):
|
||||
return Sam3VideoPredictorMultiGPU(
|
||||
*model_args, gpus_to_use=gpus_to_use, **model_kwargs
|
||||
)
|
||||
@@ -0,0 +1,941 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import matplotlib.patches as patches
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pycocotools.mask as mask_utils
|
||||
import torch
|
||||
from matplotlib.colors import to_rgb
|
||||
from PIL import Image
|
||||
from skimage.color import lab2rgb, rgb2lab
|
||||
from sklearn.cluster import KMeans
|
||||
from torchvision.ops import masks_to_boxes
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def generate_colors(n_colors=256, n_samples=5000):
|
||||
# Step 1: Random RGB samples
|
||||
np.random.seed(42)
|
||||
rgb = np.random.rand(n_samples, 3)
|
||||
# Step 2: Convert to LAB for perceptual uniformity
|
||||
# print(f"Converting {n_samples} RGB samples to LAB color space...")
|
||||
lab = rgb2lab(rgb.reshape(1, -1, 3)).reshape(-1, 3)
|
||||
# print("Conversion to LAB complete.")
|
||||
# Step 3: k-means clustering in LAB
|
||||
kmeans = KMeans(n_clusters=n_colors, n_init=10)
|
||||
# print(f"Fitting KMeans with {n_colors} clusters on {n_samples} samples...")
|
||||
kmeans.fit(lab)
|
||||
# print("KMeans fitting complete.")
|
||||
centers_lab = kmeans.cluster_centers_
|
||||
# Step 4: Convert LAB back to RGB
|
||||
colors_rgb = lab2rgb(centers_lab.reshape(1, -1, 3)).reshape(-1, 3)
|
||||
colors_rgb = np.clip(colors_rgb, 0, 1)
|
||||
return colors_rgb
|
||||
|
||||
|
||||
COLORS = generate_colors(n_colors=128, n_samples=5000)
|
||||
|
||||
|
||||
def show_img_tensor(img_batch, vis_img_idx=0):
|
||||
MEAN_IMG = np.array([0.485, 0.456, 0.406])
|
||||
STD_IMG = np.array([0.229, 0.224, 0.225])
|
||||
im_tensor = img_batch[vis_img_idx].detach().cpu()
|
||||
assert im_tensor.dim() == 3
|
||||
im_tensor = im_tensor.numpy().transpose((1, 2, 0))
|
||||
im_tensor = (im_tensor * STD_IMG) + MEAN_IMG
|
||||
im_tensor = np.clip(im_tensor, 0, 1)
|
||||
plt.imshow(im_tensor)
|
||||
|
||||
|
||||
def draw_box_on_image(image, box, color=(0, 255, 0)):
|
||||
"""
|
||||
Draws a rectangle on a given PIL image using the provided box coordinates in xywh format.
|
||||
:param image: PIL.Image - The image on which to draw the rectangle.
|
||||
:param box: tuple - A tuple (x, y, w, h) representing the top-left corner, width, and height of the rectangle.
|
||||
:param color: tuple - A tuple (R, G, B) representing the color of the rectangle. Default is red.
|
||||
:return: PIL.Image - The image with the rectangle drawn on it.
|
||||
"""
|
||||
# Ensure the image is in RGB mode
|
||||
image = image.convert("RGB")
|
||||
# Unpack the box coordinates
|
||||
x, y, w, h = box
|
||||
x, y, w, h = int(x), int(y), int(w), int(h)
|
||||
# Get the pixel data
|
||||
pixels = image.load()
|
||||
# Draw the top and bottom edges
|
||||
for i in range(x, x + w):
|
||||
pixels[i, y] = color
|
||||
pixels[i, y + h - 1] = color
|
||||
pixels[i, y + 1] = color
|
||||
pixels[i, y + h] = color
|
||||
pixels[i, y - 1] = color
|
||||
pixels[i, y + h - 2] = color
|
||||
# Draw the left and right edges
|
||||
for j in range(y, y + h):
|
||||
pixels[x, j] = color
|
||||
pixels[x + 1, j] = color
|
||||
pixels[x - 1, j] = color
|
||||
pixels[x + w - 1, j] = color
|
||||
pixels[x + w, j] = color
|
||||
pixels[x + w - 2, j] = color
|
||||
return image
|
||||
|
||||
|
||||
def plot_bbox(
|
||||
img_height,
|
||||
img_width,
|
||||
box,
|
||||
box_format="XYXY",
|
||||
relative_coords=True,
|
||||
color="r",
|
||||
linestyle="solid",
|
||||
text=None,
|
||||
ax=None,
|
||||
):
|
||||
if box_format == "XYXY":
|
||||
x, y, x2, y2 = box
|
||||
w = x2 - x
|
||||
h = y2 - y
|
||||
elif box_format == "XYWH":
|
||||
x, y, w, h = box
|
||||
elif box_format == "CxCyWH":
|
||||
cx, cy, w, h = box
|
||||
x = cx - w / 2
|
||||
y = cy - h / 2
|
||||
else:
|
||||
raise RuntimeError(f"Invalid box_format {box_format}")
|
||||
|
||||
if relative_coords:
|
||||
x *= img_width
|
||||
w *= img_width
|
||||
y *= img_height
|
||||
h *= img_height
|
||||
|
||||
if ax is None:
|
||||
ax = plt.gca()
|
||||
rect = patches.Rectangle(
|
||||
(x, y),
|
||||
w,
|
||||
h,
|
||||
linewidth=1.5,
|
||||
edgecolor=color,
|
||||
facecolor="none",
|
||||
linestyle=linestyle,
|
||||
)
|
||||
ax.add_patch(rect)
|
||||
if text is not None:
|
||||
facecolor = "w"
|
||||
ax.text(
|
||||
x,
|
||||
y - 5,
|
||||
text,
|
||||
color=color,
|
||||
weight="bold",
|
||||
fontsize=8,
|
||||
bbox={"facecolor": facecolor, "alpha": 0.75, "pad": 2},
|
||||
)
|
||||
|
||||
|
||||
def plot_mask(mask, color="r", ax=None):
|
||||
im_h, im_w = mask.shape
|
||||
mask_img = np.zeros((im_h, im_w, 4), dtype=np.float32)
|
||||
mask_img[..., :3] = to_rgb(color)
|
||||
mask_img[..., 3] = mask * 0.5
|
||||
# Use the provided ax or the current axis
|
||||
if ax is None:
|
||||
ax = plt.gca()
|
||||
ax.imshow(mask_img)
|
||||
|
||||
|
||||
def normalize_bbox(bbox_xywh, img_w, img_h):
|
||||
# Assumes bbox_xywh is in XYWH format
|
||||
if isinstance(bbox_xywh, list):
|
||||
assert (
|
||||
len(bbox_xywh) == 4
|
||||
), "bbox_xywh list must have 4 elements. Batching not support except for torch tensors."
|
||||
normalized_bbox = bbox_xywh.copy()
|
||||
normalized_bbox[0] /= img_w
|
||||
normalized_bbox[1] /= img_h
|
||||
normalized_bbox[2] /= img_w
|
||||
normalized_bbox[3] /= img_h
|
||||
else:
|
||||
assert isinstance(
|
||||
bbox_xywh, torch.Tensor
|
||||
), "Only torch tensors are supported for batching."
|
||||
normalized_bbox = bbox_xywh.clone()
|
||||
assert (
|
||||
normalized_bbox.size(-1) == 4
|
||||
), "bbox_xywh tensor must have last dimension of size 4."
|
||||
normalized_bbox[..., 0] /= img_w
|
||||
normalized_bbox[..., 1] /= img_h
|
||||
normalized_bbox[..., 2] /= img_w
|
||||
normalized_bbox[..., 3] /= img_h
|
||||
return normalized_bbox
|
||||
|
||||
|
||||
def visualize_frame_output(frame_idx, video_frames, outputs, figsize=(12, 8)):
|
||||
plt.figure(figsize=figsize)
|
||||
plt.title(f"frame {frame_idx}")
|
||||
img = load_frame(video_frames[frame_idx])
|
||||
img_H, img_W, _ = img.shape
|
||||
plt.imshow(img)
|
||||
for i in range(len(outputs["out_probs"])):
|
||||
box_xywh = outputs["out_boxes_xywh"][i]
|
||||
prob = outputs["out_probs"][i]
|
||||
obj_id = outputs["out_obj_ids"][i]
|
||||
binary_mask = outputs["out_binary_masks"][i]
|
||||
color = COLORS[obj_id % len(COLORS)]
|
||||
plot_bbox(
|
||||
img_H,
|
||||
img_W,
|
||||
box_xywh,
|
||||
text=f"(id={obj_id}, {prob=:.2f})",
|
||||
box_format="XYWH",
|
||||
color=color,
|
||||
)
|
||||
plot_mask(binary_mask, color=color)
|
||||
|
||||
|
||||
def visualize_formatted_frame_output(
|
||||
frame_idx,
|
||||
video_frames,
|
||||
outputs_list,
|
||||
titles=None,
|
||||
points_list=None,
|
||||
points_labels_list=None,
|
||||
figsize=(12, 8),
|
||||
title_suffix="",
|
||||
prompt_info=None,
|
||||
):
|
||||
"""Visualize up to three sets of segmentation masks on a video frame.
|
||||
|
||||
Args:
|
||||
frame_idx: Frame index to visualize
|
||||
image_files: List of image file paths
|
||||
outputs_list: List of {frame_idx: {obj_id: mask_tensor}} or single dict {obj_id: mask_tensor}
|
||||
titles: List of titles for each set of outputs_list
|
||||
points_list: Optional list of point coordinates
|
||||
points_labels_list: Optional list of point labels
|
||||
figsize: Figure size tuple
|
||||
save: Whether to save the visualization to file
|
||||
output_dir: Base output directory when saving
|
||||
scenario_name: Scenario name for organizing saved files
|
||||
title_suffix: Additional title suffix
|
||||
prompt_info: Dictionary with prompt information (boxes, points, etc.)
|
||||
"""
|
||||
# Handle single output dict case
|
||||
if isinstance(outputs_list, dict) and frame_idx in outputs_list:
|
||||
# This is a single outputs dict with frame indices as keys
|
||||
outputs_list = [outputs_list]
|
||||
elif isinstance(outputs_list, dict) and not any(
|
||||
isinstance(k, int) for k in outputs_list.keys()
|
||||
):
|
||||
# This is a single frame's outputs {obj_id: mask}
|
||||
single_frame_outputs = {frame_idx: outputs_list}
|
||||
outputs_list = [single_frame_outputs]
|
||||
|
||||
num_outputs = len(outputs_list)
|
||||
if titles is None:
|
||||
titles = [f"Set {i+1}" for i in range(num_outputs)]
|
||||
assert (
|
||||
len(titles) == num_outputs
|
||||
), "length of `titles` should match that of `outputs_list` if not None."
|
||||
|
||||
_, axes = plt.subplots(1, num_outputs, figsize=figsize)
|
||||
if num_outputs == 1:
|
||||
axes = [axes] # Make it iterable
|
||||
|
||||
img = load_frame(video_frames[frame_idx])
|
||||
img_H, img_W, _ = img.shape
|
||||
|
||||
for idx in range(num_outputs):
|
||||
ax, outputs_set, ax_title = axes[idx], outputs_list[idx], titles[idx]
|
||||
ax.set_title(f"Frame {frame_idx} - {ax_title}{title_suffix}")
|
||||
ax.imshow(img)
|
||||
|
||||
if frame_idx in outputs_set:
|
||||
_outputs = outputs_set[frame_idx]
|
||||
else:
|
||||
print(f"Warning: Frame {frame_idx} not found in outputs_set")
|
||||
continue
|
||||
|
||||
if prompt_info and frame_idx == 0: # Show prompts on first frame
|
||||
if "boxes" in prompt_info:
|
||||
for box in prompt_info["boxes"]:
|
||||
# box is in [x, y, w, h] normalized format
|
||||
x, y, w, h = box
|
||||
plot_bbox(
|
||||
img_H,
|
||||
img_W,
|
||||
[x, y, x + w, y + h], # Convert to XYXY
|
||||
box_format="XYXY",
|
||||
relative_coords=True,
|
||||
color="yellow",
|
||||
linestyle="dashed",
|
||||
text="PROMPT BOX",
|
||||
ax=ax,
|
||||
)
|
||||
|
||||
if "points" in prompt_info and "point_labels" in prompt_info:
|
||||
points = np.array(prompt_info["points"])
|
||||
labels = np.array(prompt_info["point_labels"])
|
||||
# Convert normalized to pixel coordinates
|
||||
points_pixel = points * np.array([img_W, img_H])
|
||||
|
||||
# Draw positive points (green stars)
|
||||
pos_points = points_pixel[labels == 1]
|
||||
if len(pos_points) > 0:
|
||||
ax.scatter(
|
||||
pos_points[:, 0],
|
||||
pos_points[:, 1],
|
||||
color="lime",
|
||||
marker="*",
|
||||
s=200,
|
||||
edgecolor="white",
|
||||
linewidth=2,
|
||||
label="Positive Points",
|
||||
zorder=10,
|
||||
)
|
||||
|
||||
# Draw negative points (red stars)
|
||||
neg_points = points_pixel[labels == 0]
|
||||
if len(neg_points) > 0:
|
||||
ax.scatter(
|
||||
neg_points[:, 0],
|
||||
neg_points[:, 1],
|
||||
color="red",
|
||||
marker="*",
|
||||
s=200,
|
||||
edgecolor="white",
|
||||
linewidth=2,
|
||||
label="Negative Points",
|
||||
zorder=10,
|
||||
)
|
||||
|
||||
objects_drawn = 0
|
||||
for obj_id, binary_mask in _outputs.items():
|
||||
mask_sum = (
|
||||
binary_mask.sum()
|
||||
if hasattr(binary_mask, "sum")
|
||||
else np.sum(binary_mask)
|
||||
)
|
||||
|
||||
if mask_sum > 0: # Only draw if mask has content
|
||||
# Convert to torch tensor if it's not already
|
||||
if not isinstance(binary_mask, torch.Tensor):
|
||||
binary_mask = torch.tensor(binary_mask)
|
||||
|
||||
# Find bounding box from mask
|
||||
if binary_mask.any():
|
||||
box_xyxy = masks_to_boxes(binary_mask.unsqueeze(0)).squeeze()
|
||||
box_xyxy = normalize_bbox(box_xyxy, img_W, img_H)
|
||||
else:
|
||||
# Fallback: create a small box at center
|
||||
box_xyxy = [0.45, 0.45, 0.55, 0.55]
|
||||
|
||||
color = COLORS[obj_id % len(COLORS)]
|
||||
|
||||
plot_bbox(
|
||||
img_H,
|
||||
img_W,
|
||||
box_xyxy,
|
||||
text=f"(id={obj_id})",
|
||||
box_format="XYXY",
|
||||
color=color,
|
||||
ax=ax,
|
||||
)
|
||||
|
||||
# Convert back to numpy for plotting
|
||||
mask_np = (
|
||||
binary_mask.numpy()
|
||||
if isinstance(binary_mask, torch.Tensor)
|
||||
else binary_mask
|
||||
)
|
||||
plot_mask(mask_np, color=color, ax=ax)
|
||||
objects_drawn += 1
|
||||
|
||||
if objects_drawn == 0:
|
||||
ax.text(
|
||||
0.5,
|
||||
0.5,
|
||||
"No objects detected",
|
||||
transform=ax.transAxes,
|
||||
fontsize=16,
|
||||
ha="center",
|
||||
va="center",
|
||||
color="red",
|
||||
weight="bold",
|
||||
)
|
||||
|
||||
# Draw additional points if provided
|
||||
if points_list is not None and points_list[idx] is not None:
|
||||
show_points(
|
||||
points_list[idx], points_labels_list[idx], ax=ax, marker_size=200
|
||||
)
|
||||
|
||||
ax.axis("off")
|
||||
|
||||
plt.tight_layout()
|
||||
plt.show()
|
||||
|
||||
|
||||
def render_masklet_frame(img, outputs, frame_idx=None, alpha=0.5):
|
||||
"""
|
||||
Overlays masklets and bounding boxes on a single image frame.
|
||||
Args:
|
||||
img: np.ndarray, shape (H, W, 3), uint8 or float32 in [0,255] or [0,1]
|
||||
outputs: dict with keys: out_boxes_xywh, out_probs, out_obj_ids, out_binary_masks
|
||||
frame_idx: int or None, for overlaying frame index text
|
||||
alpha: float, mask overlay alpha
|
||||
Returns:
|
||||
overlay: np.ndarray, shape (H, W, 3), uint8
|
||||
"""
|
||||
if img.dtype == np.float32 or img.max() <= 1.0:
|
||||
img = (img * 255).astype(np.uint8)
|
||||
img = img[..., :3] # drop alpha if present
|
||||
height, width = img.shape[:2]
|
||||
overlay = img.copy()
|
||||
|
||||
for i in range(len(outputs["out_probs"])):
|
||||
obj_id = outputs["out_obj_ids"][i]
|
||||
color = COLORS[obj_id % len(COLORS)]
|
||||
color255 = (color * 255).astype(np.uint8)
|
||||
mask = outputs["out_binary_masks"][i]
|
||||
if mask.shape != img.shape[:2]:
|
||||
mask = cv2.resize(
|
||||
mask.astype(np.float32),
|
||||
(img.shape[1], img.shape[0]),
|
||||
interpolation=cv2.INTER_NEAREST,
|
||||
)
|
||||
mask_bool = mask > 0.5
|
||||
for c in range(3):
|
||||
overlay[..., c][mask_bool] = (
|
||||
alpha * color255[c] + (1 - alpha) * overlay[..., c][mask_bool]
|
||||
).astype(np.uint8)
|
||||
|
||||
# Draw bounding boxes and text
|
||||
for i in range(len(outputs["out_probs"])):
|
||||
box_xywh = outputs["out_boxes_xywh"][i]
|
||||
obj_id = outputs["out_obj_ids"][i]
|
||||
prob = outputs["out_probs"][i]
|
||||
color = COLORS[obj_id % len(COLORS)]
|
||||
color255 = tuple(int(x * 255) for x in color)
|
||||
x, y, w, h = box_xywh
|
||||
x1 = int(x * width)
|
||||
y1 = int(y * height)
|
||||
x2 = int((x + w) * width)
|
||||
y2 = int((y + h) * height)
|
||||
cv2.rectangle(overlay, (x1, y1), (x2, y2), color255, 2)
|
||||
if prob is not None:
|
||||
label = f"id={obj_id}, p={prob:.2f}"
|
||||
else:
|
||||
label = f"id={obj_id}"
|
||||
cv2.putText(
|
||||
overlay,
|
||||
label,
|
||||
(x1, max(y1 - 10, 0)),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.5,
|
||||
color255,
|
||||
1,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
# Overlay frame index at the top-left corner
|
||||
if frame_idx is not None:
|
||||
cv2.putText(
|
||||
overlay,
|
||||
f"Frame {frame_idx}",
|
||||
(10, 30),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
1.0,
|
||||
(255, 255, 255),
|
||||
2,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
return overlay
|
||||
|
||||
|
||||
def save_masklet_video(video_frames, outputs, out_path, alpha=0.5, fps=10):
|
||||
# Each outputs dict has keys: "out_boxes_xywh", "out_probs", "out_obj_ids", "out_binary_masks"
|
||||
# video_frames: list of video frame data, same length as outputs_list
|
||||
|
||||
# Read first frame to get size
|
||||
first_img = load_frame(video_frames[0])
|
||||
height, width = first_img.shape[:2]
|
||||
if first_img.dtype == np.float32 or first_img.max() <= 1.0:
|
||||
first_img = (first_img * 255).astype(np.uint8)
|
||||
# Use 'mp4v' for best compatibility with VSCode playback (.mp4 files)
|
||||
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
writer = cv2.VideoWriter("temp.mp4", fourcc, fps, (width, height))
|
||||
|
||||
outputs_list = [
|
||||
(video_frames[frame_idx], frame_idx, outputs[frame_idx])
|
||||
for frame_idx in sorted(outputs.keys())
|
||||
]
|
||||
|
||||
for frame, frame_idx, frame_outputs in tqdm(outputs_list):
|
||||
img = load_frame(frame)
|
||||
overlay = render_masklet_frame(
|
||||
img, frame_outputs, frame_idx=frame_idx, alpha=alpha
|
||||
)
|
||||
writer.write(cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
|
||||
|
||||
writer.release()
|
||||
|
||||
# Re-encode the video for VSCode compatibility using ffmpeg
|
||||
subprocess.run(["ffmpeg", "-y", "-i", "temp.mp4", out_path])
|
||||
print(f"Re-encoded video saved to {out_path}")
|
||||
|
||||
os.remove("temp.mp4") # Clean up temporary file
|
||||
|
||||
|
||||
def save_masklet_image(frame, outputs, out_path, alpha=0.5, frame_idx=None):
|
||||
"""
|
||||
Save a single image with masklet overlays.
|
||||
"""
|
||||
img = load_frame(frame)
|
||||
overlay = render_masklet_frame(img, outputs, frame_idx=frame_idx, alpha=alpha)
|
||||
Image.fromarray(overlay).save(out_path)
|
||||
print(f"Overlay image saved to {out_path}")
|
||||
|
||||
|
||||
def prepare_masks_for_visualization(frame_to_output):
|
||||
# frame_to_obj_masks --> {frame_idx: {'output_probs': np.array, `out_obj_ids`: np.array, `out_binary_masks`: np.array}}
|
||||
for frame_idx, out in frame_to_output.items():
|
||||
_processed_out = {}
|
||||
for idx, obj_id in enumerate(out["out_obj_ids"].tolist()):
|
||||
if out["out_binary_masks"][idx].any():
|
||||
_processed_out[obj_id] = out["out_binary_masks"][idx]
|
||||
frame_to_output[frame_idx] = _processed_out
|
||||
return frame_to_output
|
||||
|
||||
|
||||
def convert_coco_to_masklet_format(
|
||||
annotations, img_info, is_prediction=False, score_threshold=0.5
|
||||
):
|
||||
"""
|
||||
Convert COCO format annotations to format expected by render_masklet_frame
|
||||
"""
|
||||
outputs = {
|
||||
"out_boxes_xywh": [],
|
||||
"out_probs": [],
|
||||
"out_obj_ids": [],
|
||||
"out_binary_masks": [],
|
||||
}
|
||||
|
||||
img_h, img_w = img_info["height"], img_info["width"]
|
||||
|
||||
for idx, ann in enumerate(annotations):
|
||||
# Get bounding box in relative XYWH format
|
||||
if "bbox" in ann:
|
||||
bbox = ann["bbox"]
|
||||
if max(bbox) > 1.0: # Convert absolute to relative coordinates
|
||||
bbox = [
|
||||
bbox[0] / img_w,
|
||||
bbox[1] / img_h,
|
||||
bbox[2] / img_w,
|
||||
bbox[3] / img_h,
|
||||
]
|
||||
else:
|
||||
mask = mask_utils.decode(ann["segmentation"])
|
||||
rows = np.any(mask, axis=1)
|
||||
cols = np.any(mask, axis=0)
|
||||
if np.any(rows) and np.any(cols):
|
||||
rmin, rmax = np.where(rows)[0][[0, -1]]
|
||||
cmin, cmax = np.where(cols)[0][[0, -1]]
|
||||
# Convert to relative XYWH
|
||||
bbox = [
|
||||
cmin / img_w,
|
||||
rmin / img_h,
|
||||
(cmax - cmin + 1) / img_w,
|
||||
(rmax - rmin + 1) / img_h,
|
||||
]
|
||||
else:
|
||||
bbox = [0, 0, 0, 0]
|
||||
|
||||
outputs["out_boxes_xywh"].append(bbox)
|
||||
|
||||
# Get probability/score
|
||||
if is_prediction:
|
||||
prob = ann["score"]
|
||||
else:
|
||||
prob = 1.0 # GT has no probability
|
||||
outputs["out_probs"].append(prob)
|
||||
|
||||
outputs["out_obj_ids"].append(idx)
|
||||
mask = mask_utils.decode(ann["segmentation"])
|
||||
mask = (mask > score_threshold).astype(np.uint8)
|
||||
|
||||
outputs["out_binary_masks"].append(mask)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
def save_side_by_side_visualization(img, gt_anns, pred_anns, noun_phrase):
|
||||
"""
|
||||
Create side-by-side visualization of GT and predictions
|
||||
"""
|
||||
|
||||
# Create side-by-side visualization
|
||||
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 7))
|
||||
|
||||
main_title = f"Noun phrase: '{noun_phrase}'"
|
||||
fig.suptitle(main_title, fontsize=16, fontweight="bold")
|
||||
|
||||
gt_overlay = render_masklet_frame(img, gt_anns, alpha=0.5)
|
||||
ax1.imshow(gt_overlay)
|
||||
ax1.set_title("Ground Truth", fontsize=14, fontweight="bold")
|
||||
ax1.axis("off")
|
||||
|
||||
pred_overlay = render_masklet_frame(img, pred_anns, alpha=0.5)
|
||||
ax2.imshow(pred_overlay)
|
||||
ax2.set_title("Predictions", fontsize=14, fontweight="bold")
|
||||
ax2.axis("off")
|
||||
|
||||
plt.subplots_adjust(top=0.88)
|
||||
plt.tight_layout()
|
||||
|
||||
|
||||
def bitget(val, idx):
|
||||
return (val >> idx) & 1
|
||||
|
||||
|
||||
def pascal_color_map():
|
||||
colormap = np.zeros((512, 3), dtype=int)
|
||||
ind = np.arange(512, dtype=int)
|
||||
for shift in reversed(list(range(8))):
|
||||
for channel in range(3):
|
||||
colormap[:, channel] |= bitget(ind, channel) << shift
|
||||
ind >>= 3
|
||||
|
||||
return colormap.astype(np.uint8)
|
||||
|
||||
|
||||
def draw_masks_to_frame(
|
||||
frame: np.ndarray, masks: np.ndarray, colors: np.ndarray
|
||||
) -> np.ndarray:
|
||||
masked_frame = frame
|
||||
for mask, color in zip(masks, colors):
|
||||
curr_masked_frame = np.where(mask[..., None], color, masked_frame)
|
||||
masked_frame = cv2.addWeighted(masked_frame, 0.75, curr_masked_frame, 0.25, 0)
|
||||
|
||||
if int(cv2.__version__[0]) > 3:
|
||||
contours, _ = cv2.findContours(
|
||||
np.array(mask, dtype=np.uint8).copy(),
|
||||
cv2.RETR_TREE,
|
||||
cv2.CHAIN_APPROX_NONE,
|
||||
)
|
||||
else:
|
||||
_, contours, _ = cv2.findContours(
|
||||
np.array(mask, dtype=np.uint8).copy(),
|
||||
cv2.RETR_TREE,
|
||||
cv2.CHAIN_APPROX_NONE,
|
||||
)
|
||||
|
||||
cv2.drawContours(
|
||||
masked_frame, contours, -1, (255, 255, 255), 7
|
||||
) # White outer contour
|
||||
cv2.drawContours(
|
||||
masked_frame, contours, -1, (0, 0, 0), 5
|
||||
) # Black middle contour
|
||||
cv2.drawContours(
|
||||
masked_frame, contours, -1, color.tolist(), 3
|
||||
) # Original color inner contour
|
||||
return masked_frame
|
||||
|
||||
|
||||
def get_annot_df(file_path: str):
|
||||
with open(file_path, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
dfs = {}
|
||||
|
||||
for k, v in data.items():
|
||||
if k in ("info", "licenses"):
|
||||
dfs[k] = v
|
||||
continue
|
||||
df = pd.DataFrame(v)
|
||||
dfs[k] = df
|
||||
|
||||
return dfs
|
||||
|
||||
|
||||
def get_annot_dfs(file_list: list[str]):
|
||||
dfs = {}
|
||||
for annot_file in tqdm(file_list):
|
||||
dataset_name = Path(annot_file).stem
|
||||
dfs[dataset_name] = get_annot_df(annot_file)
|
||||
|
||||
return dfs
|
||||
|
||||
|
||||
def get_media_dir(media_dir: str, dataset: str):
|
||||
if dataset in ["saco_veval_sav_test", "saco_veval_sav_val"]:
|
||||
return os.path.join(media_dir, "saco_sav", "JPEGImages_24fps")
|
||||
elif dataset in ["saco_veval_yt1b_test", "saco_veval_yt1b_val"]:
|
||||
return os.path.join(media_dir, "saco_yt1b", "JPEGImages_6fps")
|
||||
elif dataset in ["saco_veval_smartglasses_test", "saco_veval_smartglasses_val"]:
|
||||
return os.path.join(media_dir, "saco_sg", "JPEGImages_6fps")
|
||||
elif dataset == "sa_fari_test":
|
||||
return os.path.join(media_dir, "sa_fari", "JPEGImages_6fps")
|
||||
else:
|
||||
raise ValueError(f"Dataset {dataset} not found")
|
||||
|
||||
|
||||
def get_all_annotations_for_frame(
|
||||
dataset_df: pd.DataFrame, video_id: int, frame_idx: int, data_dir: str, dataset: str
|
||||
):
|
||||
media_dir = os.path.join(data_dir, "media")
|
||||
|
||||
# Load the annotation and video data
|
||||
annot_df = dataset_df["annotations"]
|
||||
video_df = dataset_df["videos"]
|
||||
|
||||
# Get the frame
|
||||
video_df_current = video_df[video_df.id == video_id]
|
||||
assert (
|
||||
len(video_df_current) == 1
|
||||
), f"Expected 1 video row, got {len(video_df_current)}"
|
||||
video_row = video_df_current.iloc[0]
|
||||
file_name = video_row.file_names[frame_idx]
|
||||
file_path = os.path.join(
|
||||
get_media_dir(media_dir=media_dir, dataset=dataset), file_name
|
||||
)
|
||||
frame = cv2.imread(file_path)
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# Get the masks and noun phrases annotated in this video in this frame
|
||||
annot_df_current_video = annot_df[annot_df.video_id == video_id]
|
||||
if len(annot_df_current_video) == 0:
|
||||
print(f"No annotations found for video_id {video_id}")
|
||||
return frame, None, None
|
||||
else:
|
||||
empty_mask = np.zeros(frame.shape[:2], dtype=np.uint8)
|
||||
mask_np_pairs = annot_df_current_video.apply(
|
||||
lambda row: (
|
||||
(
|
||||
mask_utils.decode(row.segmentations[frame_idx])
|
||||
if row.segmentations[frame_idx]
|
||||
else empty_mask
|
||||
),
|
||||
row.noun_phrase,
|
||||
),
|
||||
axis=1,
|
||||
)
|
||||
# sort based on noun_phrases
|
||||
mask_np_pairs = sorted(mask_np_pairs, key=lambda x: x[1])
|
||||
masks, noun_phrases = zip(*mask_np_pairs)
|
||||
|
||||
return frame, masks, noun_phrases
|
||||
|
||||
|
||||
def visualize_prompt_overlay(
|
||||
frame_idx,
|
||||
video_frames,
|
||||
title="Prompt Visualization",
|
||||
text_prompt=None,
|
||||
point_prompts=None,
|
||||
point_labels=None,
|
||||
bounding_boxes=None,
|
||||
box_labels=None,
|
||||
obj_id=None,
|
||||
):
|
||||
"""Simple prompt visualization function"""
|
||||
img = Image.fromarray(load_frame(video_frames[frame_idx]))
|
||||
fig, ax = plt.subplots(1, figsize=(6, 4))
|
||||
ax.imshow(img)
|
||||
|
||||
img_w, img_h = img.size
|
||||
|
||||
if text_prompt:
|
||||
ax.text(
|
||||
0.02,
|
||||
0.98,
|
||||
f'Text: "{text_prompt}"',
|
||||
transform=ax.transAxes,
|
||||
fontsize=12,
|
||||
color="white",
|
||||
weight="bold",
|
||||
bbox=dict(boxstyle="round,pad=0.3", facecolor="red", alpha=0.7),
|
||||
verticalalignment="top",
|
||||
)
|
||||
|
||||
if point_prompts:
|
||||
for i, point in enumerate(point_prompts):
|
||||
x, y = point
|
||||
# Convert relative to absolute coordinates
|
||||
x_img, y_img = x * img_w, y * img_h
|
||||
|
||||
# Use different colors for positive/negative points
|
||||
if point_labels and len(point_labels) > i:
|
||||
color = "green" if point_labels[i] == 1 else "red"
|
||||
marker = "o" if point_labels[i] == 1 else "x"
|
||||
else:
|
||||
color = "green"
|
||||
marker = "o"
|
||||
|
||||
ax.plot(
|
||||
x_img,
|
||||
y_img,
|
||||
marker=marker,
|
||||
color=color,
|
||||
markersize=10,
|
||||
markeredgewidth=2,
|
||||
markeredgecolor="white",
|
||||
)
|
||||
ax.text(
|
||||
x_img + 5,
|
||||
y_img - 5,
|
||||
f"P{i+1}",
|
||||
color=color,
|
||||
fontsize=10,
|
||||
weight="bold",
|
||||
bbox=dict(boxstyle="round,pad=0.2", facecolor="white", alpha=0.8),
|
||||
)
|
||||
|
||||
if bounding_boxes:
|
||||
for i, box in enumerate(bounding_boxes):
|
||||
x, y, w, h = box
|
||||
# Convert relative to absolute coordinates
|
||||
x_img, y_img = x * img_w, y * img_h
|
||||
w_img, h_img = w * img_w, h * img_h
|
||||
|
||||
# Use different colors for positive/negative boxes
|
||||
if box_labels and len(box_labels) > i:
|
||||
color = "green" if box_labels[i] == 1 else "red"
|
||||
else:
|
||||
color = "green"
|
||||
|
||||
rect = patches.Rectangle(
|
||||
(x_img, y_img),
|
||||
w_img,
|
||||
h_img,
|
||||
linewidth=2,
|
||||
edgecolor=color,
|
||||
facecolor="none",
|
||||
)
|
||||
ax.add_patch(rect)
|
||||
ax.text(
|
||||
x_img,
|
||||
y_img - 5,
|
||||
f"B{i+1}",
|
||||
color=color,
|
||||
fontsize=10,
|
||||
weight="bold",
|
||||
bbox=dict(boxstyle="round,pad=0.2", facecolor="white", alpha=0.8),
|
||||
)
|
||||
|
||||
# Add object ID info if provided
|
||||
if obj_id is not None:
|
||||
ax.text(
|
||||
0.02,
|
||||
0.02,
|
||||
f"Object ID: {obj_id}",
|
||||
transform=ax.transAxes,
|
||||
fontsize=10,
|
||||
color="white",
|
||||
weight="bold",
|
||||
bbox=dict(boxstyle="round,pad=0.3", facecolor="blue", alpha=0.7),
|
||||
verticalalignment="bottom",
|
||||
)
|
||||
|
||||
ax.set_title(title)
|
||||
ax.axis("off")
|
||||
plt.tight_layout()
|
||||
plt.show()
|
||||
|
||||
|
||||
def plot_results(img, results):
|
||||
plt.figure(figsize=(12, 8))
|
||||
plt.imshow(img)
|
||||
nb_objects = len(results["scores"])
|
||||
print(f"found {nb_objects} object(s)")
|
||||
for i in range(nb_objects):
|
||||
color = COLORS[i % len(COLORS)]
|
||||
plot_mask(results["masks"][i].squeeze(0).cpu(), color=color)
|
||||
w, h = img.size
|
||||
prob = results["scores"][i].item()
|
||||
plot_bbox(
|
||||
h,
|
||||
w,
|
||||
results["boxes"][i].cpu(),
|
||||
text=f"(id={i}, {prob=:.2f})",
|
||||
box_format="XYXY",
|
||||
color=color,
|
||||
relative_coords=False,
|
||||
)
|
||||
|
||||
|
||||
def single_visualization(img, anns, title):
|
||||
"""
|
||||
Create a single image visualization with overlays.
|
||||
"""
|
||||
fig, ax = plt.subplots(figsize=(7, 7))
|
||||
fig.suptitle(title, fontsize=16, fontweight="bold")
|
||||
overlay = render_masklet_frame(img, anns, alpha=0.5)
|
||||
ax.imshow(overlay)
|
||||
ax.axis("off")
|
||||
plt.tight_layout()
|
||||
|
||||
|
||||
def show_mask(mask, ax, obj_id=None, random_color=False):
|
||||
if random_color:
|
||||
color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
|
||||
else:
|
||||
cmap = plt.get_cmap("tab10")
|
||||
cmap_idx = 0 if obj_id is None else obj_id
|
||||
color = np.array([*cmap(cmap_idx)[:3], 0.6])
|
||||
h, w = mask.shape[-2:]
|
||||
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
|
||||
ax.imshow(mask_image)
|
||||
|
||||
|
||||
def show_box(box, ax):
|
||||
x0, y0 = box[0], box[1]
|
||||
w, h = box[2] - box[0], box[3] - box[1]
|
||||
ax.add_patch(
|
||||
plt.Rectangle((x0, y0), w, h, edgecolor="green", facecolor=(0, 0, 0, 0), lw=2)
|
||||
)
|
||||
|
||||
|
||||
def show_points(coords, labels, ax, marker_size=375):
|
||||
pos_points = coords[labels == 1]
|
||||
neg_points = coords[labels == 0]
|
||||
ax.scatter(
|
||||
pos_points[:, 0],
|
||||
pos_points[:, 1],
|
||||
color="green",
|
||||
marker="*",
|
||||
s=marker_size,
|
||||
edgecolor="white",
|
||||
linewidth=1.25,
|
||||
)
|
||||
ax.scatter(
|
||||
neg_points[:, 0],
|
||||
neg_points[:, 1],
|
||||
color="red",
|
||||
marker="*",
|
||||
s=marker_size,
|
||||
edgecolor="white",
|
||||
linewidth=1.25,
|
||||
)
|
||||
|
||||
|
||||
def load_frame(frame):
|
||||
if isinstance(frame, np.ndarray):
|
||||
img = frame
|
||||
elif isinstance(frame, Image.Image):
|
||||
img = np.array(frame)
|
||||
elif isinstance(frame, str) and os.path.isfile(frame):
|
||||
img = plt.imread(frame)
|
||||
else:
|
||||
raise ValueError(f"Invalid video frame type: {type(frame)=}")
|
||||
return img
|
||||
Reference in New Issue
Block a user