Fix loading from config

Author: qubvel

segmentation_models_pytorch/base/hub_mixin.py

@@ -136,3 +136,15 @@ def from_pretrained(pretrained_model_name_or_path: str, *args, **kwargs):
 
     model_class = getattr(smp, model_class_name)
     return model_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+
+def supports_config_loading(func):
+    """Decorator to filter special config kwargs"""
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        kwargs = {k: v for k, v in kwargs.items() if not k.startswith("_")}
+        return func(self, *args, **kwargs)
+
+    return wrapper
+from segmentation_models_pytorch.base.hub_mixin import supports_config_loading

segmentation_models_pytorch/decoders/deeplabv3/model.py

@@ -8,6 +8,7 @@
     SegmentationModel,
 )
 from segmentation_models_pytorch.encoders import get_encoder
+from segmentation_models_pytorch.base.hub_mixin import supports_config_loading
 
 from .decoder import DeepLabV3Decoder, DeepLabV3PlusDecoder
 
@@ -54,6 +55,7 @@ class DeepLabV3(SegmentationModel):
 
     """
 
+    @supports_config_loading
     def __init__(
         self,
         encoder_name: str = "resnet34",
@@ -163,6 +165,7 @@ class DeepLabV3Plus(SegmentationModel):
 
     """
 
+    @supports_config_loading
     def __init__(
         self,
         encoder_name: str = "resnet34",

segmentation_models_pytorch/decoders/fpn/model.py

@@ -6,6 +6,7 @@
     SegmentationModel,
 )
 from segmentation_models_pytorch.encoders import get_encoder
+from segmentation_models_pytorch.base.hub_mixin import supports_config_loading
 
 from .decoder import FPNDecoder
 
@@ -51,6 +52,7 @@ class FPN(SegmentationModel):
 
     """
 
+    @supports_config_loading
     def __init__(
         self,
         encoder_name: str = "resnet34",

segmentation_models_pytorch/decoders/linknet/model.py

@@ -6,6 +6,7 @@
     SegmentationModel,
 )
 from segmentation_models_pytorch.encoders import get_encoder
+from segmentation_models_pytorch.base.hub_mixin import supports_config_loading
 
 from .decoder import LinknetDecoder
 
@@ -53,6 +54,7 @@ class Linknet(SegmentationModel):
         https://arxiv.org/abs/1707.03718
     """
 
+    @supports_config_loading
     def __init__(
         self,
         encoder_name: str = "resnet34",

segmentation_models_pytorch/decoders/manet/model.py

@@ -6,6 +6,7 @@
     SegmentationModel,
 )
 from segmentation_models_pytorch.encoders import get_encoder
+from segmentation_models_pytorch.base.hub_mixin import supports_config_loading
 
 from .decoder import MAnetDecoder
 
@@ -56,6 +57,7 @@ class MAnet(SegmentationModel):
 
     """
 
+    @supports_config_loading
     def __init__(
         self,
         encoder_name: str = "resnet34",

segmentation_models_pytorch/decoders/pan/model.py

@@ -6,6 +6,7 @@
     SegmentationModel,
 )
 from segmentation_models_pytorch.encoders import get_encoder
+from segmentation_models_pytorch.base.hub_mixin import supports_config_loading
 
 from .decoder import PANDecoder
 
@@ -53,6 +54,7 @@ class PAN(SegmentationModel):
 
     """
 
+    @supports_config_loading
     def __init__(
         self,
         encoder_name: str = "resnet34",

segmentation_models_pytorch/decoders/pspnet/model.py

@@ -6,6 +6,7 @@
     SegmentationModel,
 )
 from segmentation_models_pytorch.encoders import get_encoder
+from segmentation_models_pytorch.base.hub_mixin import supports_config_loading
 
 from .decoder import PSPDecoder
 
@@ -54,6 +55,7 @@ class PSPNet(SegmentationModel):
         https://arxiv.org/abs/1612.01105
     """
 
+    @supports_config_loading
     def __init__(
         self,
         encoder_name: str = "resnet34",

segmentation_models_pytorch/decoders/segformer/model.py

@@ -6,6 +6,7 @@
     SegmentationModel,
 )
 from segmentation_models_pytorch.encoders import get_encoder
+from segmentation_models_pytorch.base.hub_mixin import supports_config_loading
 
 from .decoder import SegformerDecoder
 
@@ -46,6 +47,7 @@ class Segformer(SegmentationModel):
 
     """
 
+    @supports_config_loading
     def __init__(
         self,
         encoder_name: str = "resnet34",

segmentation_models_pytorch/decoders/unet/model.py

@@ -1,11 +1,12 @@
-from typing import Any, List, Optional, Union
+from typing import Any, Optional, Union, Tuple, Callable
 
 from segmentation_models_pytorch.base import (
     ClassificationHead,
     SegmentationHead,
     SegmentationModel,
 )
 from segmentation_models_pytorch.encoders import get_encoder
+from segmentation_models_pytorch.base.hub_mixin import supports_config_loading
 
 from .decoder import UnetDecoder
 
@@ -55,17 +56,18 @@ class Unet(SegmentationModel):
 
     """
 
+    @supports_config_loading
     def __init__(
         self,
         encoder_name: str = "resnet34",
         encoder_depth: int = 5,
         encoder_weights: Optional[str] = "imagenet",
         decoder_use_batchnorm: bool = True,
-        decoder_channels: List[int] = (256, 128, 64, 32, 16),
+        decoder_channels: Tuple[int, ...] = (256, 128, 64, 32, 16),
         decoder_attention_type: Optional[str] = None,
         in_channels: int = 3,
         classes: int = 1,
-        activation: Optional[Union[str, callable]] = None,
+        activation: Optional[Union[str, Callable]] = None,
         aux_params: Optional[dict] = None,
         **kwargs: dict[str, Any],
     ):

segmentation_models_pytorch/decoders/unetplusplus/model.py

@@ -6,6 +6,7 @@
     SegmentationModel,
 )
 from segmentation_models_pytorch.encoders import get_encoder
+from segmentation_models_pytorch.base.hub_mixin import supports_config_loading
 
 from .decoder import UnetPlusPlusDecoder
 
@@ -55,6 +56,7 @@ class UnetPlusPlus(SegmentationModel):
 
     """
 
+    @supports_config_loading
     def __init__(
         self,
         encoder_name: str = "resnet34",