segmodels-keras 0.1.0.dev4__tar.gz → 0.1.0.dev6__tar.gz

{segmodels_keras-0.1.0.dev4 → segmodels_keras-0.1.0.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: segmodels_keras
-Version: 0.1.0.dev4
+Version: 0.1.0.dev6
 Summary: Image segmentation models with pre-trained backbones with Keras.
 Home-page: https://github.com/orthoseg/segmodels_keras
 Author: Pieter Roggemans

{segmodels_keras-0.1.0.dev4 → segmodels_keras-0.1.0.dev6}/README.rst

@@ -1,8 +1,11 @@
-|Actions Status| |PyPI version|
+|Actions Status| |Code coverage| |PyPI version|
 
 .. |Actions Status| image:: https://github.com/orthoseg/segmodels_keras/actions/workflows/tests.yml/badge.svg?branch=main
    :target: https://github.com/orthoseg/segmodels_keras/actions/workflows/tests.yml?query=workflow%3ATests
 
+.. |Code coverage| image:: https://codecov.io/gh/orthoseg/segmodels_keras/graph/badge.svg?token=o01OyTMDuj 
+   :target: https://codecov.io/gh/orthoseg/segmodels_keras
+
 .. |PyPI version| image:: https://img.shields.io/pypi/v/segmodels-keras.svg
    :target: https://pypi.org/project/segmodels-keras
 

{segmodels_keras-0.1.0.dev4 → segmodels_keras-0.1.0.dev6}/segmodels_keras/_version.py

@@ -1,2 +1,2 @@
-__versioninfo__ = (0, 1, 0, "dev4")
+__versioninfo__ = (0, 1, 0, "dev6")
 __version__ = ".".join(map(str, __versioninfo__))

{segmodels_keras-0.1.0.dev4 → segmodels_keras-0.1.0.dev6}/segmodels_keras/models/fpn.py

@@ -107,6 +107,8 @@ def build_fpn(
     use_batchnorm=True,
     aggregation="sum",
     dropout=None,
+    weights_notop=None,
+    freeze_notop=False,
 ):
     input_ = backbone.input
     x = backbone.output
@@ -172,6 +174,14 @@ def build_fpn(
         size=(2, 2), interpolation="bilinear", name="final_upsampling"
     )(x)
 
+    # load weights without top if path provided
+    if weights_notop is not None:
+        model_notop = models.Model(input_, x)
+        model_notop.load_weights(weights_notop)
+        if freeze_notop:
+            freeze_model(model_notop)
+        x = model_notop.output
+
     # model head (define number of output classes)
     x = layers.Conv2D(
         filters=classes,
@@ -200,6 +210,8 @@ def FPN(
     classes=21,
     activation="softmax",
     weights=None,
+    weights_notop=None,
+    freeze_notop=False,
     encoder_weights="imagenet",
     encoder_freeze=False,
     encoder_features="default",
@@ -219,7 +231,11 @@ def FPN(
             your model be able to process images af any size, but ``H`` and ``W`` of
             input images should be divisible by factor ``32``.
         classes: a number of classes for output (output shape - ``(h, w, classes)``).
-        weights: optional, path to model weights.
+        weights: optional, path to model weights to be loaded.
+        weights_notop: optional, path to model weights without top (without segmentation
+            head) to be loaded.
+        freeze_notop: if ``True``, set all layers of the model except the top as
+            non-trainable.
         activation: name of one of ``keras.activations`` for last model layer
             (e.g. ``sigmoid``, ``softmax``, ``linear``).
         encoder_weights: one of ``None`` (random initialization), ``imagenet``
@@ -265,6 +281,8 @@ def FPN(
         activation=activation,
         classes=classes,
         aggregation=pyramid_aggregation,
+        weights_notop=weights_notop,
+        freeze_notop=freeze_notop,
     )
 
     # lock encoder weights for fine-tuning

{segmodels_keras-0.1.0.dev4 → segmodels_keras-0.1.0.dev6}/segmodels_keras/models/linknet.py

@@ -141,6 +141,8 @@ def build_linknet(
     classes=1,
     activation="sigmoid",
     use_batchnorm=True,
+    weights_notop=None,
+    freeze_notop=False,
 ):
     input_ = backbone.input
     x = backbone.output
@@ -169,6 +171,14 @@ def build_linknet(
             x, skip
         )
 
+    # load weights without top if path provided
+    if weights_notop is not None:
+        model_notop = models.Model(input_, x)
+        model_notop.load_weights(weights_notop)
+        if freeze_notop:
+            freeze_model(model_notop)
+        x = model_notop.output
+
     # model head (define number of output classes)
     x = layers.Conv2D(
         filters=classes,
@@ -196,6 +206,8 @@ def Linknet(
     classes=1,
     activation="sigmoid",
     weights=None,
+    weights_notop=None,
+    freeze_notop=False,
     encoder_weights="imagenet",
     encoder_freeze=False,
     encoder_features="default",
@@ -219,7 +231,11 @@ def Linknet(
         classes: a number of classes for output (output shape - ``(h, w, classes)``).
         activation: name of one of ``keras.activations`` for last model layer
             (e.g. ``sigmoid``, ``softmax``, ``linear``).
-        weights: optional, path to model weights.
+        weights: optional, path to model weights to be loaded.
+        weights_notop: optional, path to model weights without top (without segmentation
+            head) to be loaded.
+        freeze_notop: if ``True``, set all layers of the model except the top as
+            non-trainable.
         encoder_weights: one of ``None`` (random initialization), ``imagenet``
             (pre-training on ImageNet).
         encoder_freeze: if ``True`` set all layers of encoder (backbone model) as
@@ -274,6 +290,8 @@ def Linknet(
         activation=activation,
         n_upsample_blocks=len(decoder_filters),
         use_batchnorm=decoder_use_batchnorm,
+        weights_notop=weights_notop,
+        freeze_notop=freeze_notop,
     )
 
     # lock encoder weights for fine-tuning

{segmodels_keras-0.1.0.dev4 → segmodels_keras-0.1.0.dev6}/segmodels_keras/models/pspnet.py

@@ -124,6 +124,8 @@ def build_psp(
     classes=21,
     activation="softmax",
     dropout=None,
+    weights_notop=None,
+    freeze_notop=False,
 ):
     input_ = backbone.input
     x = (
@@ -147,6 +149,14 @@ def build_psp(
     if dropout is not None:
         x = layers.SpatialDropout2D(dropout, name="spatial_dropout")(x)
 
+    # load weights without top if path provided
+    if weights_notop is not None:
+        model_notop = models.Model(input_, x)
+        model_notop.load_weights(weights_notop)
+        if freeze_notop:
+            freeze_model(model_notop)
+        x = model_notop.output
+
     # model head
     x = layers.Conv2D(
         filters=classes,
@@ -177,6 +187,8 @@ def PSPNet(
     classes=21,
     activation="softmax",
     weights=None,
+    weights_notop=None,
+    freeze_notop=False,
     encoder_weights="imagenet",
     encoder_freeze=False,
     downsample_factor=8,
@@ -197,7 +209,11 @@ def PSPNet(
         classes: a number of classes for output (output shape - ``(h, w, classes)``).
         activation: name of one of ``keras.activations`` for last model layer
             (e.g. ``sigmoid``, ``softmax``, ``linear``).
-        weights: optional, path to model weights.
+        weights: optional, path to model weights to be loaded.
+        weights_notop: optional, path to model weights without top (without segmentation
+            head) to be loaded.
+        freeze_notop: if ``True``, set all layers of the model except the top as
+            non-trainable.
         encoder_weights: one of ``None`` (random initialization), ``imagenet``
             (pre-training on ImageNet).
         encoder_freeze: if ``True`` set all layers of encoder (backbone model) as
@@ -250,6 +266,8 @@ def PSPNet(
         classes=classes,
         activation=activation,
         dropout=psp_dropout,
+        weights_notop=weights_notop,
+        freeze_notop=freeze_notop,
     )
 
     # lock encoder weights for fine-tuning

{segmodels_keras-0.1.0.dev4 → segmodels_keras-0.1.0.dev6}/segmodels_keras/models/unet.py

@@ -113,6 +113,8 @@ def build_unet(
     classes=1,
     activation="sigmoid",
     use_batchnorm=True,
+    weights_notop=None,
+    freeze_notop=False,
 ):
     input_ = backbone.input
     x = backbone.output
@@ -141,6 +143,14 @@ def build_unet(
             x, skip
         )
 
+    # load weights without top if path provided
+    if weights_notop is not None:
+        model_notop = models.Model(input_, x)
+        model_notop.load_weights(weights_notop)
+        if freeze_notop:
+            freeze_model(model_notop)
+        x = model_notop.output
+
     # model head (define number of output classes)
     x = layers.Conv2D(
         filters=classes,
@@ -169,6 +179,8 @@ def Unet(
     classes=1,
     activation="sigmoid",
     weights=None,
+    weights_notop=None,
+    freeze_notop=False,
     encoder_weights="imagenet",
     encoder_freeze=False,
     encoder_features="default",
@@ -190,7 +202,11 @@ def Unet(
         classes: a number of classes for output (output shape - ``(h, w, classes)``).
         activation: name of one of ``keras.activations`` for last model layer
             (e.g. ``sigmoid``, ``softmax``, ``linear``).
-        weights: optional, path to model weights.
+        weights: optional, path to model weights to be loaded.
+        weights_notop: optional, path to model weights without top (without segmentation
+            head) to be loaded.
+        freeze_notop: if ``True``, set all layers of the model except the top as
+            non-trainable.
         encoder_weights: one of ``None`` (random initialization), ``imagenet``
             (pre-training on ImageNet).
         encoder_freeze: if ``True`` set all layers of encoder (backbone model) as
@@ -245,6 +261,8 @@ def Unet(
         activation=activation,
         n_upsample_blocks=len(decoder_filters),
         use_batchnorm=decoder_use_batchnorm,
+        weights_notop=weights_notop,
+        freeze_notop=freeze_notop,
     )
 
     # lock encoder weights for fine-tuning

{segmodels_keras-0.1.0.dev4 → segmodels_keras-0.1.0.dev6}/segmodels_keras/utils.py

@@ -1,8 +1,40 @@
 """Utility functions for segmentation models"""
 
+from pathlib import Path
+
 from keras import models
 
 
+def save_model_weights_notop(
+    model: models.Model, decoder: str, path: str | Path, overwrite: bool = True
+):
+    """Save model weights without top (without segmentation head).
+
+    The weights saved like this can be used to preload a segmentation model for
+    fine-tuning by passing the path to these weights to ``weights_notop`` argument of
+    the model constructor, e.g. ``Unet(weights_notop="path/to/weights.h5")``.
+
+    Args:
+        model (``keras.models.Model``): instance of keras model
+        decoder: type of the decoder part of the model. Should be one of ``fpn``,
+            ``linknet``, ``unet``, ``pspnet``.
+        path (str | Path): path to save model weights
+        overwrite (bool): whether to overwrite existing file at ``path``.
+            Defaults to ``True``.
+
+    """
+    decoder_top_layers = {"fpn": 2, "linknet": 2, "unet": 2, "pspnet": 3}
+    if decoder not in decoder_top_layers:
+        raise ValueError(
+            f"Decoder should be one of {decoder_top_layers.keys()}, got {decoder}"
+        )
+
+    # Add 1 to the slice to actually remove the number of top layers specified
+    nb_top_layers = decoder_top_layers[decoder] + 1
+    model_notop = models.Model(model.input, model.layers[-nb_top_layers].output)
+    model_notop.save_weights(str(path), overwrite=overwrite)
+
+
 def set_trainable(model, recompile=True, **kwargs):  # noqa: ARG001
     """Set all layers of model trainable and recompile it
 

{segmodels_keras-0.1.0.dev4 → segmodels_keras-0.1.0.dev6}/segmodels_keras.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: segmodels_keras
-Version: 0.1.0.dev4
+Version: 0.1.0.dev6
 Summary: Image segmentation models with pre-trained backbones with Keras.
 Home-page: https://github.com/orthoseg/segmodels_keras
 Author: Pieter Roggemans

{segmodels_keras-0.1.0.dev4 → segmodels_keras-0.1.0.dev6}/tests/test_utils.py

@@ -106,5 +106,45 @@ def test_activity_reg(case):
     _test_regularizer(model, reg_model, x, y)
 
 
+@pytest.mark.parametrize("decoder", ["fpn", "linknet", "unet", "pspnet"])
+def test_save_model_weights_notop(tmp_path, decoder):
+    backbone_name = "mobilenetv2"
+    if decoder == "pspnet":
+        model = sm.PSPNet(backbone_name, encoder_weights=None)
+    elif decoder == "fpn":
+        model = sm.FPN(backbone_name, encoder_weights=None)
+    elif decoder == "linknet":
+        model = sm.Linknet(backbone_name, encoder_weights=None)
+    elif decoder == "unet":
+        model = sm.Unet(backbone_name, encoder_weights=None)
+    else:
+        raise ValueError(f"Incorrect decoder {decoder}")
+
+    # Test saving the model weights without the top layers
+    path = tmp_path / f"{decoder}.weights.h5"
+    sm.utils.save_model_weights_notop(model, decoder=decoder, path=path)
+
+    # Test loading the model weights without the top layers again
+    assert path.exists()
+    if decoder == "pspnet":
+        model = sm.PSPNet(backbone_name, weights_notop=path, freeze_notop=True)
+    elif decoder == "fpn":
+        model = sm.FPN(backbone_name, weights_notop=path, freeze_notop=True)
+    elif decoder == "linknet":
+        model = sm.Linknet(backbone_name, weights_notop=path, freeze_notop=True)
+    elif decoder == "unet":
+        model = sm.Unet(backbone_name, weights_notop=path, freeze_notop=True)
+    else:
+        raise ValueError(f"Incorrect decoder {decoder}")
+
+
+def test_save_model_weights_notop_invalid_decoder():
+    model = keras.models.Model()
+    with pytest.raises(ValueError, match="Decoder should be one of"):
+        sm.utils.save_model_weights_notop(
+            model, decoder="invalid_decoder", path="path/to/weights.h5"
+        )
+
+
 if __name__ == "__main__":
     pytest.main([__file__])