careamics 0.0.10__py3-none-any.whl → 0.0.12__py3-none-any.whl

careamics/lightning/dataset_ng/lightning_modules/unet_module.py

@@ -0,0 +1,143 @@
+from typing import Any, Union
+
+import pytorch_lightning as L
+import torch
+from torch import nn
+from torchmetrics import MetricCollection
+from torchmetrics.image import PeakSignalNoiseRatio
+
+from careamics.config import algorithm_factory
+from careamics.config.algorithms import CAREAlgorithm, N2NAlgorithm, N2VAlgorithm
+from careamics.dataset_ng.dataset import ImageRegionData
+from careamics.models.unet import UNet
+from careamics.transforms import Denormalize
+from careamics.utils.logging import get_logger
+from careamics.utils.torch_utils import get_optimizer, get_scheduler
+
+logger = get_logger(__name__)
+
+
+class UnetModule(L.LightningModule):
+    """CAREamics PyTorch Lightning module for UNet based algorithms."""
+
+    def __init__(
+        self, algorithm_config: Union[CAREAlgorithm, N2VAlgorithm, N2NAlgorithm, dict]
+    ) -> None:
+        super().__init__()
+
+        if isinstance(algorithm_config, dict):
+            algorithm_config = algorithm_factory(algorithm_config)
+
+        self.config = algorithm_config
+        self.model: nn.Module = UNet(**algorithm_config.model.model_dump())
+
+        self._best_checkpoint_loaded = False
+
+        # TODO: how to support metric evaluation better
+        self.metrics = MetricCollection(PeakSignalNoiseRatio())
+
+    def forward(self, x: Any) -> Any:
+        """Default forward method."""
+        return self.model(x)
+
+    def _log_training_stats(self, loss: Any, batch_size: Any) -> None:
+        self.log(
+            "train_loss",
+            loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            batch_size=batch_size,
+        )
+
+        optimizer = self.optimizers()
+        if isinstance(optimizer, list):
+            current_lr = optimizer[0].param_groups[0]["lr"]
+        else:
+            current_lr = optimizer.param_groups[0]["lr"]
+        self.log(
+            "learning_rate",
+            current_lr,
+            on_step=False,
+            on_epoch=True,
+            logger=True,
+            batch_size=batch_size,
+        )
+
+    def _log_validation_stats(self, loss: Any, batch_size: Any) -> None:
+        self.log(
+            "val_loss",
+            loss,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            batch_size=batch_size,
+        )
+        self.log_dict(self.metrics, on_step=False, on_epoch=True, batch_size=batch_size)
+
+    def _load_best_checkpoint(self) -> None:
+        if (
+            not hasattr(self.trainer, "checkpoint_callback")
+            or self.trainer.checkpoint_callback is None
+        ):
+            logger.warning("No checkpoint callback found, cannot load best checkpoint.")
+            return
+
+        best_model_path = self.trainer.checkpoint_callback.best_model_path
+        if best_model_path and best_model_path != "":
+            logger.info(f"Loading best checkpoint from: {best_model_path}")
+            model_state = torch.load(best_model_path, weights_only=True)["state_dict"]
+            self.load_state_dict(model_state)
+        else:
+            logger.warning("No best checkpoint found.")
+
+    def predict_step(
+        self,
+        batch: Union[tuple[ImageRegionData], tuple[ImageRegionData, ImageRegionData]],
+        batch_idx: Any,
+        load_best_checkpoint=False,
+    ) -> Any:
+        """Default predict step."""
+        if self._best_checkpoint_loaded is False and load_best_checkpoint:
+            self._load_best_checkpoint()
+            self._best_checkpoint_loaded = True
+
+        x = batch[0]
+        # TODO: add TTA
+        prediction = self.model(x.data).cpu().numpy()
+
+        means = self._trainer.datamodule.stats.means
+        stds = self._trainer.datamodule.stats.stds
+        denormalize = Denormalize(
+            image_means=means,
+            image_stds=stds,
+        )
+        denormalized_output = denormalize(prediction)
+
+        output_batch = ImageRegionData(
+            data=denormalized_output,
+            source=x.source,
+            data_shape=x.data_shape,
+            dtype=x.dtype,
+            axes=x.axes,
+            region_spec=x.region_spec,
+        )
+        return output_batch
+
+    def configure_optimizers(self) -> Any:
+        """Configure optimizers."""
+        optimizer_func = get_optimizer(self.config.optimizer.name)
+        optimizer = optimizer_func(
+            self.model.parameters(), **self.config.optimizer.parameters
+        )
+
+        scheduler_func = get_scheduler(self.config.lr_scheduler.name)
+        scheduler = scheduler_func(optimizer, **self.config.lr_scheduler.parameters)
+
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": scheduler,
+            "monitor": "val_loss",  # otherwise triggers MisconfigurationException
+        }

careamics/lightning/lightning_module.py

@@ -148,6 +148,9 @@ class FCNModule(L.LightningModule):
         self.log(
             "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True
         )
+        optimizer = self.optimizers()
+        current_lr = optimizer.param_groups[0]["lr"]
+        self.log("learning_rate", current_lr, on_step=False, on_epoch=True, logger=True)
         return loss
 
     def validation_step(self, batch: Tensor, batch_idx: Any) -> None:

careamics/lvae_training/dataset/__init__.py

@@ -1,14 +1,19 @@
-from .multich_dataset import MultiChDloader
+from .config import DatasetConfig
 from .lc_dataset import LCMultiChDloader
+from .ms_dataset_ref import MultiChDloaderRef
+from .multich_dataset import MultiChDloader
+from .multicrop_dset import MultiCropDset
 from .multifile_dataset import MultiFileDset
-from .config import DatasetConfig
-from .types import DataType, DataSplitType, TilingMode
+from .types import DataSplitType, DataType, TilingMode
 
 __all__ = [
     "DatasetConfig",
     "MultiChDloader",
     "LCMultiChDloader",
     "MultiFileDset",
+    "MultiCropDset",
+    "MultiChDloaderRef",
+    "LCMultiChDloaderRef",
     "DataType",
     "DataSplitType",
     "TilingMode",

careamics/lvae_training/dataset/config.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from pydantic import BaseModel, ConfigDict
 
@@ -43,7 +43,7 @@ class DatasetConfig(BaseModel):
     image_size: tuple  # TODO: revisit, new model_config uses tuple
     """Size of one patch of data"""
 
-    grid_size: Optional[int] = None
+    grid_size: Optional[Union[int, tuple[int, int, int]]] = None
     """Frame is divided into square grids of this size. A patch centered on a grid 
     having size `image_size` is returned. Grid size not used in training,
     used only during val / test, grid size controls the overlap of the patches"""
@@ -82,7 +82,7 @@ class DatasetConfig(BaseModel):
     # TODO: why is this not used?
     enable_rotation_aug: Optional[bool] = False
 
-    max_val: Optional[float] = None
+    max_val: Optional[Union[float, tuple]] = None
     """Maximum data in the dataset. Is calculated for train split, and should be 
     externally set for val and test splits."""