careamics 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

careamics/careamist.py

@@ -13,10 +13,7 @@ from pytorch_lightning.callbacks import (
 )
 from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 
-from careamics.config import (
-    Configuration,
-    load_configuration,
-)
+from careamics.config import Configuration, FCNAlgorithmConfig, load_configuration
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -25,7 +22,7 @@ from careamics.config.support import (
 )
 from careamics.dataset.dataset_utils import reshape_array
 from careamics.lightning import (
-    CAREamicsModule,
+    FCNModule,
     HyperParametersCallback,
     PredictDataModule,
     ProgressBarCallback,
@@ -51,8 +48,6 @@ class CAREamist:
     work_dir : str, optional
         Path to working directory in which to save checkpoints and logs,
         by default None.
-    experiment_name : str, by default "CAREamics"
-        Experiment name used for checkpoints.
     callbacks : list of Callback, optional
         List of callbacks to use during training and prediction, by default None.
 
@@ -78,8 +73,7 @@ class CAREamist:
     def __init__(  # numpydoc ignore=GL08
         self,
         source: Union[Path, str],
-        work_dir: Optional[str] = None,
-        experiment_name: str = "CAREamics",
+        work_dir: Optional[Union[Path, str]] = None,
         callbacks: Optional[list[Callback]] = None,
     ) -> None: ...
 
@@ -87,8 +81,7 @@ class CAREamist:
     def __init__(  # numpydoc ignore=GL08
         self,
         source: Configuration,
-        work_dir: Optional[str] = None,
-        experiment_name: str = "CAREamics",
+        work_dir: Optional[Union[Path, str]] = None,
         callbacks: Optional[list[Callback]] = None,
     ) -> None: ...
 
@@ -96,7 +89,6 @@ class CAREamist:
         self,
         source: Union[Path, str, Configuration],
         work_dir: Optional[Union[Path, str]] = None,
-        experiment_name: str = "CAREamics",
         callbacks: Optional[list[Callback]] = None,
     ) -> None:
         """
@@ -109,18 +101,13 @@ class CAREamist:
 
         If no working directory is provided, the current working directory is used.
 
-        If `source` is a checkpoint, then `experiment_name` is used to name the
-        checkpoint, and is recorded in the configuration.
-
         Parameters
         ----------
         source : pathlib.Path or str or CAREamics Configuration
             Path to a configuration file or a trained model.
-        work_dir : str, optional
+        work_dir : str or pathlib.Path, optional
             Path to working directory in which to save checkpoints and logs,
             by default None.
-        experiment_name : str, optional
-            Experiment name used for checkpoints, by default "CAREamics".
         callbacks : list of Callback, optional
             List of callbacks to use during training and prediction, by default None.
 
@@ -148,9 +135,12 @@ class CAREamist:
             self.cfg = source
 
             # instantiate model
-            self.model = CAREamicsModule(
-                algorithm_config=self.cfg.algorithm_config,
-            )
+            if isinstance(self.cfg.algorithm_config, FCNAlgorithmConfig):
+                self.model = FCNModule(
+                    algorithm_config=self.cfg.algorithm_config,
+                )
+            else:
+                raise NotImplementedError("Architecture not supported.")
 
         # path to configuration file or model
         else:
@@ -164,9 +154,12 @@ class CAREamist:
                 self.cfg = load_configuration(source)
 
                 # instantiate model
-                self.model = CAREamicsModule(
-                    algorithm_config=self.cfg.algorithm_config,
-                )
+                if isinstance(self.cfg.algorithm_config, FCNAlgorithmConfig):
+                    self.model = FCNModule(
+                        algorithm_config=self.cfg.algorithm_config,
+                    )  # type: ignore
+                else:
+                    raise NotImplementedError("Architecture not supported.")
 
             # attempt loading a pre-trained model
             else:
@@ -192,6 +185,13 @@ class CAREamist:
         # instantiate trainer
         self.trainer = Trainer(
             max_epochs=self.cfg.training_config.num_epochs,
+            precision=self.cfg.training_config.precision,
+            max_steps=self.cfg.training_config.max_steps,
+            check_val_every_n_epoch=self.cfg.training_config.check_val_every_n_epoch,
+            enable_progress_bar=self.cfg.training_config.enable_progress_bar,
+            accumulate_grad_batches=self.cfg.training_config.accumulate_grad_batches,
+            gradient_clip_val=self.cfg.training_config.gradient_clip_val,
+            gradient_clip_algorithm=self.cfg.training_config.gradient_clip_algorithm,
             callbacks=self.callbacks,
             default_root_dir=self.work_dir,
             logger=self.experiment_logger,
@@ -247,6 +247,12 @@ class CAREamist:
                 EarlyStopping(self.cfg.training_config.early_stopping_callback)
             )
 
+    def stop_training(self) -> None:
+        """Stop the training loop."""
+        # raise stop training flag
+        self.trainer.should_stop = True
+        self.trainer.limit_val_batches = 0  # skip  validation
+
     # TODO: is there are more elegant way than calling train again after _train_on_paths
     def train(
         self,
@@ -393,9 +399,14 @@ class CAREamist:
         datamodule : TrainDataModule
             Datamodule to train on.
         """
-        # record datamodule
+        # register datamodule
         self.train_datamodule = datamodule
 
+        # set defaults (in case `stop_training` was called before)
+        self.trainer.should_stop = False
+        self.trainer.limit_val_batches = 1.0  # 100%
+
+        # train
         self.trainer.fit(self.model, datamodule=datamodule)
 
     def _train_on_array(
@@ -511,7 +522,7 @@ class CAREamist:
         tile_overlap: tuple[int, ...] = (48, 48),
         axes: Optional[str] = None,
         data_type: Optional[Literal["tiff", "custom"]] = None,
-        tta_transforms: bool = True,
+        tta_transforms: bool = False,
         dataloader_params: Optional[dict] = None,
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
@@ -527,7 +538,7 @@ class CAREamist:
         tile_overlap: tuple[int, ...] = (48, 48),
         axes: Optional[str] = None,
         data_type: Optional[Literal["array"]] = None,
-        tta_transforms: bool = True,
+        tta_transforms: bool = False,
         dataloader_params: Optional[dict] = None,
     ) -> Union[list[NDArray], NDArray]: ...
 
@@ -540,7 +551,7 @@ class CAREamist:
         tile_overlap: Optional[tuple[int, ...]] = (48, 48),
         axes: Optional[str] = None,
         data_type: Optional[Literal["array", "tiff", "custom"]] = None,
-        tta_transforms: bool = True,
+        tta_transforms: bool = False,
         dataloader_params: Optional[dict] = None,
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",

careamics/cli/__init__.py

@@ -0,0 +1,5 @@
+"""
+Package containing functions called by the careamics cli.
+
+Built using third party package Typer.
+"""

careamics/cli/conf.py

@@ -0,0 +1,391 @@
+"""Configuration building convenience functions for the CAREamics CLI."""
+
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Tuple
+
+import click
+import typer
+import yaml
+from typing_extensions import Annotated
+
+from ..config import (
+    Configuration,
+    create_care_configuration,
+    create_n2n_configuration,
+    create_n2v_configuration,
+    save_configuration,
+)
+
+WORK_DIR = Path.cwd()
+
+app = typer.Typer()
+
+
+def _config_builder_exit(ctx: typer.Context, config: Configuration) -> None:
+    """
+    Function to be called at the end of a CLI configuration builder.
+
+    Saves the `config` object and performs other functionality depending on the command
+    context.
+
+    Parameters
+    ----------
+    ctx : typer.Context
+        Typer Context.
+    config : Configuration
+        CAREamics configuration.
+    """
+    conf_path = (ctx.obj.dir / ctx.obj.name).with_suffix(".yaml")
+    save_configuration(config, conf_path)
+    if ctx.obj.print:
+        print(yaml.dump(config.model_dump(), indent=2))
+
+
+@dataclass
+class ConfOptions:
+    """Data class for containing CLI `conf` command option values."""
+
+    dir: Path
+    name: str
+    force: bool
+    print: bool
+
+
+@app.callback()
+def conf_options(  # numpydoc ignore=PR01
+    ctx: typer.Context,
+    dir: Annotated[
+        Path,
+        typer.Option(
+            "--dir", "-d", exists=True, help="Directory to save the config file to."
+        ),
+    ] = WORK_DIR,
+    name: Annotated[
+        str, typer.Option("--name", "-n", help="The config file name.")
+    ] = "config",
+    force: Annotated[
+        bool,
+        typer.Option(
+            "--force", "-f", help="Whether to overwrite existing config files."
+        ),
+    ] = False,
+    print: Annotated[
+        bool,
+        typer.Option(
+            "--print",
+            "-p",
+            help="Whether to print the config file to the console.",
+        ),
+    ] = False,
+):
+    """Build and save CAREamics configuration files."""
+    # Callback is called still on --help command
+    # If a config exists it will complain that you need to use the -f flag
+    if "--help" in sys.argv:
+        return
+    conf_path = (dir / name).with_suffix(".yaml")
+    if conf_path.exists() and not force:
+        raise FileExistsError(f"To overwrite '{conf_path}' use flag --force/-f.")
+
+    ctx.obj = ConfOptions(dir, name, force, print)
+
+
+def patch_size_callback(value: Tuple[int, int, int]) -> Tuple[int, ...]:
+    """
+    Callback for --patch-size option.
+
+    Parameters
+    ----------
+    value : (int, int, int)
+        Patch size value.
+
+    Returns
+    -------
+    (int, int, int) | (int, int)
+        If the last element in `value` is -1 the tuple is reduced to the first two
+        values.
+    """
+    if value[2] == -1:
+        return value[:2]
+    return value
+
+
+# TODO: Need to decide how to parse model kwargs
+#   - Could be json style string to be loaded as dict e.g. {"depth": 3}
+#        - Cons: Annoying to type, easily have syntax errors
+#   - Could parse all unknown options as model kwargs
+#       - Cons: There could be argument name clashes
+
+
+@app.command()
+def care(  # numpydoc ignore=PR01
+    ctx: typer.Context,
+    experiment_name: Annotated[str, typer.Option(help="Name of the experiment.")],
+    axes: Annotated[str, typer.Option(help="Axes of the data (e.g. SYX).")],
+    patch_size: Annotated[
+        click.Tuple,
+        typer.Option(
+            help=(
+                "Size of the patches along the spatial dimensions (if the data "
+                "is not 3D pass the last value as -1 e.g. --patch-size 64 64 -1)."
+            ),
+            click_type=click.Tuple([int, int, int]),
+            callback=patch_size_callback,
+        ),
+    ],
+    batch_size: Annotated[int, typer.Option(help="Batch size.")],
+    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")],
+    data_type: Annotated[
+        click.Choice,
+        typer.Option(click_type=click.Choice(["tiff"]), help="Type of the data."),
+    ] = "tiff",
+    use_augmentations: Annotated[
+        bool, typer.Option(help="Whether to use augmentations.")
+    ] = True,
+    independent_channels: Annotated[
+        bool, typer.Option(help="Whether to train all channels independently.")
+    ] = False,
+    loss: Annotated[
+        click.Choice,
+        typer.Option(
+            click_type=click.Choice(["mae", "mse"]),
+            help="Loss function to use.",
+        ),
+    ] = "mae",
+    n_channels_in: Annotated[int, typer.Option(help="Number of channels in")] = 1,
+    n_channels_out: Annotated[int, typer.Option(help="Number of channels out")] = -1,
+    logger: Annotated[
+        click.Choice,
+        typer.Option(
+            click_type=click.Choice(["wandb", "tensorboard", "none"]),
+            help="Logger to use.",
+        ),
+    ] = "none",
+    # TODO: How to address model kwargs
+) -> None:
+    """
+    Create a configuration for training CARE.
+
+    If "Z" is present in `axes`, then `path_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
+    channels. Likewise, if you set the number of channels, then "C" must be present in
+    `axes`.
+
+    To set the number of output channels, use the `n_channels_out` parameter. If it is
+    not specified, it will be assumed to be equal to `n_channels_in`.
+
+    By default, all channels are trained together. To train all channels independently,
+    set `independent_channels` to True.
+
+    By setting `use_augmentations` to False, the only transformation applied will be
+    normalization.
+    """
+    config = create_care_configuration(
+        experiment_name=experiment_name,
+        data_type=data_type,
+        axes=axes,
+        patch_size=patch_size,
+        batch_size=batch_size,
+        num_epochs=num_epochs,
+        # TODO: fix choosing augmentations
+        augmentations=None if use_augmentations else [],
+        independent_channels=independent_channels,
+        loss=loss,
+        n_channels_in=n_channels_in,
+        n_channels_out=n_channels_out,
+        logger=logger,
+    )
+    _config_builder_exit(ctx, config)
+
+
+@app.command()
+def n2n(  # numpydoc ignore=PR01
+    ctx: typer.Context,
+    experiment_name: Annotated[str, typer.Option(help="Name of the experiment.")],
+    axes: Annotated[str, typer.Option(help="Axes of the data (e.g. SYX).")],
+    patch_size: Annotated[
+        click.Tuple,
+        typer.Option(
+            help=(
+                "Size of the patches along the spatial dimensions (if the data "
+                "is not 3D pass the last value as -1 e.g. --patch-size 64 64 -1)."
+            ),
+            click_type=click.Tuple([int, int, int]),
+            callback=patch_size_callback,
+        ),
+    ],
+    batch_size: Annotated[int, typer.Option(help="Batch size.")],
+    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")],
+    data_type: Annotated[
+        click.Choice,
+        typer.Option(click_type=click.Choice(["tiff"]), help="Type of the data."),
+    ] = "tiff",
+    use_augmentations: Annotated[
+        bool, typer.Option(help="Whether to use augmentations.")
+    ] = True,
+    independent_channels: Annotated[
+        bool, typer.Option(help="Whether to train all channels independently.")
+    ] = False,
+    loss: Annotated[
+        click.Choice,
+        typer.Option(
+            click_type=click.Choice(["mae", "mse"]),
+            help="Loss function to use.",
+        ),
+    ] = "mae",
+    n_channels_in: Annotated[int, typer.Option(help="Number of channels in")] = 1,
+    n_channels_out: Annotated[int, typer.Option(help="Number of channels out")] = -1,
+    logger: Annotated[
+        click.Choice,
+        typer.Option(
+            click_type=click.Choice(["wandb", "tensorboard", "none"]),
+            help="Logger to use.",
+        ),
+    ] = "none",
+    # TODO: How to address model kwargs
+) -> None:
+    """
+    Create a configuration for training Noise2Noise.
+
+    If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
+    2.
+
+    If "C" is present in `axes`, then you need to set `n_channels` to the number of
+    channels. Likewise, if you set the number of channels, then "C" must be present in
+    `axes`.
+
+    By default, all channels are trained together. To train all channels independently,
+    set `independent_channels` to True.
+
+    By setting `use_augmentations` to False, the only transformation applied will be
+    normalization.
+    """
+    config = create_n2n_configuration(
+        experiment_name=experiment_name,
+        data_type=data_type,
+        axes=axes,
+        patch_size=patch_size,
+        batch_size=batch_size,
+        num_epochs=num_epochs,
+        # TODO: fix choosing augmentations
+        augmentations=None if use_augmentations else [],
+        independent_channels=independent_channels,
+        loss=loss,
+        n_channels_in=n_channels_in,
+        n_channels_out=n_channels_out,
+        logger=logger,
+    )
+    _config_builder_exit(ctx, config)
+
+
+@app.command()
+def n2v(  # numpydoc ignore=PR01
+    ctx: typer.Context,
+    experiment_name: Annotated[str, typer.Option(help="Name of the experiment.")],
+    axes: Annotated[str, typer.Option(help="Axes of the data (e.g. SYX).")],
+    patch_size: Annotated[
+        click.Tuple,
+        typer.Option(
+            help=(
+                "Size of the patches along the spatial dimensions (if the data "
+                "is not 3D pass the last value as -1 e.g. --patch-size 64 64 -1)."
+            ),
+            click_type=click.Tuple([int, int, int]),
+            callback=patch_size_callback,
+        ),
+    ],
+    batch_size: Annotated[int, typer.Option(help="Batch size.")],
+    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")],
+    data_type: Annotated[
+        click.Choice,
+        typer.Option(click_type=click.Choice(["tiff"]), help="Type of the data."),
+    ] = "tiff",
+    use_augmentations: Annotated[
+        bool, typer.Option(help="Whether to use augmentations.")
+    ] = True,
+    independent_channels: Annotated[
+        bool, typer.Option(help="Whether to train all channels independently.")
+    ] = True,
+    use_n2v2: Annotated[bool, typer.Option(help="Whether to use N2V2")] = False,
+    n_channels: Annotated[
+        int, typer.Option(help="Number of channels (in and out)")
+    ] = 1,
+    roi_size: Annotated[int, typer.Option(help="N2V pixel manipulation area.")] = 11,
+    masked_pixel_percentage: Annotated[
+        float, typer.Option(help="Percentage of pixels masked in each patch.")
+    ] = 0.2,
+    struct_n2v_axis: Annotated[
+        click.Choice,
+        typer.Option(click_type=click.Choice(["horizontal", "vertical", "none"])),
+    ] = "none",
+    struct_n2v_span: Annotated[
+        int, typer.Option(help="Span of the structN2V mask.")
+    ] = 5,
+    logger: Annotated[
+        click.Choice,
+        typer.Option(
+            click_type=click.Choice(["wandb", "tensorboard", "none"]),
+            help="Logger to use.",
+        ),
+    ] = "none",
+    # TODO: How to address model kwargs
+) -> None:
+    """
+    Create a configuration for training Noise2Void.
+
+    N2V uses a UNet model to denoise images in a self-supervised manner. To use its
+    variants structN2V and N2V2, set the `struct_n2v_axis` and `struct_n2v_span`
+    (structN2V) parameters, or set `use_n2v2` to True (N2V2).
+
+    N2V2 modifies the UNet architecture by adding blur pool layers and removes the skip
+    connections, thus removing checkboard artefacts. StructN2V is used when vertical
+    or horizontal correlations are present in the noise; it applies an additional mask
+    to the manipulated pixel neighbors.
+
+    If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
+    2.
+
+    If "C" is present in `axes`, then you need to set `n_channels` to the number of
+    channels.
+
+    By default, all channels are trained independently. To train all channels together,
+    set `independent_channels` to False.
+
+    By setting `use_augmentations` to False, the only transformations applied will be
+    normalization and N2V manipulation.
+
+    The `roi_size` parameter specifies the size of the area around each pixel that will
+    be manipulated by N2V. The `masked_pixel_percentage` parameter specifies how many
+    pixels per patch will be manipulated.
+
+    The parameters of the UNet can be specified in the `model_kwargs` (passed as a
+    parameter-value dictionary). Note that `use_n2v2` and 'n_channels' override the
+    corresponding parameters passed in `model_kwargs`.
+
+    If you pass "horizontal" or "vertical" to `struct_n2v_axis`, then structN2V mask
+    will be applied to each manipulated pixel.
+    """
+    config = create_n2v_configuration(
+        experiment_name=experiment_name,
+        data_type=data_type,
+        axes=axes,
+        patch_size=patch_size,
+        batch_size=batch_size,
+        num_epochs=num_epochs,
+        # TODO: fix choosing augmentations
+        augmentations=None if use_augmentations else [],
+        independent_channels=independent_channels,
+        use_n2v2=use_n2v2,
+        n_channels=n_channels,
+        roi_size=roi_size,
+        masked_pixel_percentage=masked_pixel_percentage,
+        struct_n2v_axis=struct_n2v_axis,
+        struct_n2v_span=struct_n2v_span,
+        logger=logger,
+        # TODO: Model kwargs
+    )
+    _config_builder_exit(ctx, config)

careamics/cli/main.py

@@ -0,0 +1,134 @@
+"""
+Module for CLI functionality and entrypoint.
+
+Contains the CLI entrypoint, the `run` function; and first level subcommands `train`
+and `predict`. The `conf` subcommand is added through the `app.add_typer` function, and
+its implementation is contained in the conf.py file.
+"""
+
+from pathlib import Path
+from typing import Optional
+
+import typer
+from typing_extensions import Annotated
+
+from ..careamist import CAREamist
+from . import conf
+
+app = typer.Typer(
+    help="Run CAREamics algorithms from the command line, including Noise2Void "
+    "and its many variants and cousins"
+)
+app.add_typer(
+    conf.app,
+    name="conf",
+    # callback=conf.conf_options
+)
+
+
+@app.command()
+def train(  # numpydoc ignore=PR01
+    source: Annotated[
+        Path,
+        typer.Argument(
+            help="Path to a configuration file or a trained model.",
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+        ),
+    ],
+    train_source: Annotated[
+        Path,
+        typer.Option(
+            "--train-source",
+            "-ts",
+            help="Path to the training data.",
+            exists=True,
+            file_okay=True,
+            dir_okay=True,
+        ),
+    ],
+    train_target: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--train-target",
+            "-tt",
+            help="Path to train target data.",
+            exists=True,
+            file_okay=True,
+            dir_okay=True,
+        ),
+    ] = None,
+    val_source: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--val-source",
+            "-vs",
+            help="Path to validation data.",
+            exists=True,
+            file_okay=True,
+            dir_okay=True,
+        ),
+    ] = None,
+    val_target: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--val-target",
+            "-vt",
+            help="Path to validation target data.",
+            exists=True,
+            file_okay=True,
+            dir_okay=True,
+        ),
+    ] = None,
+    use_in_memory: Annotated[
+        bool,
+        typer.Option(
+            "--use-in-memory/--not-in-memory",
+            "-m/-M",
+            help="Use in memory dataset if possible.",
+        ),
+    ] = True,
+    val_percentage: Annotated[
+        float,
+        typer.Option(help="Percentage of files to use for validation."),
+    ] = 0.1,
+    val_minimum_split: Annotated[
+        int,
+        typer.Option(help="Minimum number of files to use for validation,"),
+    ] = 1,
+    work_dir: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--work-dir",
+            "-wd",
+            help=("Path to working directory in which to save checkpoints and " "logs"),
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+        ),
+    ] = None,
+):
+    """Train CAREamics models."""
+    engine = CAREamist(source=source, work_dir=work_dir)
+    engine.train(
+        train_source=train_source,
+        val_source=val_source,
+        train_target=train_target,
+        val_target=val_target,
+        use_in_memory=use_in_memory,
+        val_percentage=val_percentage,
+        val_minimum_split=val_minimum_split,
+    )
+
+
+@app.command()
+def predict():  # numpydoc ignore=PR01
+    """Create and save predictions from CAREamics models."""
+    # TODO: Need a save predict to workdir function
+    raise NotImplementedError
+
+
+def run():
+    """CLI Entry point."""
+    app()