build_m.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch

from functools import partial

# from .modeling import ImageEncoderViT, MaskDecoder, PromptEncoder, M, TwoWayTransformer
from modeling import ImageEncoderViT, MaskDecoder, PromptEncoder, M, TwoWayTransformer


def build_m_vit_h(resume=None,
                  ck_image_encoder=None,
                  ck_prompt_encoder=None,
                  ck_mask_decoder=None):
    return _build_m(
        encoder_embed_dim=1280,
        encoder_depth=32,
        encoder_num_heads=16,
        encoder_global_attn_indexes=[7, 15, 23, 31],
        resume=resume,
        ck_image_encoder=ck_image_encoder,
        ck_prompt_encoder=ck_prompt_encoder,
        ck_mask_decoder=ck_mask_decoder,
    )


build_m = build_m_vit_h


def build_m_vit_l(resume=None,
                  ck_image_encoder=None,
                  ck_prompt_encoder=None,
                  ck_mask_decoder=None):
    return _build_m(
        encoder_embed_dim=1024,
        encoder_depth=24,
        encoder_num_heads=16,
        encoder_global_attn_indexes=[5, 11, 17, 23],
        resume=resume,
        ck_image_encoder=ck_image_encoder,
        ck_prompt_encoder=ck_prompt_encoder,
        ck_mask_decoder=ck_mask_decoder,
    )


def build_m_vit_b(resume=None,
                  ck_image_encoder=None,
                  ck_prompt_encoder=None,
                  ck_mask_decoder=None):
    return _build_m(
        encoder_embed_dim=768,
        encoder_depth=12,
        encoder_num_heads=12,
        encoder_global_attn_indexes=[2, 5, 8, 11],
        resume=resume,
        ck_image_encoder=ck_image_encoder,
        ck_prompt_encoder=ck_prompt_encoder,
        ck_mask_decoder=ck_mask_decoder,
    )


m_model_registry = {
    "default": build_m_vit_h,
    "vit_h": build_m_vit_h,
    "vit_l": build_m_vit_l,
    "vit_b": build_m_vit_b,
}


def _build_m(
    encoder_embed_dim,
    encoder_depth,
    encoder_num_heads,
    encoder_global_attn_indexes,
    resume=None,
    ck_image_encoder=None,
    ck_prompt_encoder=None,
    ck_mask_decoder=None,
):
    prompt_embed_dim = 256
    image_size = 1024
    vit_patch_size = 16
    image_embedding_size = image_size // vit_patch_size
    m = M(
        image_encoder=ImageEncoderViT(
            depth=encoder_depth,
            embed_dim=encoder_embed_dim,
            img_size=image_size,
            mlp_ratio=4,
            norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
            num_heads=encoder_num_heads,
            patch_size=vit_patch_size,
            qkv_bias=True,
            use_rel_pos=True,
            global_attn_indexes=encoder_global_attn_indexes,
            window_size=14,
            out_chans=prompt_embed_dim,
        ),
        prompt_encoder=PromptEncoder(
            embed_dim=prompt_embed_dim,
            model_name='ViT-L-14-336',
            pretrained='openai',
            text_dim=768,
            depth=3,
            image_embedding_size=(image_embedding_size, image_embedding_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,
                dropout=0.1,
                droppath=0.4,
            ),
            transformer_dim=prompt_embed_dim,
            iou_head_depth=3,
            iou_head_hidden_dim=256,
        ),
    )
    
    if resume:
        with open(resume, "rb") as f:
            state_dict = torch.load(f)
        m.load_state_dict(state_dict['model'])
    else:
        if ck_image_encoder:
            with open(ck_image_encoder, "rb") as f:
                state_dict = torch.load(f)
            m.image_encoder.load_state_dict(state_dict)
        if ck_prompt_encoder:
            with open(ck_prompt_encoder, "rb") as f:
                state_dict = torch.load(f)
            m.prompt_encoder.load_state_dict(state_dict, strict=False)
        if ck_mask_decoder:
            with open(ck_mask_decoder, "rb") as f:
                state_dict = torch.load(f)
            m.mask_decoder.load_state_dict(state_dict, strict=False)
    return m


# model = build_sam_vit_h()
# for name, param in model.named_parameters():
#     print(name)

# model = build_sam_vit_h(checkpoint='checkpoints/sam_vit_h_4b8939.pth')

# model = build_m_vit_h(ck_image_encoder='./pretrained/image_encoder/sam_vit_h.pth',
#                       ck_mask_decoder='./pretrained/mask_decoder/sam_vit_h_decoder.pth')