15 lines
342 B
Python
15 lines
342 B
Python
from .swin_v2 import SwinTransformerV2MSL
|
|
from .vit import VisionTransformerMSL
|
|
|
|
__all__ = [
|
|
'SwinTransformerV2MSL', 'VisionTransformerMSL'
|
|
]
|
|
|
|
type_mapping = {
|
|
'SwinTransformerV2MSL': SwinTransformerV2MSL,
|
|
'VisionTransformerMSL': VisionTransformerMSL
|
|
}
|
|
|
|
def build_backbone(type, **kwargs):
|
|
return type_mapping[type](**kwargs)
|