careamics 0.1.0rc1__py3-none-any.whl → 0.1.0rc3__py3-none-any.whl

careamics/config/training.py

@@ -1,534 +0,0 @@
-"""Training configuration."""
-from __future__ import annotations
-
-from typing import Dict, List
-
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field,
-    FieldValidationInfo,
-    field_validator,
-    model_validator,
-)
-from torch import optim
-
-from .config_filter import remove_default_optionals
-from .torch_optim import TorchLRScheduler, TorchOptimizer, get_parameters
-
-
-class Optimizer(BaseModel):
-    """
-    Torch optimizer.
-
-    Only parameters supported by the corresponding torch optimizer will be taken
-    into account. For more details, check:
-    https://pytorch.org/docs/stable/optim.html#algorithms
-
-    Note that mandatory parameters (see the specific Optimizer signature in the
-    link above) must be provided. For example, SGD requires `lr`.
-
-    Attributes
-    ----------
-    name : TorchOptimizer
-        Name of the optimizer.
-    parameters : dict
-        Parameters of the optimizer (see torch documentation).
-    """
-
-    # Pydantic class configuration
-    model_config = ConfigDict(
-        use_enum_values=True,
-        validate_assignment=True,
-    )
-
-    # Mandatory field
-    name: TorchOptimizer
-
-    # Optional parameters
-    parameters: dict = {}
-
-    @field_validator("parameters")
-    def filter_parameters(cls, user_params: dict, values: FieldValidationInfo) -> Dict:
-        """
-        Validate optimizer parameters.
-
-        This method filters out unknown parameters, given the optimizer name.
-
-        Parameters
-        ----------
-        user_params : dict
-            Parameters passed on to the torch optimizer.
-        values : FieldValidationInfo
-            Pydantic field validation info, used to get the optimizer name.
-
-        Returns
-        -------
-        Dict
-            Filtered optimizer parameters.
-
-        Raises
-        ------
-        ValueError
-            If the optimizer name is not specified.
-        """
-        if "name" in values.data:
-            optimizer_name = values.data["name"]
-
-            # retrieve the corresponding optimizer class
-            optimizer_class = getattr(optim, optimizer_name)
-
-            # filter the user parameters according to the optimizer's signature
-            return get_parameters(optimizer_class, user_params)
-        else:
-            raise ValueError(
-                "Cannot validate optimizer parameters without `name`, check that it "
-                "has correctly been specified."
-            )
-
-    @model_validator(mode="after")
-    def sgd_lr_parameter(cls, optimizer: Optimizer) -> Optimizer:
-        """
-        Check that SGD optimizer has the mandatory `lr` parameter specified.
-
-        Parameters
-        ----------
-        optimizer : Optimizer
-            Optimizer to validate.
-
-        Returns
-        -------
-        Optimizer
-            Validated optimizer.
-
-        Raises
-        ------
-        ValueError
-            If the optimizer is SGD and the lr parameter is not specified.
-        """
-        if optimizer.name == TorchOptimizer.SGD and "lr" not in optimizer.parameters:
-            raise ValueError(
-                "SGD optimizer requires `lr` parameter, check that it has correctly "
-                "been specified in `parameters`."
-            )
-
-        return optimizer
-
-    def model_dump(
-        self, exclude_optionals: bool = True, *args: List, **kwargs: Dict
-    ) -> Dict:
-        """
-        Override model_dump method.
-
-        The purpose of this method is to ensure smooth export to yaml. It
-        includes:
-            - removing entries with None value.
-            - removing optional values if they have the default value.
-
-        Parameters
-        ----------
-        exclude_optionals : bool, optional
-            Whether to exclude optional arguments if they are default, by default True.
-        *args : List
-            Positional arguments, unused.
-        **kwargs : Dict
-            Keyword arguments, unused.
-
-        Returns
-        -------
-        dict
-            Dictionary containing the model parameters.
-        """
-        dictionary = super().model_dump(exclude_none=True)
-
-        if exclude_optionals:
-            # remove optional arguments if they are default
-            default_optionals: dict = {"parameters": {}}
-
-            remove_default_optionals(dictionary, default_optionals)
-
-        return dictionary
-
-
-class LrScheduler(BaseModel):
-    """
-    Torch learning rate scheduler.
-
-    Only parameters supported by the corresponding torch lr scheduler will be taken
-    into account. For more details, check:
-    https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
-
-    Note that mandatory parameters (see the specific LrScheduler signature in the
-    link above) must be provided. For example, StepLR requires `step_size`.
-
-    Attributes
-    ----------
-    name : TorchLRScheduler
-        Name of the learning rate scheduler.
-    parameters : dict
-        Parameters of the learning rate scheduler (see torch documentation).
-    """
-
-    # Pydantic class configuration
-    model_config = ConfigDict(
-        use_enum_values=True,
-        validate_assignment=True,
-    )
-
-    # Mandatory field
-    name: TorchLRScheduler
-
-    # Optional parameters
-    parameters: dict = {}
-
-    @field_validator("parameters")
-    def filter_parameters(cls, user_params: dict, values: FieldValidationInfo) -> Dict:
-        """
-        Validate lr scheduler parameters.
-
-        This method filters out unknown parameters, given the lr scheduler name.
-
-        Parameters
-        ----------
-        user_params : dict
-            Parameters passed on to the torch lr scheduler.
-        values : FieldValidationInfo
-            Pydantic field validation info, used to get the lr scheduler name.
-
-        Returns
-        -------
-        Dict
-            Filtered lr scheduler parameters.
-
-        Raises
-        ------
-        ValueError
-            If the lr scheduler name is not specified.
-        """
-        if "name" in values.data:
-            lr_scheduler_name = values.data["name"]
-
-            # retrieve the corresponding lr scheduler class
-            lr_scheduler_class = getattr(optim.lr_scheduler, lr_scheduler_name)
-
-            # filter the user parameters according to the lr scheduler's signature
-            return get_parameters(lr_scheduler_class, user_params)
-        else:
-            raise ValueError(
-                "Cannot validate lr scheduler parameters without `name`, check that it "
-                "has correctly been specified."
-            )
-
-    @model_validator(mode="after")
-    def step_lr_step_size_parameter(cls, lr_scheduler: LrScheduler) -> LrScheduler:
-        """
-        Check that StepLR lr scheduler has `step_size` parameter specified.
-
-        Parameters
-        ----------
-        lr_scheduler : LrScheduler
-            Lr scheduler to validate.
-
-        Returns
-        -------
-        LrScheduler
-            Validated lr scheduler.
-
-        Raises
-        ------
-        ValueError
-            If the lr scheduler is StepLR and the step_size parameter is not specified.
-        """
-        if (
-            lr_scheduler.name == TorchLRScheduler.StepLR
-            and "step_size" not in lr_scheduler.parameters
-        ):
-            raise ValueError(
-                "StepLR lr scheduler requires `step_size` parameter, check that it has "
-                "correctly been specified in `parameters`."
-            )
-
-        return lr_scheduler
-
-    def model_dump(
-        self, exclude_optionals: bool = True, *args: List, **kwargs: Dict
-    ) -> Dict:
-        """
-        Override model_dump method.
-
-        The purpose of this method is to ensure smooth export to yaml. It includes:
-            - removing entries with None value.
-            - removing optional values if they have the default value.
-
-        Parameters
-        ----------
-        exclude_optionals : bool, optional
-            Whether to exclude optional arguments if they are default, by default True.
-        *args : List
-            Positional arguments, unused.
-        **kwargs : Dict
-            Keyword arguments, unused.
-
-        Returns
-        -------
-        dict
-            Dictionary containing the model parameters.
-        """
-        dictionary = super().model_dump(exclude_none=True)
-
-        if exclude_optionals:
-            # remove optional arguments if they are default
-            default_optionals: dict = {"parameters": {}}
-            remove_default_optionals(dictionary, default_optionals)
-
-        return dictionary
-
-
-class AMP(BaseModel):
-    """
-    Automatic mixed precision (AMP) parameters.
-
-    See: https://pytorch.org/docs/stable/amp.html.
-
-    Attributes
-    ----------
-    use : bool, optional
-        Whether to use AMP or not, default False.
-    init_scale : int, optional
-        Initial scale used for loss scaling, default 1024.
-    """
-
-    model_config = ConfigDict(
-        validate_assignment=True,
-    )
-
-    use: bool = False
-
-    # TODO review init_scale and document better
-    init_scale: int = Field(default=1024, ge=512, le=65536)
-
-    @field_validator("init_scale")
-    def power_of_two(cls, scale: int) -> int:
-        """
-        Validate that init_scale is a power of two.
-
-        Parameters
-        ----------
-        scale : int
-            Initial scale used for loss scaling.
-
-        Returns
-        -------
-        int
-            Validated initial scale.
-
-        Raises
-        ------
-        ValueError
-            If the init_scale is not a power of two.
-        """
-        if not scale & (scale - 1) == 0:
-            raise ValueError(f"Init scale must be a power of two (got {scale}).")
-
-        return scale
-
-    def model_dump(
-        self, exclude_optionals: bool = True, *args: List, **kwargs: Dict
-    ) -> Dict:
-        """
-        Override model_dump method.
-
-        The purpose is to ensure export smooth import to yaml. It includes:
-            - remove entries with None value.
-            - remove optional values if they have the default value.
-
-        Parameters
-        ----------
-        exclude_optionals : bool, optional
-            Whether to exclude optional arguments if they are default, by default True.
-        *args : List
-            Positional arguments, unused.
-        **kwargs : Dict
-            Keyword arguments, unused.
-
-        Returns
-        -------
-        dict
-            Dictionary containing the model parameters.
-        """
-        dictionary = super().model_dump(exclude_none=True)
-
-        if exclude_optionals:
-            # remove optional arguments if they are default
-            defaults = {
-                "init_scale": 1024,
-            }
-
-            remove_default_optionals(dictionary, defaults)
-
-        return dictionary
-
-
-class Training(BaseModel):
-    """
-    Parameters related to the training.
-
-    Mandatory parameters are:
-        - num_epochs: number of epochs, greater than 0.
-        - patch_size: patch size, 2D or 3D, non-zero and divisible by 2.
-        - batch_size: batch size, greater than 0.
-        - optimizer: optimizer, see `Optimizer`.
-        - lr_scheduler: learning rate scheduler, see `LrScheduler`.
-        - augmentation: whether to use data augmentation or not (True or False).
-
-    The other fields are optional:
-        - use_wandb: whether to use wandb or not (default True).
-        - num_workers: number of workers (default 0).
-        - amp: automatic mixed precision parameters (disabled by default).
-
-    Attributes
-    ----------
-    num_epochs : int
-        Number of epochs, greater than 0.
-    patch_size : conlist(int, min_length=2, max_length=3)
-        Patch size, 2D or 3D, non-zero and divisible by 2.
-    batch_size : int
-        Batch size, greater than 0.
-    optimizer : Optimizer
-        Optimizer.
-    lr_scheduler : LrScheduler
-        Learning rate scheduler.
-    augmentation : bool
-        Whether to use data augmentation or not.
-    use_wandb : bool
-        Optional, whether to use wandb or not (default True).
-    num_workers : int
-        Optional, number of workers (default 0).
-    amp : AMP
-        Optional, automatic mixed precision parameters (disabled by default).
-    """
-
-    # Pydantic class configuration
-    model_config = ConfigDict(
-        use_enum_values=True,
-        validate_assignment=True,
-    )
-
-    # Mandatory fields
-    num_epochs: int
-    patch_size: List[int] = Field(..., min_length=2, max_length=3)
-    batch_size: int
-
-    optimizer: Optimizer
-    lr_scheduler: LrScheduler
-
-    augmentation: bool
-
-    # Optional fields
-    use_wandb: bool = False
-    num_workers: int = Field(default=0, ge=0)
-    amp: AMP = AMP()
-
-    @field_validator("num_epochs", "batch_size")
-    def greater_than_0(cls, val: int) -> int:
-        """
-        Validate number of epochs.
-
-        Number of epochs must be greater than 0.
-
-        Parameters
-        ----------
-        val : int
-            Number of epochs.
-
-        Returns
-        -------
-        int
-            Validated number of epochs.
-
-        Raises
-        ------
-        ValueError
-            If the number of epochs is 0.
-        """
-        if val < 1:
-            raise ValueError(f"Number of epochs must be greater than 0 (got {val}).")
-
-        return val
-
-    @field_validator("patch_size")
-    def all_elements_non_zero_divisible_by_2(cls, patch_list: List[int]) -> List[int]:
-        """
-        Validate patch size.
-
-        Patch size must be non-zero, positive and divisible by 2.
-
-        Parameters
-        ----------
-        patch_list : List[int]
-            Patch size.
-
-        Returns
-        -------
-        List[int]
-            Validated patch size.
-
-        Raises
-        ------
-        ValueError
-            If the patch size is 0.
-        ValueError
-            If the patch size is not divisible by 2.
-        """
-        for dim in patch_list:
-            if dim < 1:
-                raise ValueError(f"Patch size must be non-zero positive (got {dim}).")
-
-            if dim % 2 != 0:
-                raise ValueError(f"Patch size must be divisible by 2 (got {dim}).")
-
-        return patch_list
-
-    def model_dump(
-        self, exclude_optionals: bool = True, *args: List, **kwargs: Dict
-    ) -> Dict:
-        """
-        Override model_dump method.
-
-        The purpose is to ensure export smooth import to yaml. It includes:
-            - remove entries with None value.
-            - remove optional values if they have the default value.
-
-        Parameters
-        ----------
-        exclude_optionals : bool, optional
-            Whether to exclude optional arguments if they are default, by default True.
-        *args : List
-            Positional arguments, unused.
-        **kwargs : Dict
-            Keyword arguments, unused.
-
-        Returns
-        -------
-        dict
-            Dictionary containing the model parameters.
-        """
-        dictionary = super().model_dump(exclude_none=True)
-
-        dictionary["optimizer"] = self.optimizer.model_dump(exclude_optionals)
-        dictionary["lr_scheduler"] = self.lr_scheduler.model_dump(exclude_optionals)
-
-        if self.amp is not None:
-            dictionary["amp"] = self.amp.model_dump(exclude_optionals)
-
-        if exclude_optionals:
-            # remove optional arguments if they are default
-            defaults = {
-                "use_wandb": False,
-                "num_workers": 0,
-                "amp": AMP().model_dump(),
-            }
-
-            remove_default_optionals(dictionary, defaults)
-
-        return dictionary

careamics/dataset/dataset_utils.py

@@ -1,115 +0,0 @@
-"""Convenience methods for datasets."""
-import logging
-from pathlib import Path
-from typing import List, Union
-
-import numpy as np
-import tifffile
-
-
-def list_files(data_path: Union[str, Path], data_format: str) -> List[Path]:
-    """
-    Return a list of path to files in a directory.
-
-    Parameters
-    ----------
-    data_path : str
-        Path to the folder containing the data.
-    data_format : str
-        Extension of the files to load, without period, e.g. `tif`.
-
-    Returns
-    -------
-    List[Path]
-        List of pathlib.Path objects.
-    """
-    files = sorted(Path(data_path).rglob(f"*.{data_format}*"))
-    return files
-
-
-def _update_axes(array: np.ndarray, axes: str) -> np.ndarray:
-    """
-    Update axes of the sample to match the config axes.
-
-    This method concatenate the S and T axes.
-
-    Parameters
-    ----------
-    array : np.ndarray
-        Input array.
-    axes : str
-        Description of axes in format STCZYX.
-
-    Returns
-    -------
-    np.ndarray
-        Updated array.
-    """
-    # concatenate ST axes to N, return NCZYX
-    if ("S" in axes or "T" in axes) and array.dtype != "O":
-        new_axes_len = len(axes.replace("Z", "").replace("YX", ""))
-        # TODO test reshape as it can scramble data, moveaxis is probably better
-        array = array.reshape(-1, *array.shape[new_axes_len:]).astype(np.float32)
-
-    elif array.dtype == "O":
-        for i in range(len(array)):
-            array[i] = np.expand_dims(array[i], axis=0).astype(np.float32)
-
-    else:
-        array = np.expand_dims(array, axis=0).astype(np.float32)
-
-    return array
-
-
-def read_tiff(file_path: Path, axes: str) -> np.ndarray:
-    """
-    Read a tiff file and return a numpy array.
-
-    Parameters
-    ----------
-    file_path : Path
-        Path to a file.
-    axes : str
-        Description of axes in format STCZYX.
-
-    Returns
-    -------
-    np.ndarray
-        Resulting array.
-
-    Raises
-    ------
-    ValueError
-        If the file failed to open.
-    OSError
-        If the file failed to open.
-    ValueError
-        If the file is not a valid tiff.
-    ValueError
-        If the data dimensions are incorrect.
-    ValueError
-        If the axes length is incorrect.
-    """
-    if file_path.suffix[:4] == ".tif":
-        try:
-            sample = tifffile.imread(file_path)
-        except (ValueError, OSError) as e:
-            logging.exception(f"Exception in file {file_path}: {e}, skipping it.")
-            raise e
-    else:
-        raise ValueError(f"File {file_path} is not a valid tiff.")
-
-    sample = sample.squeeze()
-
-    if len(sample.shape) < 2 or len(sample.shape) > 4:
-        raise ValueError(
-            f"Incorrect data dimensions. Must be 2, 3 or 4 (got {sample.shape} for"
-            f"file {file_path})."
-        )
-
-    # check number of axes
-    if len(axes) != len(sample.shape):
-        raise ValueError(f"Incorrect axes length (got {axes} for file {file_path}).")
-    sample = _update_axes(sample, axes)
-
-    return sample