diff --git a/src/careamics/config/__init__.py b/src/careamics/config/__init__.py index 642319980..f3541532c 100644 --- a/src/careamics/config/__init__.py +++ b/src/careamics/config/__init__.py @@ -15,6 +15,7 @@ "N2NAlgorithm", "N2VAlgorithm", "PN2VAlgorithm", + "SegAlgorithm", "ShannonPatchFilterConfig", "UNetBasedAlgorithm", "UNetConfig", @@ -22,10 +23,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", ] @@ -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 diff --git a/src/careamics/config/algorithms/__init__.py b/src/careamics/config/algorithms/__init__.py index 81d2175cc..b21a16f0b 100644 --- a/src/careamics/config/algorithms/__init__.py +++ b/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 diff --git a/src/careamics/config/algorithms/care_algorithm_config.py b/src/careamics/config/algorithms/care_algorithm_config.py index e62879371..cd672a498 100644 --- a/src/careamics/config/algorithms/care_algorithm_config.py +++ b/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 ------- diff --git a/src/careamics/config/algorithms/n2n_algorithm_config.py b/src/careamics/config/algorithms/n2n_algorithm_config.py index 599da5118..629c70378 100644 --- a/src/careamics/config/algorithms/n2n_algorithm_config.py +++ b/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 ------- diff --git a/src/careamics/config/algorithms/n2v_algorithm_config.py b/src/careamics/config/algorithms/n2v_algorithm_config.py index 625854c2f..565fd331a 100644 --- a/src/careamics/config/algorithms/n2v_algorithm_config.py +++ b/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 ------- diff --git a/src/careamics/config/algorithms/seg_unet_algorithm_config.py b/src/careamics/config/algorithms/seg_unet_algorithm_config.py new file mode 100644 index 000000000..0dd386969 --- /dev/null +++ b/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 diff --git a/src/careamics/config/algorithms/unet_algorithm_config.py b/src/careamics/config/algorithms/unet_algorithm_config.py index 1d25040fd..9a8540dc0 100644 --- a/src/careamics/config/algorithms/unet_algorithm_config.py +++ b/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 diff --git a/src/careamics/config/architectures/unet_config.py b/src/careamics/config/architectures/unet_config.py index f0bb78a8f..a8cb05fd0 100644 --- a/src/careamics/config/architectures/unet_config.py +++ b/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") diff --git a/src/careamics/config/configuration.py b/src/careamics/config/configuration.py index 8cb39024b..8c726cafc 100644 --- a/src/careamics/config/configuration.py +++ b/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. diff --git a/src/careamics/config/factories/__init__.py b/src/careamics/config/factories/__init__.py index 0343e3baa..d984ed24c 100644 --- a/src/careamics/config/factories/__init__.py +++ b/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 diff --git a/src/careamics/config/factories/algorithm_factory.py b/src/careamics/config/factories/algorithm_factory.py index 1bf70233d..ecbd37639 100644 --- a/src/careamics/config/factories/algorithm_factory.py +++ b/src/careamics/config/factories/algorithm_factory.py @@ -8,15 +8,17 @@ CAREAlgorithm, N2NAlgorithm, N2VAlgorithm, + SegAlgorithm, + # PN2VAlgorithm, # TODO not yet compatible with NG Dataset ) from careamics.config.architectures import UNetConfig from careamics.config.support.supported_architectures import SupportedArchitecture -# TODO rename so that it does not bear the same name as the module? +# TODO replace by config_discriminator.instantiate_algorithm_config ? def algorithm_factory( algorithm: dict[str, Any], -) -> Union[N2VAlgorithm, N2NAlgorithm, CAREAlgorithm]: +) -> Union[N2VAlgorithm, N2NAlgorithm, CAREAlgorithm, SegAlgorithm]: """ Create an algorithm model for training CAREamics. @@ -27,12 +29,12 @@ def algorithm_factory( Returns ------- - N2VAlgorithm or N2NAlgorithm or CAREAlgorithm + N2VAlgorithm or N2NAlgorithm or CAREAlgorithm or SegAlgorithm Algorithm model for training CAREamics. """ adapter: TypeAdapter = TypeAdapter( Annotated[ - Union[N2VAlgorithm, N2NAlgorithm, CAREAlgorithm], + Union[N2VAlgorithm, N2NAlgorithm, CAREAlgorithm, SegAlgorithm], Field(discriminator="algorithm"), ] ) @@ -41,8 +43,8 @@ def algorithm_factory( def create_algorithm_configuration( dimensions: Literal[2, 3], - algorithm: Literal["n2v", "care", "n2n"], - loss: Literal["n2v", "mae", "mse"], + algorithm: Literal["n2v", "care", "n2n", "seg"], + loss: Literal["n2v", "mae", "mse", "dice", "ce", "dice_ce"], independent_channels: bool, n_channels_in: int, n_channels_out: int, @@ -60,10 +62,11 @@ def create_algorithm_configuration( ---------- dimensions : {2, 3} Dimension of the model, either 2D or 3D. - algorithm : {"n2v", "care", "n2n"} + algorithm : {"n2v", "care", "n2n", "seg"} Algorithm to use. - loss : {"n2v", "mae", "mse"} - Loss function to use. + loss : {"n2v", "mae", "mse", "dice", "ce", "dice_ce"} + Loss function to use. Choose `n2v` for N2V, `mae` or `mse` for CARE and N2N, + and `dice`, `ce`, or `dice_ce` for segmentation. independent_channels : bool Whether to train all channels independently. n_channels_in : int diff --git a/src/careamics/config/factories/config_discriminators.py b/src/careamics/config/factories/config_discriminators.py index 5c25d00da..ab88f1cf0 100644 --- a/src/careamics/config/factories/config_discriminators.py +++ b/src/careamics/config/factories/config_discriminators.py @@ -4,7 +4,12 @@ from pydantic import Discriminator, Tag, TypeAdapter -from careamics.config.algorithms import CAREAlgorithm, N2NAlgorithm, N2VAlgorithm +from careamics.config.algorithms import ( + CAREAlgorithm, + N2NAlgorithm, + N2VAlgorithm, + SegAlgorithm, +) from careamics.config.configuration import Configuration from careamics.config.data.normalization_config import NormalizationConfig from careamics.config.n2v_configuration import N2VConfiguration @@ -59,6 +64,7 @@ def _algo_discriminator(algo: Any) -> SupportedAlgorithm | None: Annotated[N2VConfiguration, Tag(SupportedAlgorithm.N2V)], Annotated[Configuration, Tag(SupportedAlgorithm.CARE)], Annotated[Configuration, Tag(SupportedAlgorithm.N2N)], + Annotated[Configuration, Tag(SupportedAlgorithm.SEG)], ], Discriminator(_config_discriminator), ] @@ -68,6 +74,7 @@ def _algo_discriminator(algo: Any) -> SupportedAlgorithm | None: Annotated[N2VAlgorithm, Tag(SupportedAlgorithm.N2V)], Annotated[CAREAlgorithm, Tag(SupportedAlgorithm.CARE)], Annotated[N2NAlgorithm, Tag(SupportedAlgorithm.N2N)], + Annotated[SegAlgorithm, Tag(SupportedAlgorithm.SEG)], ], Discriminator(_algo_discriminator), ] diff --git a/src/careamics/config/factories/seg_factory.py b/src/careamics/config/factories/seg_factory.py new file mode 100644 index 000000000..8f1129227 --- /dev/null +++ b/src/careamics/config/factories/seg_factory.py @@ -0,0 +1,363 @@ +"""Convenience function to create UNet-based segmentation configurations.""" + +from collections.abc import Sequence +from typing import Any, Literal + +from careamics.config.seg_configuration import SegConfiguration +from careamics.utils import get_logger + +from .algorithm_factory import create_algorithm_configuration +from .data_factory import ( + SupportedPatchFilterConfig, + create_data_configuration, +) +from .factory_utils import assemble_augmentations, validate_input_channels +from .training_factory import create_training_configuration, update_trainer_params + +logging = get_logger("Segmentation factory") + + +def _get_expected_target_axes(axes: str) -> str: + """Return expected target axes from input axes. + + Parameters + ---------- + axes : str + Expected target axes. + + Returns + ------- + str + Expected target axes given inputs. + """ + return "".join([ax for ax in axes if ax != "C"]) + + +def _get_input_size( + axes: str, + channels: Sequence[int] | None, + n_channels_in: int | None, +) -> int: + """Validate channel dimensions and return model input size. + + Parameters + ---------- + axes : str + Axes of the data (e.g. YX). + channels : Sequence[int] or None + Indices of the channels to use. + n_channels_in : int or None + Number of input channels. + + Returns + ------- + int + Adjusted number of input channels. + """ + validate_input_channels( + axes=axes, channels=channels, n_channels=n_channels_in, attr_name="n_channels" + ) + + # resolve number of input channels + if n_channels_in is None and channels is None: + resolved_n_channels_in = 1 + elif n_channels_in is not None: + resolved_n_channels_in = n_channels_in + else: + assert channels is not None + resolved_n_channels_in = len(channels) + + return resolved_n_channels_in + + +def create_seg_config( + *, + experiment_name: str, + data_type: Literal["array", "tiff", "zarr", "czi", "custom"], + axes: str, + patch_size: Sequence[int], + batch_size: int, + n_classes: int, + # optional parameters + num_epochs: int = 30, + num_steps: int | None = None, + augmentations: Sequence[Literal["x_flip", "y_flip", "rotate_90"]] | None = None, + n_val_patches: int = 8, + n_channels_in: int | None = None, +) -> SegConfiguration: + """ + Create a configuration for training a UNet for semantic segmentation. + + The `axes` parameters must reflect the actual axes and axis order from the data, + and should be the same throughout all images. The accepted axes are STCZYX. If "C" + is in `axes`, then you need to set `n_channels_in` to the number of channels + expected in the input. + + By default, CAREamics will go through the entire training data once per epoch. For + large datasets, this can lead to very long epochs. To limit the number of batches + per epoch, set the `num_steps` parameter to the desired number of batches. + + If the content of your data is expected to always have the same orientation, + consider disabling certain augmentations. By default `augmentations=None` will apply + random flips along X and Y, and random 90 degrees rotations in the XY plane. To + disable augmentations, set `augmentations=[]`. + + See `create_advanced_seg_config` for more parameters. + + Parameters + ---------- + experiment_name : str + Name of the experiment. A valid experiment name is a non-empty string that only + contains letters, numbers, underscores, dashes and spaces. + data_type : Literal["array", "tiff", "zarr", "czi", "custom"] + Type of the data. + axes : str + Axes of the data (e.g. YX). + patch_size : Sequence[int] + Size of the patches along the spatial dimensions (e.g. [64, 64]). + batch_size : int + Batch size. + n_classes : int + Number of foreground segmentation classes. + num_epochs : int, default=30 + Number of epochs to train for. + num_steps : int, default=None + Number of batches in 1 epoch. + augmentations : Sequence of {"x_flip", "y_flip", "rotate_90"}, default=None + List of augmentations to apply. If `None`, all augmentations are applied. + n_val_patches : int, default=8, + The number of patches to set aside for validation during training. This + parameter will be ignored if separate validation data is specified for training. + n_channels_in : int or None, default=None + Number of input channels. + + Returns + ------- + SegConfiguration + Configuration for training a UNet for semantic segmentation. + """ + return create_advanced_seg_config(**locals()) + + +def create_advanced_seg_config( + experiment_name: str, + data_type: Literal["array", "tiff", "zarr", "czi", "custom"], + axes: str, + patch_size: Sequence[int], + batch_size: int, + n_classes: int, + # optional parameters + num_epochs: int = 30, + num_steps: int | None = None, + n_channels_in: int | None = None, + augmentations: Sequence[Literal["x_flip", "y_flip", "rotate_90"]] | None = None, + n_val_patches: int = 8, + # advanced parameters + in_memory: bool | None = None, + channels: Sequence[int] | None = None, + normalization: Literal["mean_std", "min_max", "quantile", "none"] = "mean_std", + normalization_params: dict[str, Any] | None = None, + patch_filter_config: SupportedPatchFilterConfig | None = None, + # lightning parameters + num_workers: int = -1, + loss: Literal["dice", "ce", "dice_ce"] = "dice", + trainer_params: dict | None = None, + model_params: dict | None = None, + optimizer: Literal["Adam", "Adamax", "SGD"] = "Adam", + optimizer_params: dict[str, Any] | None = None, + lr_scheduler: Literal["ReduceLROnPlateau", "StepLR"] = "ReduceLROnPlateau", + lr_scheduler_params: dict[str, Any] | None = None, + train_dataloader_params: dict[str, Any] | None = None, + val_dataloader_params: dict[str, Any] | None = None, + checkpoint_params: dict[str, Any] | None = None, + early_stopping_params: dict[str, Any] | None = None, + logger: Literal["wandb", "tensorboard", "none"] = "none", + # reproducibility + seed: int | None = None, +) -> SegConfiguration: + """ + Create a configuration for training segmentation using a UNet model. + + If "Z" is present in `axes`, then `patch_size` must be a list of length 3, otherwise + 2. + + If "C" is present in `axes`, then you need to set `n_channels_in` to the number + of input channels. + + By default, the transformations applied are a random flip along X or Y, and a random + 90 degrees rotation in the XY plane. Normalization is always applied. + + The parameters of the UNet can be specified in the `model_params` (passed as a + parameter-value dictionary). + + Parameters + ---------- + experiment_name : str + Name of the experiment. A valid experiment name is a non-empty string that only + contains letters, numbers, underscores, dashes and spaces. + data_type : Literal["array", "tiff", "zarr", "czi", "custom"] + Type of the data. + axes : str + Axes of the data (e.g. SYX). + patch_size : Sequence[int] + Size of the patches along the spatial dimensions (e.g. [64, 64]). + batch_size : int + Batch size. + n_classes : int + Number of foreground segmentation classes. + num_epochs : int, default=30 + Number of epochs to train for. If provided, this will be added to + trainer_params. + num_steps : int | None, default=None + Number of batches in 1 epoch. If provided, this will be added to trainer_params. + Translates to `limit_train_batches` in PyTorch Lightning Trainer. See relevant + documentation for more details. + n_channels_in : int | None, default=None + Number of input channels. If `channels` is specified, then the number of + channels is inferred from its length and this parameter is ignored. + augmentations : Sequence[{"x_flip", "y_flip", "rotate_90"}] | None, default=None + List of transforms to apply, either both or one of XYFlipConfig and + XYRandomRotate90Config. By default, it applies both XYFlip (on X and Y) + and XYRandomRotate90 (in XY) to the images. + n_val_patches : int, default=8, + The number of patches to set aside for validation during training. This + parameter will be ignored if separate validation data is specified for training. + in_memory : bool | None, default=None + Whether to load all data into memory. This is only supported for 'array', + 'tiff' and 'custom' data types. If `None`, defaults to `True` for 'array', + 'tiff' and `custom`, and `False` for 'zarr' and 'czi' data types. Must be `True` + for `array`. + channels : Sequence[int] | None, default=None + List of channels to use. If `None`, all channels are used. + normalization : {"mean_std", "min_max", "quantile", "none"}, default="mean_std" + Normalization strategy to use. + normalization_params : dict[str, Any] | None, default=None + Strategy-specific normalization parameters. If None, default values are used. + For "mean_std": {"input_means": [...], "input_stds": [...]} (optional) + For "min_max": {"input_mins": [...], "input_maxes": [...]} (optional) + For "quantile": {"lower_quantiles": 0.01, "upper_quantiles": 0.99} (optional) + For "none": No parameters needed. + patch_filter_config : SupportedPatchFilterConfig | None, default=None + Specify the configuration for patch filtering. Patch filtering reduces the + probability of background patches being selected during training. If `None`, + no patch filter is applied. + num_workers : int, default=-1 + Number of workers for data loading. Use `-1` to automatically choose based + on the number of available CPUs. Unless explicitly overridden in + `train_dataloader_params` and `val_dataloader_params`, this will be applied to + all dataloaders. + loss : Literal["dice", "ce", "dice_ce"], default="dice" + Loss function to use for training. + trainer_params : dict | None, default=None + Parameters for the trainer, see the relevant documentation. + model_params : dict | None, default=None + UNetModel parameters. + optimizer : Literal["Adam", "Adamax", "SGD"], default="Adam" + Optimizer to use. + optimizer_params : dict[str, Any] | None, default=None + Parameters for the optimizer, see PyTorch documentation for more details. + lr_scheduler : Literal["ReduceLROnPlateau", "StepLR"], default="ReduceLROnPlateau" + Learning rate scheduler to use. + lr_scheduler_params : dict[str, Any] | None, default=None + Parameters for the learning rate scheduler, see PyTorch documentation for more + details. + train_dataloader_params : dict[str, Any] | None, default=None + Parameters for the training dataloader, see the PyTorch docs for `DataLoader`. + If left as `None`, `{"shuffle": True}` will be used. + val_dataloader_params : dict[str, Any] | None, default=None + Parameters for the validation dataloader, see PyTorch the docs for `DataLoader`. + checkpoint_params : dict[str, Any] | None, default=None + Parameters for the checkpoint callback, see PyTorch Lightning documentation + (`ModelCheckpoint`) for the list of available parameters. + early_stopping_params : dict[str, Any] | None, default=None + Parameters for the early stopping callback, see PyTorch Lightning documentation + (`EarlyStopping`) for the list of available parameters. + logger : Literal["wandb", "tensorboard", "none"], default="none" + Logger to use. + seed : int | None, default=None + Random seed for reproducibility. + + Returns + ------- + SegConfiguration + Configuration for training a segmentation model. + """ + n_channels_in = _get_input_size( + axes=axes, + channels=channels, + n_channels_in=n_channels_in, + ) + + # normalization + norm_config = {"name": normalization, "skip_target": True} + if normalization_params is not None: + if ( + "skip_target" in normalization_params + and not normalization_params["skip_target"] + ): + logging.warning( + msg=( + "Parameter `skip_target` in `normalization_params` must be `True`. " + "Current value will be ignored." + ), + stacklevel=2, + ) + del normalization_params["skip_target"] + norm_config.update(normalization_params) + + # data + data_config = create_data_configuration( + data_type=data_type, + axes=axes, + patch_size=patch_size, + batch_size=batch_size, + target_axes=_get_expected_target_axes(axes), + augmentations=assemble_augmentations(augmentations, seed), + n_val_patches=n_val_patches, + normalization=norm_config, + patch_filter_config=patch_filter_config, + channels=channels, + in_memory=in_memory, + num_workers=num_workers, + train_dataloader_params=train_dataloader_params, + val_dataloader_params=val_dataloader_params, + seed=seed, + ) + + # algorithm + algorithm_params = create_algorithm_configuration( + dimensions=3 if data_config.is_3D() else 2, + algorithm="seg", + loss=loss, + independent_channels=False, + n_channels_in=n_channels_in, + n_channels_out=n_classes + 1, # add background channel + use_n2v2=False, + model_params=model_params, + optimizer=optimizer, + optimizer_params=optimizer_params, + lr_scheduler=lr_scheduler, + lr_scheduler_params=lr_scheduler_params, + ) + + # training + final_trainer_params = update_trainer_params( + trainer_params=trainer_params, + num_epochs=num_epochs, + num_steps=num_steps, + ) + training_params = create_training_configuration( + algorithm="seg", + trainer_params=final_trainer_params, + logger=logger, + checkpoint_params=checkpoint_params, + early_stopping_params=early_stopping_params, + monitor_metric="val_loss", + ) + + return SegConfiguration( + experiment_name=experiment_name, + algorithm_config=algorithm_params, + data_config=data_config, + training_config=training_params, + ) diff --git a/src/careamics/config/factories/training_factory.py b/src/careamics/config/factories/training_factory.py index cdc6c6e7c..e03484cab 100644 --- a/src/careamics/config/factories/training_factory.py +++ b/src/careamics/config/factories/training_factory.py @@ -9,7 +9,7 @@ def create_training_configuration( - algorithm: Literal["care", "n2n", "n2v"], + algorithm: Literal["care", "n2n", "n2v", "seg"], trainer_params: dict, logger: Literal["wandb", "tensorboard", "none"], checkpoint_params: dict[str, Any] | None = None, diff --git a/src/careamics/config/lightning/training_configuration.py b/src/careamics/config/lightning/training_configuration.py index b192c234c..f47864699 100644 --- a/src/careamics/config/lightning/training_configuration.py +++ b/src/careamics/config/lightning/training_configuration.py @@ -62,7 +62,7 @@ class SelfSupervisedCheckpointing: def default_training_dict( - algorithm: Literal["care", "n2n", "n2v"], + algorithm: Literal["care", "n2n", "n2v", "seg"], trainer_params: dict[str, Any] | None = None, logger: Literal["wandb", "tensorboard", "none"] = "none", checkpoint_params: dict[str, Any] | None = None, @@ -77,7 +77,7 @@ def default_training_dict( Parameters ---------- - algorithm : {"care", "n2n", "n2v"} + algorithm : {"care", "n2n", "n2v", "seg"} Algorithm type, used to select the default checkpointing preset. trainer_params : dict, optional Parameters for Lightning Trainer class, by default None. @@ -104,7 +104,7 @@ def default_training_dict( # select default checkpointing preset based on algorithm default_ckpt_preset = ( SupervisedCheckpointing - if algorithm == "care" + if algorithm == "care" or algorithm == "seg" # since Noise2Noise is comparing noisy pixels to other noisy pixels, it # cannot be monitored based on a metric, we use the self-supervised preset else SelfSupervisedCheckpointing diff --git a/src/careamics/config/n2v_configuration.py b/src/careamics/config/n2v_configuration.py index aba27aa55..ce7f9a886 100644 --- a/src/careamics/config/n2v_configuration.py +++ b/src/careamics/config/n2v_configuration.py @@ -15,6 +15,8 @@ class N2VConfiguration(Configuration): """N2V-specific configuration.""" algorithm_config: N2VAlgorithm + """Algorithm configuration, holding all parameters required to configure the + model.""" @model_validator(mode="after") def target_axes_must_be_none(self: Self) -> Self: diff --git a/src/careamics/config/seg_configuration.py b/src/careamics/config/seg_configuration.py new file mode 100644 index 000000000..1d69a0aab --- /dev/null +++ b/src/careamics/config/seg_configuration.py @@ -0,0 +1,39 @@ +"""Configuration for semantic segmentation.""" + +from typing import Self + +from pydantic import model_validator + +from careamics.config.algorithms import SegAlgorithm + +from .configuration import Configuration + + +class SegConfiguration(Configuration): + """Segmentation-specific configuration.""" + + algorithm_config: SegAlgorithm + """Algorithm configuration, holding all parameters required to configure the + model.""" + + @model_validator(mode="after") + def target_normalization_is_skipped(self: Self) -> Self: + """Ensure that the target are skipped in normalization calculation. + + Returns + ------- + Self + Validated configuration. + + Raises + ------ + ValueError + If `data_config.normalization.skip_target` is not `True`. + """ + norm = self.data_config.normalization + if norm.name != "none" and not norm.skip_target: + raise ValueError( + f"Normalization {norm} must have parameter `skip_target` set to `False`" + f" for segmentation tasks." + ) + return self diff --git a/src/careamics/config/support/supported_losses.py b/src/careamics/config/support/supported_losses.py index 3f297d507..0105f6da9 100644 --- a/src/careamics/config/support/supported_losses.py +++ b/src/careamics/config/support/supported_losses.py @@ -3,7 +3,6 @@ from enum import StrEnum -# TODO register loss with custom_loss decorator? class SupportedLoss(StrEnum): """Supported losses. @@ -15,12 +14,32 @@ class SupportedLoss(StrEnum): Mean Absolute Error loss. N2V : str Noise2Void loss. + PN2V : str + Probabilistic Noise2Void loss. + CE : str + Cross-Entropy loss. + DICE : str + Dice loss. + DICE_CE : str + Combined Dice and Cross-Entropy loss. + HDN : str + Hierarchical DivNoising loss. """ + # --- CARE and Noise2Noise losses MSE = "mse" MAE = "mae" + + # --- Noise2Void losses N2V = "n2v" PN2V = "pn2v" + + # --- Segmentation losses + CE = "ce" + DICE = "dice" + DICE_CE = "dice_ce" + + # --- VAE losses HDN = "hdn" MUSPLIT = "musplit" MICROSPLIT = "microsplit" @@ -28,5 +47,3 @@ class SupportedLoss(StrEnum): DENOISPLIT_MUSPLIT = ( "denoisplit_musplit" # TODO refac losses, leave only microsplit ) - # CE = "ce" - # DICE = "dice" diff --git a/src/careamics/lightning/modules/__init__.py b/src/careamics/lightning/modules/__init__.py index 133bf5d22..a623baf8d 100644 --- a/src/careamics/lightning/modules/__init__.py +++ b/src/careamics/lightning/modules/__init__.py @@ -7,11 +7,13 @@ get_module_cls, ) from .n2v_module import N2VModule +from .seg_unet_module import SegModule __all__ = [ "CAREModule", "CAREamicsModule", "N2VModule", + "SegModule", "create_module", "get_module_cls", ] diff --git a/src/careamics/lightning/modules/care_module.py b/src/careamics/lightning/modules/care_module.py index 6b0596bf6..b8557a3b0 100644 --- a/src/careamics/lightning/modules/care_module.py +++ b/src/careamics/lightning/modules/care_module.py @@ -57,8 +57,10 @@ def __init__(self, algorithm_config: CAREAlgorithm | N2NAlgorithm | dict) -> Non config = algorithm_config if not isinstance(config, (CAREAlgorithm, N2NAlgorithm)): - raise TypeError( - "algorithm_config must be a CAREAlgorithm or a N2NAlgorithm" + raise ValueError( + f"Parameter `algorithm_config` must be a CAREAlgorithm, N2NAlgorithm, " + f"or a dict that represents a valid CAREAlgorithm or N2NAlgorithm " + f"Pydantic model (got {type(config).__name__})." ) self.save_hyperparameters({"algorithm_config": config.model_dump(mode="json")}) @@ -68,6 +70,10 @@ def __init__(self, algorithm_config: CAREAlgorithm | N2NAlgorithm | dict) -> Non MSELoss() if self.config.loss == SupportedLoss.MSE else L1Loss() ) + # TODO an alternative to logging a SIPSNR per channel would be to remove the + # channel mechanism from SIPSNR, have it return a tensor of shape (n_channels,) + # and log the channels in an overload of on_validation_epoch_end, similarly to + # the segmentation module self.metrics: MetricCollection = MetricCollection( { f"SIPSNR_{i}": SIPSNR( @@ -134,7 +140,6 @@ def training_step( torch.Tensor The loss value computed for the current batch. """ - # TODO: add validation to determine if target is initialized x, target = batch[0], batch[1] prediction = self.model(x.data) diff --git a/src/careamics/lightning/modules/get_module.py b/src/careamics/lightning/modules/get_module.py index 7aec17b5d..30315fa47 100644 --- a/src/careamics/lightning/modules/get_module.py +++ b/src/careamics/lightning/modules/get_module.py @@ -1,14 +1,15 @@ """Factory functions for lightning modules.""" -from careamics.config import CAREAlgorithm, N2NAlgorithm, N2VAlgorithm +from careamics.config import CAREAlgorithm, N2NAlgorithm, N2VAlgorithm, SegAlgorithm from careamics.config.algorithms.unet_algorithm_config import UNetBasedAlgorithm from careamics.config.support import SupportedAlgorithm from .care_module import CAREModule from .n2v_module import N2VModule +from .seg_unet_module import SegModule -CAREamicsModuleCls = type[N2VModule] | type[CAREModule] -CAREamicsModule = N2VModule | CAREModule +CAREamicsModuleCls = type[N2VModule] | type[CAREModule] | type[SegModule] +CAREamicsModule = N2VModule | CAREModule | SegModule # TODO: update to accept all algorithm configs @@ -35,6 +36,8 @@ def create_module(algorithm_config: UNetBasedAlgorithm) -> CAREamicsModule: return CAREModule(algorithm_config) elif isinstance(algorithm_config, N2VAlgorithm): return N2VModule(algorithm_config) + elif isinstance(algorithm_config, SegAlgorithm): + return SegModule(algorithm_config) else: algorithm = algorithm_config.algorithm raise NotImplementedError( @@ -67,6 +70,8 @@ def get_module_cls(algorithm: SupportedAlgorithm) -> CAREamicsModuleCls: return CAREModule case SupportedAlgorithm.N2V: return N2VModule + case SupportedAlgorithm.SEG: + return SegModule case _: raise NotImplementedError( f"Support for {algorithm.value} has not been implemented yet." diff --git a/src/careamics/lightning/modules/module_utils.py b/src/careamics/lightning/modules/module_utils.py index 4d0978a29..f00496eea 100644 --- a/src/careamics/lightning/modules/module_utils.py +++ b/src/careamics/lightning/modules/module_utils.py @@ -54,7 +54,7 @@ def log_validation_stats( module: L.LightningModule, loss: Any, batch_size: int, - metrics: MetricCollection, + metrics: MetricCollection | None = None, ) -> None: """Log validation loss and metrics. @@ -66,8 +66,8 @@ def log_validation_stats( The loss value for the current validation step. batch_size : int The size of the batch used in the current validation step. - metrics : MetricCollection - The metrics collection to log. + metrics : MetricCollection or None, optional + The metrics collection to log. When None, only the loss is logged. """ module.log( "val_loss", @@ -78,7 +78,8 @@ def log_validation_stats( logger=True, batch_size=batch_size, ) - module.log_dict(metrics, on_step=False, on_epoch=True, batch_size=batch_size) + if metrics is not None: + module.log_dict(metrics, on_step=False, on_epoch=True, batch_size=batch_size) def get_optimizer(name: str) -> type[torch.optim.Optimizer]: diff --git a/src/careamics/lightning/modules/n2v_module.py b/src/careamics/lightning/modules/n2v_module.py index cddc0167a..e60fac9e2 100644 --- a/src/careamics/lightning/modules/n2v_module.py +++ b/src/careamics/lightning/modules/n2v_module.py @@ -52,8 +52,11 @@ def __init__(self, algorithm_config: N2VAlgorithm | dict[str, Any]) -> None: config = algorithm_config if not isinstance(config, N2VAlgorithm): - raise TypeError("algorithm_config must be a N2VAlgorithm") - + raise ValueError( + f"Parameter `algorithm_config` must be a N2VAlgorithm " + f"or a dict that represents a valid N2VAlgorithm Pydantic model " + f"(got {type(config).__name__})." + ) self.save_hyperparameters({"algorithm_config": config.model_dump(mode="json")}) self.config = config self.model: nn.Module = UNet(**self.config.model.model_dump()) diff --git a/src/careamics/lightning/modules/seg_unet_module.py b/src/careamics/lightning/modules/seg_unet_module.py new file mode 100644 index 000000000..d475abed5 --- /dev/null +++ b/src/careamics/lightning/modules/seg_unet_module.py @@ -0,0 +1,244 @@ +"""UNet-based segmentation Lightning Module.""" + +from typing import TYPE_CHECKING, Any + +import torch +from lightning.pytorch import LightningModule +from torch import nn +from torchmetrics import MetricCollection +from torchmetrics.segmentation import GeneralizedDiceScore + +from careamics.config import SegAlgorithm +from careamics.config.factories.algorithm_factory import algorithm_factory +from careamics.dataset import ImageRegionData +from careamics.dataset.factory import TrainValData, TrainValSplitData +from careamics.losses import get_seg_loss +from careamics.models.unet import UNet +from careamics.utils.logging import get_logger + +from .module_utils import ( + configure_optimizers, + log_training_stats, + log_validation_stats, +) + +if TYPE_CHECKING: + from careamics.lightning.data.data_module import CareamicsDataModule + +logger = get_logger(__name__) + + +class SegModule(LightningModule): + """CAREamics PyTorch Lightning module for UNet-based segmentation. + + Parameters + ---------- + algorithm_config : SegAlgorithm or dict + Configuration for the segmentation algorithm, either as a SegAlgorithm + instance or a dictionary. + """ + + def __init__(self, algorithm_config: SegAlgorithm | dict) -> None: + """Instantiate Segmentation Module. + + Parameters + ---------- + algorithm_config : SegAlgorithm or dict + Configuration for the segmentation algorithm, either as a SegAlgorithm + instance or a dictionary. + """ + super().__init__() + + if isinstance(algorithm_config, dict): + config = algorithm_factory(algorithm_config) + else: + config = algorithm_config + + if not isinstance(config, SegAlgorithm): + raise ValueError( + f"Parameter `algorithm_config` must be a SegAlgorithm " + f"or a dict that represents a valid SegAlgorithm Pydantic model " + f"(got {type(config).__name__})." + ) + + self.save_hyperparameters({"algorithm_config": config.model_dump(mode="json")}) + self.config = config + self.model: nn.Module = UNet(**self.config.model.model_dump()) + loss = self.config.loss + self.loss_func = get_seg_loss(loss) + + self.metrics: MetricCollection = MetricCollection( + GeneralizedDiceScore( + num_classes=self.config.model.num_classes, + per_class=True, + input_format="index", + ) + ) + + def on_fit_start(self) -> None: + """On fit start hook for Segmentation module. + + Check that training and validation target data have been supplied. + """ + assert self._trainer is not None + datamodule: CareamicsDataModule = self._trainer.datamodule # type: ignore[union-attr] + assert isinstance(datamodule._data, (TrainValData, TrainValSplitData)) + if datamodule._data.train_data_target is None: + raise ValueError( + "Training target data must be provided for supervised training." + ) + if ( + isinstance(datamodule._data, TrainValData) + and datamodule._data.val_data_target is None + ): + raise ValueError( + "Validation target data must be provided for supervised training." + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass. + + Parameters + ---------- + x : torch.Tensor + Input tensor. + + Returns + ------- + torch.Tensor + Model output tensor. + """ + return self.model(x) + + def training_step( + self, + batch: tuple[ImageRegionData, ImageRegionData], + batch_idx: int, + ) -> torch.Tensor: + """Training step for segmentation module. + + Parameters + ---------- + batch : (ImageRegionData, ImageRegionData) + A tuple containing the input data and the target data. + batch_idx : int + The index of the current batch in the training loop. + + Returns + ------- + torch.Tensor + The loss value computed for the current batch. + """ + x, target = batch[0], batch[1] + + prediction = self.model(x.data) + loss = self.loss_func(prediction, target.data) + + log_training_stats(self, loss, batch_size=x.data.shape[0]) + + return loss + + def validation_step( + self, + batch: tuple[ImageRegionData, ImageRegionData], + batch_idx: int, + ) -> None: + """Validation step for segmentation module. + + Parameters + ---------- + batch : (ImageRegionData, ImageRegionData) + A tuple containing the input data and the target data. + batch_idx : int + The index of the current batch in the validation loop. + """ + x, target = batch[0], batch[1] + + prediction = self.model(x.data) + val_loss = self.loss_func(prediction, target.data) + + # compute metrics on validation + # get index class, without C dimension + pred_classes = prediction.argmax(dim=1) # (B, [Z], Y, X) + + assert isinstance(target.data, torch.Tensor) + target_long = target.data.long().squeeze(1) # (B, [Z], Y, X) + self.metrics(pred_classes, target_long) + + # not passing metrics because GenerelizedDiceScore is a tensor of length + # num_classes, which cannot be reduced to a scalar if num_classes > 1 + # so we log the metrics ourselves in on_validation_epoch_end + log_validation_stats(self, val_loss, batch_size=x.data.shape[0]) + + def on_validation_epoch_end(self) -> None: + """Log per-class Dice scores at the end of each validation epoch.""" + scores = self.metrics.compute() + dice_per_class = scores["GeneralizedDiceScore"] + + if dice_per_class.ndim == 0: + self.log("val_dice", dice_per_class, prog_bar=True, logger=True) + else: + for i, score in enumerate(dice_per_class): + self.log(f"val_dice_class_{i}", score, prog_bar=True, logger=True) + self.metrics.reset() + + def predict_step( + self, + batch: tuple[ImageRegionData] | tuple[ImageRegionData, ImageRegionData], + batch_idx: int, + ) -> ImageRegionData: + """Prediction step for segmentation module. + + Parameters + ---------- + batch : ImageRegionData or (ImageRegionData, ImageRegionData) + A tuple containing the input data and optionally the target data. + batch_idx : int + The index of the current batch in the prediction loop. + + Returns + ------- + ImageRegionData + The output batch containing the predictions. + """ + x = batch[0] + # TODO: add TTA + prediction = self.model(x.data) + + # TODO in the future, we will probably want to also return probability map + # this is not currently not possible due to the prediction conversion + # restoring the original target shape + + # apply softmax to obtain class probabilities over all classes + # prediction = prediction.softmax(dim=1) + prediction = prediction.argmax(dim=1, keepdim=True).cpu().numpy() + + output_batch = ImageRegionData( + data=prediction, + source=x.source, + data_shape=x.data_shape, + dtype=x.dtype, + axes=x.axes, + target_axes=x.target_axes, + original_data_shape=x.original_data_shape, + region_spec=x.region_spec, + additional_metadata=x.additional_metadata, + ) + return output_batch + + def configure_optimizers(self) -> dict[str, Any]: # type: ignore[override] + """Configure optimizer and learning rate scheduler. + + Returns + ------- + dict[str, Any] + A dictionary containing the optimizer and learning rate scheduler. + """ + return configure_optimizers( + model=self.model, + optimizer_name=self.config.optimizer.name, + optimizer_parameters=self.config.optimizer.parameters, + lr_scheduler_name=self.config.lr_scheduler.name, + lr_scheduler_parameters=self.config.lr_scheduler.parameters, + ) diff --git a/src/careamics/losses/__init__.py b/src/careamics/losses/__init__.py index bddad7228..e8c6b3497 100644 --- a/src/careamics/losses/__init__.py +++ b/src/careamics/losses/__init__.py @@ -3,6 +3,7 @@ __all__ = [ "denoisplit_loss", "denoisplit_musplit_loss", + "get_seg_loss", "hdn_loss", "lvae_loss_factory", "musplit_loss", @@ -18,3 +19,4 @@ musplit_loss, ) from .n2v_losses import n2v_loss, pn2v_loss +from .segmentation_losses import get_seg_loss diff --git a/src/careamics/losses/segmentation_losses.py b/src/careamics/losses/segmentation_losses.py new file mode 100644 index 000000000..7d0bacba5 --- /dev/null +++ b/src/careamics/losses/segmentation_losses.py @@ -0,0 +1,283 @@ +"""Segmentation losses.""" + +from collections.abc import Callable + +import torch +import torch.nn.functional as F +from torch.nn import Module + + +def _targets_to_class_indices(targets: torch.Tensor, num_classes: int) -> torch.Tensor: + """Convert segmentation targets to class indices. + + This method removes the C dimension, casts the labels to long for the loss + calculation, and performs validation. + + Parameters + ---------- + targets : torch.Tensor + Target representing class labels with a singleton C dimension. + num_classes : int + Number of classes. + + Returns + ------- + torch.Tensor + Target as class indices tensor. + """ + if targets.shape[1] == 1: + class_indices = targets[:, 0].long() + else: + raise ValueError( + f"Target channel dimension must be of size 1 (class labels), got size " + f"{targets.shape[1]}." + ) + + if class_indices.min() < 0 or class_indices.max() >= num_classes: + raise ValueError( + f"Target class values must be in [0, {num_classes - 1}], got values in " + f"[{class_indices.min().item()}, {class_indices.max().item()}]." + ) + + return class_indices + + +def _targets_to_one_hot(targets: torch.Tensor, num_classes: int) -> torch.Tensor: + """Convert singleton-channel class labels to one-hot encoding. + + Parameters + ---------- + targets : torch.Tensor + Target representing class labels with a singleton C dimension. + num_classes : int + Number of classes. + + Returns + ------- + torch.Tensor + Target as one-hot encoded tensor. + """ + class_indices = _targets_to_class_indices(targets, num_classes) + one_hot = F.one_hot(class_indices, num_classes=num_classes).movedim(-1, 1) + + return one_hot.float() + + +class DiceLoss(Module): + """Dice loss for binary and multi-class segmentation. + + Applies softmax activation to the model logits and computes Dice coefficient per + class, then averages across classes. + + Parameters + ---------- + weight : Tensor, optional + A manual rescaling weight given to each class. + include_background : bool, default=True + Whether to include the background class (class 0) in the loss calculation. + """ + + def __init__(self, weight=None, include_background=True) -> None: + """Constructor. + + Parameters + ---------- + weight : Tensor, optional + A manual rescaling weight given to each class. + include_background : bool, default=True + Whether to include the background class (class 0) in the loss calculation. + """ + super().__init__() + self.weight = weight + self.include_background = include_background + + def forward(self, inputs, targets, smooth=1) -> torch.Tensor: + """Compute Dice loss. + + Parameters + ---------- + inputs : Tensor + Predicted logits of shape (B, C, [Z], Y, X) where C is the number of + classes, including background (C=2 for binary). + targets : Tensor + Ground truth of shape (B, 1, [Z], Y, X) with class indices. + smooth : float, default=1 + Smoothing constant to avoid division by zero. + + Returns + ------- + Tensor + Dice loss value (1 - Dice coefficient). + """ + num_classes = inputs.shape[1] + + probabilities = F.softmax(inputs, dim=1) + targets = _targets_to_one_hot(targets, num_classes).to(inputs.device) + + if not self.include_background: + probabilities = probabilities[:, 1:] + targets = targets[:, 1:] + + probabilities = probabilities.flatten(2) + targets = targets.flatten(2) + + intersection = (probabilities * targets).sum(dim=2) + union = probabilities.sum(dim=2) + targets.sum(dim=2) + dice_per_class = (2.0 * intersection + smooth) / (union + smooth) + + if self.weight is not None: + weight = self.weight + if weight.shape[0] != num_classes: + raise ValueError( + f"Class weights must have length {num_classes}, got " + f"{weight.shape[0]}." + ) + + if not self.include_background: + weight = weight[1:] + + dice_per_class = dice_per_class * weight.to(dice_per_class.device) + + return 1 - dice_per_class.mean() + + +class DiceCELoss(Module): + """Combined Dice and Cross-Entropy loss for binary and multi-class segmentation. + + Parameters + ---------- + weight : Tensor, default=None + A manual rescaling weight given to each class for both losses. + include_background : bool, default=True + Whether to include the background class in the Dice loss calculation. + ce_weight : float, default=1.0 + Weight for the cross-entropy component. + dice_weight : float, default=1.0 + Weight for the Dice loss component. + """ + + def __init__( + self, weight=None, include_background=True, ce_weight=1.0, dice_weight=1.0 + ) -> None: + """Constructor. + + Parameters + ---------- + weight : Tensor, default=None + A manual rescaling weight given to each class for both losses. + include_background : bool, default=True + Whether to include the background class in the Dice loss calculation. + ce_weight : float, default=1.0 + Weight for the cross-entropy component. + dice_weight : float, default=1.0 + Weight for the Dice loss component. + """ + super().__init__() + self.dice_loss = DiceLoss(weight=weight, include_background=include_background) + self.weight = weight + self.ce_weight = ce_weight + self.dice_weight = dice_weight + + def forward(self, inputs, targets, smooth=1) -> torch.Tensor: + """Compute combined Dice and Cross-Entropy loss. + + Parameters + ---------- + inputs : Tensor + Predicted logits of shape (B, C, [Z], Y, X) where C is the number of + classes, including background (C=2 for binary). + targets : Tensor + Ground truth of shape (B, 1, [Z], Y, X) with class indices. + smooth : float, default=1 + Smoothing constant for Dice loss. + + Returns + ------- + Tensor + Combined loss value. + """ + num_classes = inputs.shape[1] + + target_indices = _targets_to_class_indices(targets, num_classes).to( + inputs.device + ) + + # compute Dice loss + dice_loss = self.dice_loss(inputs, targets, smooth=smooth) + + # compute cross entropy + ce_loss = F.cross_entropy( + inputs, + target_indices, + weight=self.weight, + reduction="mean", + ) + + return self.ce_weight * ce_loss + self.dice_weight * dice_loss + + +class CrossEntropyLoss(Module): + """Cross-entropy loss for segmentation targets with singleton label channels. + + Parameters + ---------- + weight : Tensor, default=None + A manual rescaling weight given to each class for both losses. + """ + + def __init__(self, weight=None) -> None: + """Constructor. + + Parameters + ---------- + weight : Tensor, optional + A manual rescaling weight given to each class. + """ + super().__init__() + self.weight = weight + + def forward(self, inputs, targets) -> torch.Tensor: + """Compute cross-entropy loss from segmentation logits and targets. + + Parameters + ---------- + inputs : Tensor + Predicted logits of shape (B, C, [Z], Y, X) where C is the number of + classes, including background (C=2 for binary). + targets : Tensor + Ground truth of shape (B, 1, [Z], Y, X) with class indices. + + Returns + ------- + Tensor + Loss value. + """ + target_indices = _targets_to_class_indices(targets, inputs.shape[1]).to( + inputs.device + ) + return F.cross_entropy( + inputs, target_indices, weight=self.weight, reduction="mean" + ) + + +def get_seg_loss(loss: str) -> Callable: + """Get loss function by name. + + Parameters + ---------- + loss : str + Name of the loss function. Supported: "dice", "ce", "dice_ce". + + Returns + ------- + Callable + Corresponding loss function. + """ + if loss == "dice": + return DiceLoss() + elif loss == "ce": + return CrossEntropyLoss() + elif loss == "dice_ce": + return DiceCELoss() + else: + raise ValueError(f"Unsupported loss function: {loss}") diff --git a/tests/unit/config/algorithm/test_algorithm_configs.py b/tests/unit/config/algorithm/test_algorithm_configs.py index ae98444d9..5a3040f67 100644 --- a/tests/unit/config/algorithm/test_algorithm_configs.py +++ b/tests/unit/config/algorithm/test_algorithm_configs.py @@ -6,10 +6,13 @@ ) from tests.utils import unet_algo_dict_testing +# ------------------------ Test utilities -------------------------- + ALGORITHMS = ["care", "n2n", "n2v"] + ALGORITHMS_CLASSES = [CAREAlgorithm, N2NAlgorithm, N2VAlgorithm] -# ------------------------ Test utilities -------------------------- +# --- Unit tests def test_default_unet_algorithm_config(): @@ -29,23 +32,24 @@ def test_unet_algorithm_configs(algorithm, cfg_class): @pytest.mark.parametrize( - "algorithm, n_in, n_out", list(zip(ALGORITHMS, [1, 2, 3], [1, 2, 3], strict=True)) + "algorithm, n_in, n_out", + [ + # CARE + ("care", 1, 1), + ("care", 1, 2), + ("care", 2, 3), + # N2N + ("n2n", 1, 1), + ("n2n", 1, 2), + ("n2n", 2, 3), + # N2V, channels must be equal + ("n2v", 1, 1), + ("n2v", 2, 2), + ], ) def test_unet_algorithm_config_channels(algorithm, n_in, n_out): """Test that an algorithm config can be created for all UNet-based algorithms with - equal channels.""" - algo_config_dict = unet_algo_dict_testing( - algorithm=algorithm, n_channels_in=n_in, n_channels_out=n_out - ) - instantiate_algorithm_config(algo_config_dict) - - -@pytest.mark.parametrize( - "algorithm, n_in, n_out", list(zip(["care", "n2n"], [3, 2], [2, 3], strict=True)) -) -def test_unet_algorithm_config_diff_channels(algorithm, n_in, n_out): - """Test that an algorithm config can be created for all UNet-based algorithms with - different channels.""" + various channels.""" algo_config_dict = unet_algo_dict_testing( algorithm=algorithm, n_channels_in=n_in, n_channels_out=n_out ) diff --git a/tests/unit/config/algorithm/test_seg_unet_algorithm_config.py b/tests/unit/config/algorithm/test_seg_unet_algorithm_config.py new file mode 100644 index 000000000..69a7bee4c --- /dev/null +++ b/tests/unit/config/algorithm/test_seg_unet_algorithm_config.py @@ -0,0 +1,42 @@ +from contextlib import nullcontext as does_not_raise + +import pytest + +from careamics.config.algorithms.seg_unet_algorithm_config import ( + _model_with_at_least_2_classes, + _model_with_dependent_channels, +) +from careamics.config.architectures import UNetConfig + + +@pytest.mark.parametrize( + "num_classes, exp_error", + [(1, pytest.raises(ValueError, match="at least 2 classes")), (2, does_not_raise())], +) +def test_model_with_at_least_2_classes(num_classes, exp_error): + """Test the validation of a segmentation model output classes.""" + model = UNetConfig( + architecture="UNet", + num_classes=num_classes, + independent_channels=False, + ) + with exp_error: + _ = _model_with_at_least_2_classes(model) + + +@pytest.mark.parametrize( + "ind_channels, exp_error", + [ + (True, pytest.raises(ValueError, match="independent_channels")), + (False, does_not_raise()), + ], +) +def test_model_with_dependent_channels(ind_channels, exp_error): + """Test the validation of a segmentation model with dependent channels.""" + model = UNetConfig( + architecture="UNet", + num_classes=2, + independent_channels=ind_channels, + ) + with exp_error: + _ = _model_with_dependent_channels(model) diff --git a/tests/unit/config/data/test_data_config.py b/tests/unit/config/data/test_data_config.py index ea949b3cf..be901606d 100644 --- a/tests/unit/config/data/test_data_config.py +++ b/tests/unit/config/data/test_data_config.py @@ -120,8 +120,8 @@ def test_default_data_config(): """Test that the default DataConfig can be created.""" - ng_data_config_dict = data_config_dict_testing() - DataConfig(**ng_data_config_dict) + data_config_dict = data_config_dict_testing() + DataConfig(**data_config_dict) # -------------------------- Unit tests ---------------------------- @@ -300,11 +300,11 @@ def test_valid_channels(self, channels, axes): ), ) def test_validate_patching_mode(mode: str, patching: str, expectation): - ng_data_config_dict = data_config_dict_testing(mode=mode, patching=patching) + data_config_dict = data_config_dict_testing(mode=mode, patching=patching) with expectation: - cfg = DataConfig(**ng_data_config_dict) - assert cfg.mode == ng_data_config_dict["mode"] - assert cfg.patching.name == ng_data_config_dict["patching"]["name"] + cfg = DataConfig(**data_config_dict) + assert cfg.mode == data_config_dict["mode"] + assert cfg.patching.name == data_config_dict["patching"]["name"] @pytest.mark.parametrize( diff --git a/tests/unit/config/factories/test_config_factories.py b/tests/unit/config/factories/test_config_factories.py index cdde2caba..a24afb30f 100644 --- a/tests/unit/config/factories/test_config_factories.py +++ b/tests/unit/config/factories/test_config_factories.py @@ -1,6 +1,7 @@ -"""Shared tests between CARE, N2N and N2V factories.""" +"""Shared tests between CARE, N2N, N2V and Seg factories.""" import itertools +from collections.abc import Callable import pytest @@ -10,6 +11,7 @@ create_advanced_care_config, create_advanced_n2n_config, create_advanced_n2v_config, + create_advanced_seg_config, ) from careamics.config.support import ( SupportedData, @@ -45,6 +47,7 @@ def test_is_subdict(): create_advanced_care_config, create_advanced_n2n_config, create_advanced_n2v_config, + create_advanced_seg_config, ] DATA_TYPE = [d.value for d in SupportedData if d.value != "czi"] # separate czi tests @@ -62,6 +65,14 @@ def test_is_subdict(): PATCH_SIZE_3D = [(16, 32, 32)] +def required_parameters(factory: Callable) -> dict: + """Additional required parameters.""" + if factory is create_advanced_seg_config: + return {"n_classes": 2} + else: + return {} + + # --- Unit tests @@ -119,6 +130,9 @@ def test_orthogonal_params( logger = SupportedLogger.WANDB.value seed = 42 + # additional required parameters that is factory dependent + extra_required_params = required_parameters(factory) + config: Configuration = factory( experiment_name=exp_name, data_type=data_type, @@ -144,6 +158,7 @@ def test_orthogonal_params( checkpoint_params=checkpoint_params, logger=logger, seed=seed, + **extra_required_params, ) assert config.experiment_name == exp_name @@ -184,6 +199,9 @@ def test_orthogonal_params( @pytest.mark.parametrize("factory", FACTORIES) def test_no_augmentation(factory): """Test that no augmentation is correctly passed to the configuration.""" + # additional required parameters that is factory dependent + extra_required_params = required_parameters(factory) + config: Configuration = factory( experiment_name="test_no_aug", data_type="tiff", @@ -191,5 +209,6 @@ def test_no_augmentation(factory): patch_size=[32, 32], batch_size=8, augmentations=[], + **extra_required_params, ) assert len(config.data_config.augmentations) == 0 diff --git a/tests/unit/config/factories/test_seg_factory.py b/tests/unit/config/factories/test_seg_factory.py new file mode 100644 index 000000000..34e92a979 --- /dev/null +++ b/tests/unit/config/factories/test_seg_factory.py @@ -0,0 +1,71 @@ +import pytest + +from careamics.config.factories.seg_factory import ( + _get_expected_target_axes, + _get_input_size, + create_advanced_seg_config, +) +from careamics.config.seg_configuration import SegConfiguration + +# --- Test utilities + + +def create_configuration(**kwargs) -> SegConfiguration: + """Wrapper around `create_advanced_seg_config`.""" + min_params = { + "experiment_name": "test_seg", + "data_type": "array", + "axes": "YX", + "patch_size": (16, 16), + "batch_size": 2, + "n_classes": 1, + } + + min_params.update(**kwargs) + return create_advanced_seg_config(**min_params) + + +# --- Unit tests + + +@pytest.mark.parametrize( + "axes, exp_axes", + [("YX", "YX"), ("YXC", "YX"), ("STZYXC", "STZYX"), ("SCYX", "SYX")], +) +def test_expected_target_axes(axes, exp_axes): + """Test expected target axes.""" + target_axes = _get_expected_target_axes(axes) + assert target_axes == exp_axes + + +@pytest.mark.parametrize( + "axes, channels, n_channels_in, exp_n_channels", + [ + ("YX", None, None, 1), + ("YX", None, 1, 1), + ("CYX", [0, 1], 2, 2), + ("CYX", [0, 1], None, 2), + ("CYX", None, 2, 2), + ], +) +def test_get_input_size(axes, channels, n_channels_in, exp_n_channels): + """Test _get_input_size.""" + result = _get_input_size(axes, channels, n_channels_in) + assert result == exp_n_channels + + +class TestSegFactory: + + @pytest.mark.parametrize("n_classes", [1, 2]) + def test_n_classes_to_model_inputs(self, n_classes): + """Test that the model inputs is background + foreground classes.""" + cfg: SegConfiguration = create_configuration(n_classes=n_classes) + assert cfg.algorithm_config.model.num_classes == n_classes + 1 + + @pytest.mark.parametrize( + "norm_params", [None, {}, {"skip_target": False}, {"skip_target": True}] + ) + def test_skip_target_enforced(self, norm_params): + """Test that `skip_target` is always enforced.""" + cfg: SegConfiguration = create_configuration(normalization_params=norm_params) + assert cfg.data_config.normalization.skip_target diff --git a/tests/unit/config/lightning/lightning/__init__.py b/tests/unit/config/lightning/lightning/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/losses/__init__.py b/tests/unit/losses/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/losses/test_segmentation_losses.py b/tests/unit/losses/test_segmentation_losses.py new file mode 100644 index 000000000..f96efefee --- /dev/null +++ b/tests/unit/losses/test_segmentation_losses.py @@ -0,0 +1,167 @@ +"""Tests for segmentation losses.""" + +from contextlib import nullcontext as does_not_raise + +import pytest +import torch + +from careamics.losses.segmentation_losses import ( + CrossEntropyLoss, + DiceCELoss, + DiceLoss, + get_seg_loss, +) + +# --- Test utilities + +LOSSES = [DiceCELoss, DiceLoss, CrossEntropyLoss] + +BIN_CLASS_LABELS = torch.tensor([[0, 1], [0, 1]]) + +BIN_ONE_HOT = torch.tensor( + [ + [[1, 0], [1, 0]], + [[0, 1], [0, 1]], + ] +) + +MUL_CLASS_LABELS = torch.tensor([[0, 1], [0, 2]]) + +MUL_ONE_HOT = torch.tensor( + [ + [[1, 0], [1, 0]], + [[0, 1], [0, 0]], + [[0, 0], [0, 1]], + ] +) + + +def to_3D(tensor: torch.Tensor) -> torch.Tensor: + """Turn a 2D label map or one-hot map into a small 3D volume.""" + if tensor.ndim == 2: + return torch.stack([tensor, torch.rot90(tensor, k=1)], dim=0) + return torch.stack([tensor, torch.rot90(tensor, k=1, dims=(1, 2))], dim=1) + + +def to_batch(tensor: torch.Tensor, batch_size: int = 1) -> torch.tensor: + """Add batch dimension by repeating the same sample.""" + return tensor.unsqueeze(0).repeat(batch_size, *([1] * tensor.ndim)) + + +def make_targets(class_labels: torch.Tensor, batch_size: int = 1) -> torch.Tensor: + """Create batched targets of shape (B, 1, ...).""" + return to_batch(class_labels.unsqueeze(0), batch_size=batch_size) + + +def low_loss_logits(one_hot, logit=10.0): + """Logits corresponding to a near-perfect prediction.""" + return one_hot * logit + (1 - one_hot) * -logit + + +def high_loss_logits(one_hot, logit=10.0): + """Logits corresponding to a confidently wrong prediction.""" + return one_hot * -logit + (1 - one_hot) * logit + + +# --- Unit tests + + +@pytest.mark.parametrize( + "loss_func,class_labels,one_hot_labels,low_threshold", + [ + (DiceLoss, BIN_CLASS_LABELS, BIN_ONE_HOT, 1e-3), + (DiceLoss, MUL_CLASS_LABELS, MUL_ONE_HOT, 1e-3), + (DiceCELoss, BIN_CLASS_LABELS, BIN_ONE_HOT, 1e-2), + (DiceCELoss, MUL_CLASS_LABELS, MUL_ONE_HOT, 1e-2), + (CrossEntropyLoss, BIN_CLASS_LABELS, BIN_ONE_HOT, 1e-3), + (CrossEntropyLoss, MUL_CLASS_LABELS, MUL_ONE_HOT, 1e-3), + ], +) +@pytest.mark.parametrize("batch_size", [1, 2]) +@pytest.mark.parametrize("is_3D", [False, True]) +def test_loss_perfect_prediction_is_low( + loss_func, class_labels, one_hot_labels, low_threshold, batch_size, is_3D +): + """Perfect logits should produce a very small loss.""" + loss = loss_func() + + inputs = one_hot_labels if not is_3D else to_3D(one_hot_labels) + targets = class_labels if not is_3D else to_3D(class_labels) + + logits = to_batch(low_loss_logits(inputs), batch_size=batch_size) + targets = make_targets(targets, batch_size=batch_size) + + loss_value = loss(logits, targets) + + assert loss_value < low_threshold + + +@pytest.mark.parametrize( + "loss_func,class_labels,one_hot_labels", + [ + (DiceLoss, BIN_CLASS_LABELS, BIN_ONE_HOT), + (DiceLoss, MUL_CLASS_LABELS, MUL_ONE_HOT), + (DiceCELoss, BIN_CLASS_LABELS, BIN_ONE_HOT), + (DiceCELoss, MUL_CLASS_LABELS, MUL_ONE_HOT), + (CrossEntropyLoss, BIN_CLASS_LABELS, BIN_ONE_HOT), + (CrossEntropyLoss, MUL_CLASS_LABELS, MUL_ONE_HOT), + ], +) +def test_loss_wrong_prediction_is_higher(loss_func, class_labels, one_hot_labels): + """Wrong logits should produce a larger loss than perfect logits.""" + loss = loss_func() + low_logits = to_batch(low_loss_logits(one_hot_labels)) + high_logits = to_batch(high_loss_logits(one_hot_labels)) + targets = make_targets(class_labels) + + low_loss = loss(low_logits, targets) + high_loss = loss(high_logits, targets) + + assert low_loss < high_loss + assert high_loss - low_loss > 0.1 + + +@pytest.mark.parametrize( + "class_labels,one_hot_labels", + [ + (BIN_CLASS_LABELS, BIN_ONE_HOT), + (MUL_CLASS_LABELS, MUL_ONE_HOT), + (BIN_CLASS_LABELS, BIN_ONE_HOT), + (MUL_CLASS_LABELS, MUL_ONE_HOT), + (BIN_CLASS_LABELS, BIN_ONE_HOT), + (MUL_CLASS_LABELS, MUL_ONE_HOT), + ], +) +@pytest.mark.parametrize("batch_size", [1, 2]) +@pytest.mark.parametrize("is_3D", [False, True]) +def test_dice_ce_sum(class_labels, one_hot_labels, batch_size, is_3D): + """Test that DiceCE is the sum of Dice and CE.""" + dice_ce = DiceCELoss() + ce = CrossEntropyLoss() + dice = DiceLoss() + + inputs = one_hot_labels if not is_3D else to_3D(one_hot_labels) + targets = class_labels if not is_3D else to_3D(class_labels) + + inp = to_batch(low_loss_logits(inputs), batch_size=batch_size) + tar = make_targets(targets, batch_size=batch_size) + + assert dice_ce.forward(inp, tar) == ce.forward(inp, tar) + dice.forward(inp, tar) + + +@pytest.mark.parametrize( + "loss_name, exp_class, exp_error", + [ + # no error + ("dice", DiceLoss, does_not_raise()), + ("ce", CrossEntropyLoss, does_not_raise()), + ("dice_ce", DiceCELoss, does_not_raise()), + # error + ("not_a_loss", None, pytest.raises(ValueError, match="Unsupported")), + ], +) +def test_get_loss(loss_name, exp_class, exp_error): + """Test loss factory.""" + with exp_error: + loss = get_seg_loss(loss_name) + assert isinstance(loss, exp_class) diff --git a/tests/unit/models/constraints/test_unet_constraints.py b/tests/unit/models/constraints/test_unet_constraints.py index 97c24a7e5..bd031ac78 100644 --- a/tests/unit/models/constraints/test_unet_constraints.py +++ b/tests/unit/models/constraints/test_unet_constraints.py @@ -72,7 +72,7 @@ def _incompatible_shapes(depth: int): ) ), ) -def test_validate_input_shape(x_shape, z_shape, depth, expected_error): +def test_validate_spatial_shape(x_shape, z_shape, depth, expected_error): cfg = UNetConfig( architecture="UNet", depth=depth, @@ -86,7 +86,7 @@ def test_validate_input_shape(x_shape, z_shape, depth, expected_error): @pytest.mark.parametrize("length", [1, 4]) -def test_validate_input_shape_wrong_length(length): +def test_validate_spatial_shape_wrong_length(length): depth = 2 shape = (_compatible_shapes(depth=depth)[0],) * length