[models] Create specialized model per library

src/mednet/libs/classification/models/alexnet.py

@@ -11,12 +11,13 @@ import torch.optim.optimizer
 import torch.utils.data
 import torchvision.models as models
 from mednet.libs.common.data.typing import TransformSequence
-from mednet.libs.common.models.model import Model
+
+from .classification_model import ClassificationModel
 
 logger = logging.getLogger(__name__)
 
 
-class Alexnet(Model):
+class Alexnet(ClassificationModel):
     """Alexnet module.
 
     Note: only usable with a normalized dataset
@@ -115,13 +116,7 @@ class Alexnet(Model):
             )
             self.normalizer = make_imagenet_normalizer()
         else:
-            from .normalizer import make_z_normalizer
-
-            logger.info(
-                f"Uninitialised {self.name} model - "
-                f"computing z-norm factors from train dataloader.",
-            )
-            self.normalizer = make_z_normalizer(dataloader)
+            super().set_normalizer(dataloader)
 
     def training_step(self, batch, _):
         images = batch[0]

src/mednet/libs/classification/models/classification_model.py

+# SPDX-FileCopyrightText: Copyright © 2024 Idiap Research Institute <contact@idiap.ch>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import logging
+import typing
+
+import torch
+import torch.nn
+import torch.optim.optimizer
+import torch.utils.data
+from mednet.libs.common.data.typing import TransformSequence
+from mednet.libs.common.models.model import Model
+
+logger = logging.getLogger("mednet")
+
+
+class ClassificationModel(Model):
+    """Base model for classification task.
+
+    Parameters
+    ----------
+    loss_type
+        The loss to be used for training and evaluation.
+
+        .. warning::
+
+           The loss should be set to always return batch averages (as opposed
+           to the batch sum), as our logging system expects it so.
+    loss_arguments
+        Arguments to the loss.
+    optimizer_type
+        The type of optimizer to use for training.
+    optimizer_arguments
+        Arguments to the optimizer after ``params``.
+    model_transforms
+        An optional sequence of torch modules containing transforms to be
+        applied on the input **before** it is fed into the network.
+    augmentation_transforms
+        An optional sequence of torch modules containing transforms to be
+        applied on the input **before** it is fed into the network.
+    num_classes
+        Number of outputs (classes) for this model.
+    """
+
+    def __init__(
+        self,
+        loss_type: torch.nn.Module = torch.nn.BCEWithLogitsLoss,
+        loss_arguments: dict[str, typing.Any] = {},
+        optimizer_type: type[torch.optim.Optimizer] = torch.optim.Adam,
+        optimizer_arguments: dict[str, typing.Any] = {},
+        model_transforms: TransformSequence = [],
+        augmentation_transforms: TransformSequence = [],
+        num_classes: int = 1,
+    ):
+        super().__init__(
+            loss_type,
+            loss_arguments,
+            optimizer_type,
+            optimizer_arguments,
+            model_transforms,
+            augmentation_transforms,
+            num_classes,
+        )
+
+    def set_normalizer(self, dataloader: torch.utils.data.DataLoader) -> None:
+        """Initialize the input normalizer for the current model.
+
+        Parameters
+        ----------
+        dataloader
+            A torch Dataloader from which to compute the mean and std.
+        """
+
+        from .normalizer import make_z_normalizer
+
+        logger.info(
+            f"Uninitialised {self.name} model - "
+            f"computing z-norm factors from train dataloader.",
+        )
+        self.normalizer = make_z_normalizer(dataloader)

src/mednet/libs/classification/models/densenet.py

@@ -11,12 +11,13 @@ import torch.optim.optimizer
 import torch.utils.data
 import torchvision.models as models
 from mednet.libs.common.data.typing import TransformSequence
-from mednet.libs.common.models.model import Model
+
+from .classification_model import ClassificationModel
 
 logger = logging.getLogger(__name__)
 
 
-class Densenet(Model):
+class Densenet(ClassificationModel):
     """Densenet-121 module.
 
     Parameters
@@ -118,13 +119,7 @@ class Densenet(Model):
             )
             self.normalizer = make_imagenet_normalizer()
         else:
-            from .normalizer import make_z_normalizer
-
-            logger.info(
-                f"Uninitialised {self.name} model - "
-                f"computing z-norm factors from train dataloader.",
-            )
-            self.normalizer = make_z_normalizer(dataloader)
+            super().set_normalizer(dataloader)
 
     def training_step(self, batch, _):
         images = batch[0]

src/mednet/libs/classification/models/pasa.py

@@ -11,12 +11,13 @@ import torch.nn.functional as F  # noqa: N812
 import torch.optim.optimizer
 import torch.utils.data
 from mednet.libs.common.data.typing import TransformSequence
-from mednet.libs.common.models.model import Model
+
+from .classification_model import ClassificationModel
 
 logger = logging.getLogger(__name__)
 
 
-class Pasa(Model):
+class Pasa(ClassificationModel):
     """Implementation of CNN by Pasa and others.
 
     Simple CNN for classification based on paper by [PASA-2019]_.
@@ -192,23 +193,6 @@ class Pasa(Model):
 
         # x = F.log_softmax(x, dim=1) # 0 is batch size
 
-    def set_normalizer(self, dataloader: torch.utils.data.DataLoader) -> None:
-        """Initialize the input normalizer for the current model.
-
-        Parameters
-        ----------
-        dataloader
-            A torch Dataloader from which to compute the mean and std.
-        """
-
-        from .normalizer import make_z_normalizer
-
-        logger.info(
-            f"Uninitialised {self.name} model - "
-            f"computing z-norm factors from train dataloader.",
-        )
-        self.normalizer = make_z_normalizer(dataloader)
-
     def training_step(self, batch, _):
         images = batch[0]
         labels = batch[1]["target"]

src/mednet/libs/segmentation/models/driu.py

@@ -9,11 +9,11 @@ import typing
 import torch
 import torch.nn
 from mednet.libs.common.data.typing import TransformSequence
-from mednet.libs.common.models.model import Model
 
 from .backbones.vgg import vgg16_for_segmentation
 from .losses import SoftJaccardBCELogitsLoss
 from .make_layers import UpsampleCropBlock, conv_with_kaiming_uniform
+from .segmentation_model import SegmentationModel
 
 logger = logging.getLogger(__name__)
 
@@ -70,7 +70,7 @@ class DRIUHead(torch.nn.Module):
         return self.concatfuse(conv1_2_16, upsample2, upsample4, upsample8)
 
 
-class DRIU(Model):
+class DRIU(SegmentationModel):
     """Implementation of the DRIU model.
 
     Parameters
@@ -158,13 +158,7 @@ class DRIU(Model):
             )
             self.normalizer = make_imagenet_normalizer()
         else:
-            from .normalizer import make_z_normalizer
-
-            logger.info(
-                f"Uninitialised {self.name} model - "
-                f"computing z-norm factors from train dataloader.",
-            )
-            self.normalizer = make_z_normalizer(dataloader)
+            super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
         images = batch[0]["image"]

src/mednet/libs/segmentation/models/driu_bn.py

@@ -9,11 +9,11 @@ import typing
 import torch
 import torch.nn
 from mednet.libs.common.data.typing import TransformSequence
-from mednet.libs.common.models.model import Model
 
 from .backbones.vgg import vgg16_for_segmentation
 from .losses import SoftJaccardBCELogitsLoss
 from .make_layers import UpsampleCropBlock, conv_with_kaiming_uniform
+from .segmentation_model import SegmentationModel
 
 logger = logging.getLogger(__name__)
 
@@ -73,7 +73,7 @@ class DRIUBNHead(torch.nn.Module):
         return self.concatfuse(conv1_2_16, upsample2, upsample4, upsample8)
 
 
-class DRIUBN(Model):
+class DRIUBN(SegmentationModel):
     """Implementation of the DRIU-BN model.
 
     Parameters
@@ -161,13 +161,7 @@ class DRIUBN(Model):
             )
             self.normalizer = make_imagenet_normalizer()
         else:
-            from .normalizer import make_z_normalizer
-
-            logger.info(
-                f"Uninitialised {self.name} model - "
-                f"computing z-norm factors from train dataloader.",
-            )
-            self.normalizer = make_z_normalizer(dataloader)
+            super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
         images = batch[0]["image"]

src/mednet/libs/segmentation/models/driu_od.py

@@ -9,12 +9,12 @@ import typing
 import torch
 import torch.nn
 from mednet.libs.common.data.typing import TransformSequence
-from mednet.libs.common.models.model import Model
 
 from .backbones.vgg import vgg16_for_segmentation
 from .driu import ConcatFuseBlock
 from .losses import SoftJaccardBCELogitsLoss
 from .make_layers import UpsampleCropBlock
+from .segmentation_model import SegmentationModel
 
 logger = logging.getLogger(__name__)
 
@@ -55,7 +55,7 @@ class DRIUODHead(torch.nn.Module):
         return self.concatfuse(upsample2, upsample4, upsample8, upsample16)
 
 
-class DRIUOD(Model):
+class DRIUOD(SegmentationModel):
     """Implementation of the DRIU-OD model.
 
     Parameters
@@ -143,13 +143,7 @@ class DRIUOD(Model):
             )
             self.normalizer = make_imagenet_normalizer()
         else:
-            from .normalizer import make_z_normalizer
-
-            logger.info(
-                f"Uninitialised {self.name} model - "
-                f"computing z-norm factors from train dataloader.",
-            )
-            self.normalizer = make_z_normalizer(dataloader)
+            super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
         images = batch[0]["image"]

src/mednet/libs/segmentation/models/driu_pix.py

@@ -9,12 +9,12 @@ import typing
 import torch
 import torch.nn
 from mednet.libs.common.data.typing import TransformSequence
-from mednet.libs.common.models.model import Model
 
 from .backbones.vgg import vgg16_for_segmentation
 from .driu import ConcatFuseBlock
 from .losses import SoftJaccardBCELogitsLoss
 from .make_layers import UpsampleCropBlock
+from .segmentation_model import SegmentationModel
 
 logger = logging.getLogger(__name__)
 
@@ -59,7 +59,7 @@ class DRIUPIXHead(torch.nn.Module):
         return self.concatfuse(conv1_2_16, upsample2, upsample4, upsample8)
 
 
-class DRIUPix(Model):
+class DRIUPix(SegmentationModel):
     """Implementation of the DRIU-BN model.
 
     Parameters
@@ -147,13 +147,7 @@ class DRIUPix(Model):
             )
             self.normalizer = make_imagenet_normalizer()
         else:
-            from .normalizer import make_z_normalizer
-
-            logger.info(
-                f"Uninitialised {self.name} model - "
-                f"computing z-norm factors from train dataloader.",
-            )
-            self.normalizer = make_z_normalizer(dataloader)
+            super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
         images = batch[0]["image"]

src/mednet/libs/segmentation/models/hed.py

@@ -8,11 +8,11 @@ import typing
 import torch
 import torch.nn
 from mednet.libs.common.data.typing import TransformSequence
-from mednet.libs.common.models.model import Model
 
 from .backbones.vgg import vgg16_for_segmentation
 from .losses import MultiSoftJaccardBCELogitsLoss
 from .make_layers import UpsampleCropBlock, conv_with_kaiming_uniform
+from .segmentation_model import SegmentationModel
 
 logger = logging.getLogger(__name__)
 
@@ -73,7 +73,7 @@ class HEDHead(torch.nn.Module):
         return (upsample2, upsample4, upsample8, upsample16, concatfuse)
 
 
-class HED(Model):
+class HED(SegmentationModel):
     """Implementation of the HED model.
 
     Parameters
@@ -162,13 +162,7 @@ class HED(Model):
             )
             self.normalizer = make_imagenet_normalizer()
         else:
-            from .normalizer import make_z_normalizer
-
-            logger.info(
-                f"Uninitialised {self.name} model - "
-                f"computing z-norm factors from train dataloader.",
-            )
-            self.normalizer = make_z_normalizer(dataloader)
+            super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
         images = batch[0]["image"]

src/mednet/libs/segmentation/models/lwnet.py

@@ -21,9 +21,10 @@ import typing
 import torch
 import torch.nn
 from mednet.libs.common.data.typing import TransformSequence
-from mednet.libs.common.models.model import Model
 from mednet.libs.segmentation.models.losses import MultiWeightedBCELogitsLoss
 
+from .segmentation_model import SegmentationModel
+
 logger = logging.getLogger("mednet")
 
 
@@ -275,7 +276,7 @@ class LittleUNet(torch.nn.Module):
         return self.final(x)
 
 
-class LittleWNet(Model):
+class LittleWNet(SegmentationModel):
     """Little W-Net model, concatenating two Little U-Net models.
 
     Parameters
@@ -341,23 +342,6 @@ class LittleWNet(Model):
             shortcut=True,
         )
 
-    def set_normalizer(self, dataloader: torch.utils.data.DataLoader) -> None:
-        """Initialize the input normalizer for the current model.
-
-        Parameters
-        ----------
-        dataloader
-            A torch Dataloader from which to compute the mean and std.
-        """
-
-        from .normalizer import make_z_normalizer
-
-        logger.info(
-            f"Uninitialised {self.name} model - "
-            f"computing z-norm factors from train dataloader.",
-        )
-        self.normalizer = make_z_normalizer(dataloader)
-
     def forward(self, x):
         xn = self.normalizer(x)
         x1 = self.unet1(xn)

src/mednet/libs/segmentation/models/m2unet.py

@@ -8,11 +8,11 @@ import typing
 import torch
 import torch.nn
 from mednet.libs.common.data.typing import TransformSequence
-from mednet.libs.common.models.model import Model
 from mednet.libs.segmentation.models.losses import SoftJaccardBCELogitsLoss
 from torchvision.models.mobilenetv2 import InvertedResidual
 
 from .backbones.mobilenetv2 import mobilenet_v2_for_segmentation
+from .segmentation_model import SegmentationModel
 
 logger = logging.getLogger(__name__)
 
@@ -121,7 +121,7 @@ class M2UNetHead(torch.nn.Module):
         return self.decode1(decode2, x[1])  # 30, 3
 
 
-class M2UNET(Model):
+class M2UNET(SegmentationModel):
     """Implementation of the M2UNET model.
 
     Parameters
@@ -210,13 +210,7 @@ class M2UNET(Model):
             )
             self.normalizer = make_imagenet_normalizer()
         else:
-            from .normalizer import make_z_normalizer
-
-            logger.info(
-                f"Uninitialised {self.name} model - "
-                f"computing z-norm factors from train dataloader.",
-            )
-            self.normalizer = make_z_normalizer(dataloader)
+            super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
         images = batch[0]["image"]

src/mednet/libs/segmentation/models/segmentation_model.py

+# SPDX-FileCopyrightText: Copyright © 2024 Idiap Research Institute <contact@idiap.ch>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import logging
+import typing
+
+import torch
+import torch.nn
+import torch.optim.optimizer
+import torch.utils.data
+from mednet.libs.common.data.typing import TransformSequence
+from mednet.libs.common.models.model import Model
+from mednet.libs.segmentation.models.losses import MultiWeightedBCELogitsLoss
+
+logger = logging.getLogger("mednet")
+
+
+class SegmentationModel(Model):
+    """Base model for segmentation task.
+
+    Parameters
+    ----------
+    loss_type
+        The loss to be used for training and evaluation.
+
+        .. warning::
+
+           The loss should be set to always return batch averages (as opposed
+           to the batch sum), as our logging system expects it so.
+    loss_arguments
+        Arguments to the loss.
+    optimizer_type
+        The type of optimizer to use for training.
+    optimizer_arguments
+        Arguments to the optimizer after ``params``.
+    model_transforms
+        An optional sequence of torch modules containing transforms to be
+        applied on the input **before** it is fed into the network.
+    augmentation_transforms
+        An optional sequence of torch modules containing transforms to be
+        applied on the input **before** it is fed into the network.
+    num_classes
+        Number of outputs (classes) for this model.
+    """
+
+    def __init__(
+        self,
+        loss_type: torch.nn.Module = MultiWeightedBCELogitsLoss,
+        loss_arguments: dict[str, typing.Any] = {},
+        optimizer_type: type[torch.optim.Optimizer] = torch.optim.Adam,
+        optimizer_arguments: dict[str, typing.Any] = {},
+        model_transforms: TransformSequence = [],
+        augmentation_transforms: TransformSequence = [],
+        num_classes: int = 1,
+    ):
+        super().__init__(
+            loss_type,
+            loss_arguments,
+            optimizer_type,
+            optimizer_arguments,
+            model_transforms,
+            augmentation_transforms,
+            num_classes,
+        )
+
+    def set_normalizer(self, dataloader: torch.utils.data.DataLoader) -> None:
+        """Initialize the input normalizer for the current model.
+
+        Parameters
+        ----------
+        dataloader
+            A torch Dataloader from which to compute the mean and std.
+        """
+
+        from .normalizer import make_z_normalizer
+
+        logger.info(
+            f"Uninitialised {self.name} model - "
+            f"computing z-norm factors from train dataloader.",
+        )
+        self.normalizer = make_z_normalizer(dataloader)

src/mednet/libs/segmentation/models/unet.py

@@ -8,11 +8,11 @@ import typing
 
 import torch.nn
 from mednet.libs.common.data.typing import TransformSequence
-from mednet.libs.common.models.model import Model
 
 from .backbones.vgg import vgg16_for_segmentation
 from .losses import SoftJaccardBCELogitsLoss
 from .make_layers import UnetBlock, conv_with_kaiming_uniform
+from .segmentation_model import SegmentationModel
 
 logger = logging.getLogger(__name__)
 
@@ -62,7 +62,7 @@ class UNetHead(torch.nn.Module):
         return self.final(decode1)
 
 
-class Unet(Model):
+class Unet(SegmentationModel):
     """Implementation of the Unet model.
 
     Parameters
@@ -151,13 +151,7 @@ class Unet(Model):
             )
             self.normalizer = make_imagenet_normalizer()
         else:
-            from .normalizer import make_z_normalizer
-
-            logger.info(
-                f"Uninitialised {self.name} model - "
-                f"computing z-norm factors from train dataloader.",
-            )
-            self.normalizer = make_z_normalizer(dataloader)
+            super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
         images = batch[0]["image"]