| |
| |
| |
|
|
| from src.efficientvit.models.efficientvit import (EfficientViTSam, |
| efficientvit_sam_l0, |
| efficientvit_sam_l1, |
| efficientvit_sam_l2, |
| efficientvit_sam_xl0, |
| efficientvit_sam_xl1) |
| from src.efficientvit.models.nn.norm import set_norm_eps |
| from src.efficientvit.models.utils import load_state_dict_from_file |
|
|
| __all__ = ["create_sam_model"] |
|
|
|
|
| REGISTERED_SAM_MODEL: dict[str, str] = { |
| "l0": "assets/checkpoints/sam/l0.pt", |
| "l1": "assets/checkpoints/sam/l1.pt", |
| "l2": "assets/checkpoints/sam/l2.pt", |
| "xl0": "assets/checkpoints/sam/xl0.pt", |
| "xl1": "assets/checkpoints/sam/xl1.pt", |
| } |
|
|
|
|
| def create_sam_model( |
| name: str, pretrained=True, weight_url: str or None = None, **kwargs |
| ) -> EfficientViTSam: |
| model_dict = { |
| "l0": efficientvit_sam_l0, |
| "l1": efficientvit_sam_l1, |
| "l2": efficientvit_sam_l2, |
| "xl0": efficientvit_sam_xl0, |
| "xl1": efficientvit_sam_xl1, |
| } |
|
|
| model_id = name.split("-")[0] |
| if model_id not in model_dict: |
| raise ValueError( |
| f"Do not find {name} in the model zoo. List of models: {list(model_dict.keys())}" |
| ) |
| else: |
| model = model_dict[model_id](**kwargs) |
| set_norm_eps(model, 1e-6) |
|
|
| if pretrained: |
| weight_url = weight_url or REGISTERED_SAM_MODEL.get(name, None) |
| if weight_url is None: |
| raise ValueError(f"Do not find the pretrained weight of {name}.") |
| else: |
| weight = load_state_dict_from_file(weight_url) |
| model.load_state_dict(weight) |
| return model |
|
|
