init

lib/models/backbones/__init__.py

@@ -0,0 +1,14 @@
+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)