careamics 0.1.0rc2__py3-none-any.whl → 0.1.0rc4__py3-none-any.whl

careamics/__init__.py

@@ -1,5 +1,4 @@
-"""Main module."""
-
+"""Main CAREamics module."""
 
 from importlib.metadata import PackageNotFoundError, version
 
@@ -8,7 +7,20 @@ try:
 except PackageNotFoundError:
     __version__ = "uninstalled"
 
-__all__ = ["Engine", "Configuration", "load_configuration", "save_configuration"]
+__all__ = [
+    "CAREamist",
+    "CAREamicsModuleWrapper",
+    "CAREamicsPredictData",
+    "CAREamicsTrainData",
+    "Configuration",
+    "load_configuration",
+    "save_configuration",
+    "TrainingDataWrapper",
+    "PredictDataWrapper",
+]
 
+from .careamist import CAREamist
 from .config import Configuration, load_configuration, save_configuration
-from .engine import Engine as Engine
+from .lightning_datamodule import CAREamicsTrainData, TrainingDataWrapper
+from .lightning_module import CAREamicsModuleWrapper
+from .lightning_prediction_datamodule import CAREamicsPredictData, PredictDataWrapper

careamics/callbacks/__init__.py

@@ -0,0 +1,6 @@
+"""Callbacks module."""
+
+__all__ = ["HyperParametersCallback", "ProgressBarCallback"]
+
+from .hyperparameters_callback import HyperParametersCallback
+from .progress_bar_callback import ProgressBarCallback

careamics/callbacks/hyperparameters_callback.py

@@ -0,0 +1,42 @@
+from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.callbacks import Callback
+
+from careamics.config import Configuration
+
+
+class HyperParametersCallback(Callback):
+    """
+    Callback allowing saving CAREamics configuration as hyperparameters in the model.
+
+    This allows saving the configuration as dictionnary in the checkpoints, and
+    loading it subsequently in a CAREamist instance.
+
+    Attributes
+    ----------
+    config : Configuration
+        CAREamics configuration to be saved as hyperparameter in the model.
+    """
+
+    def __init__(self, config: Configuration):
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        config : Configuration
+            CAREamics configuration to be saved as hyperparameter in the model.
+        """
+        self.config = config
+
+    def on_train_start(self, trainer: Trainer, pl_module: LightningModule):
+        """
+        Update the hyperparameters of the model with the configuration on train start.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            PyTorch Lightning trainer.
+        pl_module : LightningModule
+            PyTorch Lightning module.
+        """
+        pl_module.hparams.update(self.config.model_dump())

careamics/callbacks/progress_bar_callback.py

@@ -0,0 +1,57 @@
+import sys
+from typing import Dict, Union
+
+from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.callbacks import TQDMProgressBar
+from tqdm import tqdm
+
+
+class ProgressBarCallback(TQDMProgressBar):
+    """Progress bar for training and validation steps."""
+
+    def init_train_tqdm(self) -> tqdm:
+        """Override this to customize the tqdm bar for training."""
+        bar = tqdm(
+            desc="Training",
+            position=(2 * self.process_position),
+            disable=self.is_disabled,
+            leave=True,
+            dynamic_ncols=True,
+            file=sys.stdout,
+            smoothing=0,
+        )
+        return bar
+
+    def init_validation_tqdm(self) -> tqdm:
+        """Override this to customize the tqdm bar for validation."""
+        # The main progress bar doesn't exist in `trainer.validate()`
+        has_main_bar = self.train_progress_bar is not None
+        bar = tqdm(
+            desc="Validating",
+            position=(2 * self.process_position + has_main_bar),
+            disable=self.is_disabled,
+            leave=False,
+            dynamic_ncols=True,
+            file=sys.stdout,
+        )
+        return bar
+
+    def init_test_tqdm(self) -> tqdm:
+        """Override this to customize the tqdm bar for testing."""
+        bar = tqdm(
+            desc="Testing",
+            position=(2 * self.process_position),
+            disable=self.is_disabled,
+            leave=True,
+            dynamic_ncols=False,
+            ncols=100,
+            file=sys.stdout,
+        )
+        return bar
+
+    def get_metrics(
+        self, trainer: Trainer, pl_module: LightningModule
+    ) -> Dict[str, Union[int, str, float, Dict[str, float]]]:
+        """Override this to customize the metrics displayed in the progress bar."""
+        pbar_metrics = trainer.progress_bar_metrics
+        return {**pbar_metrics}