|
|
|
|
|
|
|
|
|
|
|
|
|
import torch |
|
from functools import partial |
|
from segment_anything.modeling import ImageEncoderViT, MaskDecoder, PromptEncoder, Sam, TwoWayTransformer |
|
from RepViTSAM import repvit |
|
from timm.models import create_model |
|
|
|
def build_sam_repvit(checkpoint=None): |
|
prompt_embed_dim = 256 |
|
image_size = 1024 |
|
vit_patch_size = 16 |
|
image_embedding_size = image_size // vit_patch_size |
|
repvit_sam = Sam( |
|
image_encoder=create_model('repvit'), |
|
prompt_encoder=PromptEncoder( |
|
embed_dim=prompt_embed_dim, |
|
image_embedding_size=(image_embedding_size, image_embedding_size), |
|
input_image_size=(image_size, image_size), |
|
mask_in_chans=16, |
|
), |
|
mask_decoder=MaskDecoder( |
|
num_multimask_outputs=3, |
|
transformer=TwoWayTransformer( |
|
depth=2, |
|
embedding_dim=prompt_embed_dim, |
|
mlp_dim=2048, |
|
num_heads=8, |
|
), |
|
transformer_dim=prompt_embed_dim, |
|
iou_head_depth=3, |
|
iou_head_hidden_dim=256, |
|
), |
|
pixel_mean=[123.675, 116.28, 103.53], |
|
pixel_std=[58.395, 57.12, 57.375], |
|
) |
|
|
|
repvit_sam.eval() |
|
if checkpoint is not None: |
|
with open(checkpoint, "rb") as f: |
|
state_dict = torch.load(f) |
|
repvit_sam.load_state_dict(state_dict) |
|
return repvit_sam |
|
|
|
from functools import partial |
|
|
|
sam_model_registry = { |
|
"repvit": partial(build_sam_repvit), |
|
} |
|
|
