Feat(Segmentation): add configs

src/careamics/config/__init__.py

@@ -15,17 +15,20 @@
     "N2NAlgorithm",
     "N2VAlgorithm",
     "PN2VAlgorithm",
+    "SegAlgorithm",
     "ShannonPatchFilterConfig",
     "UNetBasedAlgorithm",
     "UNetConfig",
     "VAEBasedAlgorithm",
     "create_advanced_care_config",
     "create_advanced_n2n_config",
     "create_advanced_n2v_config",
+    "create_advanced_seg_config",
     "create_care_config",
     "create_data_configuration",
     "create_n2n_config",
     "create_n2v_config",
+    "create_seg_config",
     "create_structn2v_config",
 ]
 
@@ -36,6 +39,7 @@
     N2NAlgorithm,
     N2VAlgorithm,
     PN2VAlgorithm,
+    SegAlgorithm,
     UNetBasedAlgorithm,
     VAEBasedAlgorithm,
 )
@@ -51,9 +55,11 @@
     create_advanced_care_config,
     create_advanced_n2n_config,
     create_advanced_n2v_config,
+    create_advanced_seg_config,
     create_care_config,
     create_n2n_config,
     create_n2v_config,
+    create_seg_config,
     create_structn2v_config,
 )
 from .factories.data_factory import create_data_configuration

src/careamics/config/algorithms/__init__.py

@@ -7,15 +7,20 @@
     "N2NAlgorithm",
     "N2VAlgorithm",
     "PN2VAlgorithm",
+    "SegAlgorithm",
     "UNetBasedAlgorithm",
     "VAEBasedAlgorithm",
 ]
 
+# --- unet based
 from .care_algorithm_config import CAREAlgorithm
 from .hdn_algorithm_config import HDNAlgorithm
 from .microsplit_algorithm_config import MicroSplitAlgorithm
 from .n2n_algorithm_config import N2NAlgorithm
 from .n2v_algorithm_config import N2VAlgorithm
 from .pn2v_algorithm_config import PN2VAlgorithm
+from .seg_unet_algorithm_config import SegAlgorithm
 from .unet_algorithm_config import UNetBasedAlgorithm
+
+# --- vae based
 from .vae_algorithm_config import VAEBasedAlgorithm

src/careamics/config/algorithms/care_algorithm_config.py

@@ -115,7 +115,7 @@ def get_algorithm_description(self) -> str:
     @classmethod
     def is_supervised(cls) -> bool:
         """
-        Return whether the algorithm is supervised.
+        Whether the algorithm is supervised.
 
         Returns
         -------

src/careamics/config/algorithms/n2n_algorithm_config.py

@@ -107,7 +107,7 @@ def get_algorithm_description(self) -> str:
     @classmethod
     def is_supervised(cls) -> bool:
         """
-        Return whether the algorithm is supervised.
+        Whether the algorithm is supervised.
 
         Returns
         -------

src/careamics/config/algorithms/n2v_algorithm_config.py

@@ -249,7 +249,7 @@ def get_algorithm_description(self) -> str:
     @classmethod
     def is_supervised(cls) -> bool:
         """
-        Return whether the algorithm is supervised.
+        Whether the algorithm is supervised.
 
         Returns
         -------

src/careamics/config/algorithms/seg_unet_algorithm_config.py

@@ -0,0 +1,166 @@
+"""Segmentation with UNet algorithm configuration."""
+
+from typing import Annotated, Literal
+
+from bioimageio.spec.generic.v0_3 import CiteEntry
+from pydantic import AfterValidator
+
+from careamics.config.algorithms.unet_algorithm_config import UNetBasedAlgorithm
+from careamics.config.architectures import UNetConfig
+from careamics.config.validators import (
+    model_without_final_activation,
+    model_without_n2v2,
+)
+
+
+def _model_with_at_least_2_classes(model: UNetConfig) -> UNetConfig:
+    """Validate that the Unet model has at least two classes.
+
+    Parameters
+    ----------
+    model : UNetConfig
+        Model to validate.
+
+    Returns
+    -------
+    UNetConfig
+        The validated model.
+
+    Raises
+    ------
+    ValueError
+        If the model has less than two classes.
+    """
+    if model.num_classes < 2:
+        raise ValueError(
+            f"U-Net model should have at least 2 classes, including background, "
+            f"got {model.num_classes} classes."
+        )
+
+    return model
+
+
+def _model_with_dependent_channels(model: UNetConfig) -> UNetConfig:
+    """Validate that the Unet model has dependent channels.
+
+    Parameters
+    ----------
+    model : UNetConfig
+        Model to validate.
+
+    Returns
+    -------
+    UNetConfig
+        The validated model.
+
+    Raises
+    ------
+    ValueError
+        If the model has independent channels.
+    """
+    if model.independent_channels:
+        raise ValueError(
+            "U-Net model should have `independent_channels` set to `False`."
+        )
+
+    return model
+
+
+class SegAlgorithm(UNetBasedAlgorithm):
+    """Configuration for segmentation algorithm."""
+
+    algorithm: Literal["seg"] = "seg"
+    """Segmentation algorithm name."""
+
+    loss: Literal["dice", "ce", "dice_ce"] = "dice"
+    """Segmentation-compatible loss function."""
+
+    model: Annotated[
+        UNetConfig,
+        AfterValidator(model_without_n2v2),
+        AfterValidator(model_without_final_activation),
+        AfterValidator(_model_with_at_least_2_classes),
+        AfterValidator(_model_with_dependent_channels),
+    ]
+    """UNet without a final activation function and without the `n2v2` modifications."""
+
+    def get_algorithm_friendly_name(self) -> str:
+        """
+        Get the friendly name of the algorithm.
+
+        Returns
+        -------
+        str
+            Friendly name.
+        """
+        return "UNet semantic segmentation"
+
+    def get_algorithm_keywords(self) -> list[str]:
+        """
+        Get algorithm keywords.
+
+        Returns
+        -------
+        list[str]
+            List of keywords.
+        """
+        keywords = [
+            "semantic segmentation",
+            "UNet",
+            "3D" if self.model.is_3D() else "2D",
+            "CAREamics",
+            "pytorch",
+        ]
+
+        return keywords
+
+    def get_algorithm_references(self) -> str:
+        """
+        Get the algorithm references.
+
+        This is used to generate the README of the BioImage Model Zoo export.
+
+        Returns
+        -------
+        str
+            Algorithm references.
+        """
+        return ""
+
+    def get_algorithm_citations(self) -> list[CiteEntry]:
+        """
+        Return a list of citation entries of the current algorithm.
+
+        This is used to generate the model description for the BioImage Model Zoo.
+
+        Returns
+        -------
+        List[CiteEntry]
+            List of citation entries.
+        """
+        return []
+
+    def get_algorithm_description(self) -> str:
+        """
+        Return a description of the algorithm.
+
+        This method is used to generate the README of the BioImage Model Zoo export.
+
+        Returns
+        -------
+        str
+            Description of the algorithm.
+        """
+        return "UNet semantic segmentation."
+
+    @classmethod
+    def is_supervised(cls) -> bool:
+        """
+        Whether the algorithm is supervised.
+
+        Returns
+        -------
+        bool
+            Whether the algorithm is supervised.
+        """
+        return True

src/careamics/config/algorithms/unet_algorithm_config.py

@@ -1,7 +1,6 @@
 """UNet-based algorithm Pydantic model."""
 
 from pprint import pformat
-from typing import Literal
 
 from pydantic import BaseModel, ConfigDict
 
@@ -19,16 +18,12 @@ class UNetBasedAlgorithm(BaseModel):
     training algorithm: which algorithm, loss function, model architecture, optimizer,
     and learning rate scheduler to use.
 
-    Currently, we only support N2V, CARE, N2N, and PN2V algorithms. In order to train
-    these algorithms, use the corresponding configuration child classes (e.g.
-    `N2VAlgorithm`) to ensure coherent parameters (e.g. specific losses).
-
 
     Attributes
     ----------
-    algorithm : {"n2v", "care", "n2n", "pn2v"}
+    algorithm : str
         Algorithm to use.
-    loss : {"n2v", "mae", "mse"}
+    loss : str
         Loss function to use.
     model : UNetConfig
         Model architecture to use.
@@ -53,10 +48,10 @@ class UNetBasedAlgorithm(BaseModel):
     )
 
     # Mandatory fields
-    algorithm: Literal["n2v", "care", "n2n", "pn2v"]
+    algorithm: str
     """Algorithm name, as defined in SupportedAlgorithm."""
 
-    loss: Literal["n2v", "mae", "mse", "pn2v"]
+    loss: str
     """Loss function to use, as defined in SupportedLoss."""
 
     model: UNetConfig

src/careamics/config/architectures/unet_config.py

@@ -9,20 +9,8 @@
 from .architecture_config import ArchitectureConfig
 
 
-# TODO tests activation <-> pydantic model, test the literals!
-# TODO annotations for the json schema?
 class UNetConfig(ArchitectureConfig):
-    """
-    Pydantic model for a N2V(2)-compatible UNet.
-
-    Attributes
-    ----------
-    depth : int
-        Depth of the model, between 1 and 10 (default 2).
-    num_channels_init : int
-        Number of filters of the first level of the network, should be even
-        and minimum 8 (default 96).
-    """
+    """Pydantic model for a N2V(2)-compatible UNet."""
 
     # pydantic model config
     model_config = ConfigDict(validate_assignment=True, extra="forbid")

src/careamics/config/configuration.py

@@ -14,6 +14,7 @@
     CAREAlgorithm,
     N2NAlgorithm,
     N2VAlgorithm,
+    SegAlgorithm,
 )
 from careamics.config.data import DataConfig
 from careamics.config.lightning.training_configuration import (
@@ -24,7 +25,9 @@
     get_model_constraints,
 )
 
-AlgorithmConfig = TypeVar("AlgorithmConfig", CAREAlgorithm, N2NAlgorithm, N2VAlgorithm)
+AlgorithmConfig = TypeVar(
+    "AlgorithmConfig", CAREAlgorithm, N2NAlgorithm, N2VAlgorithm, SegAlgorithm
+)
 
 
 class Configuration(BaseModel, Generic[AlgorithmConfig]):
@@ -317,7 +320,7 @@ def get_safe_experiment_name(self) -> str:
 
     def is_supervised(self) -> bool:
         """
-        Return whether the algorithm is supervised.
+        Whether the algorithm is supervised.
 
         This is true for CARE and N2N, and false for N2V. This is used to determine
         whether a target is required for training.

src/careamics/config/factories/__init__.py

@@ -4,10 +4,12 @@
     "create_advanced_care_config",
     "create_advanced_n2n_config",
     "create_advanced_n2v_config",
+    "create_advanced_seg_config",
     "create_care_config",
     "create_data_configuration",
     "create_n2n_config",
     "create_n2v_config",
+    "create_seg_config",
     "create_structn2v_config",
 ]
 
@@ -23,3 +25,4 @@
     create_n2v_config,
     create_structn2v_config,
 )
+from .seg_factory import create_advanced_seg_config, create_seg_config