Lightning

Instantiate CARE Module.

Parameters:

• algorithm_config (CAREAlgorithm, N2NAlgorithm, or dict) –

  Configuration for the CARE algorithm, either as a CAREAlgorithm/N2NAlgorithm instance or a dictionary.

Configure optimizer and learning rate scheduler.

Returns:

• dict[str, Any] –

  A dictionary containing the optimizer and learning rate scheduler.

Forward pass.

Parameters:

• x (Tensor) –

  Input tensor.

Returns:

• Tensor –

  Model output tensor.

On fit start hook for CARE module.

Check that training and validation target data have been supplied.

Prediction step for CARE 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.

Training step for CARE 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:

• Tensor –

  The loss value computed for the current batch.

Validation step for CARE 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.

Data module for Careamics dataset initialization.

Create a lightning datamodule that handles creating datasets for training, validation, and prediction.

Parameters:

• data_config (DataConfig) –

  Pydantic model for CAREamics data configuration.

• train_data (Any, default: None ) –

  Training data. If custom loading is provided it can be any type, otherwise it must be a pathlib.Path, str, numpy.ndarray or a sequence of these, or None.

• train_data_target (Any, default: None ) –

  Training data target. If custom loading is provided it can be any type, otherwise it must be a pathlib.Path, str, numpy.ndarray or a sequence of these, or None.

• train_data_mask (Any, default: None ) –

  Training data mask, an optional mask that can be provided to filter regions of the data during training, such as large areas of background. The mask should be a binary image where a 1 indicates a pixel should be included in the training data. If custom loading is provided it can be any type, otherwise it must be a pathlib.Path, str, numpy.ndarray or a sequence of these, or None.

• val_data (Any, default: None ) –

  Validation data. If not provided, data_config.n_val_patches patches will selected from the training data for validation. If custom loading is provided it can be any type, otherwise it must be a pathlib.Path, str, numpy.ndarray or a sequence of these, or None.

• val_data_target (Any, default: None ) –

  Validation data target. If custom loading is provided it can be any type, otherwise it must be a pathlib.Path, str, numpy.ndarray or a sequence of these, or None.

• pred_data (Any, default: None ) –

  Prediction data. If custom loading is provided it can be any type, otherwise it must be a pathlib.Path, str, numpy.ndarray or a sequence of these, or None.

• pred_data_target (Any, default: None ) –

  Prediction data target, this may be used for calculating metrics. If custom loading is provided it can be any type, otherwise it must be a pathlib.Path, str, numpy.ndarray or a sequence of these, or None.

• model_constraints (ModelConstraints, default: None ) –

  If provided, the data module will validate input and target channels and spatial shapes against the model constraints.

• loading (ReadFuncLoading | ImageStackLoading | None, default: None ) –

  The type of loading used for custom data. ReadFuncLoading is the use of a simple function that will load full images into memory. ImageStackLoading is for custom chunked or memory-mapped next-generation file formats enabling single patches to be read from disk at a time. If the data type is not custom loading should be None.

Create a dataloader for prediction.

Returns:

• DataLoader –

  Prediction dataloader.

Setup datasets.

Lightning hook that is called at the beginning of fit (train + validate), validate, test, or predict. Creates the datasets for a given stage.

Parameters:

• stage (str) –

  The stage to set up datasets for. Is either 'fit', 'validate', 'test', or 'predict'.

Raises:

• NotImplementedError –

  If stage is not one of "fit", "validate" or "predict".

Create a dataloader for training.

Returns:

• DataLoader –

  Training dataloader.

Create a dataloader for validation.

Returns:

• DataLoader –

  Validation dataloader.

Initialize the callback.

Parameters:

• careamics_version (str) –

  Version of CAREamics used for training.

• experiment_name (str) –

  Name of the experiment.

• training_config (TrainingConfig) –

  Training configuration to store in checkpoint.

Lightning hook called when saving a checkpoint.

Adds CAREamics configuration to the checkpoint dictionary.

Parameters:

• trainer (Trainer) –

  Lightning trainer instance.

• pl_module (LightningModule) –

  Lightning module being trained.

• checkpoint (dict[str, Any]) –

  Checkpoint dictionary to modify.

Called when trainer is setting up.

Parameters:

• trainer (Trainer) –

  PyTorch Lightning trainer.

• module (LightningModule) –

  Lightning module.

• stage (str) –

  Current stage (fit, validate, test, or predict).

Initialize MicroSplitDataModule.

Parameters:

• data_config (MicroSplitDataConfig) –

  Configuration for the MicroSplit dataset.

• train_data (str) –

  Path to training data directory.

• val_data (str, default: None ) –

  Path to validation data directory.

• train_data_target (str, default: None ) –

  Path to training target data.

• val_data_target (str, default: None ) –

  Path to validation target data.

• read_source_func (Callable, default: None ) –

  Function to read source data.

• extension_filter (str, default: '' ) –

  File extension filter.

• val_percentage (float, default: 0.1 ) –

  Percentage of data to use for validation, by default 0.1.

• val_minimum_split (int, default: 5 ) –

  Minimum number of samples for validation split, by default 5.

• use_in_memory (bool, default: True ) –

  Whether to use in-memory dataset, by default True.

Get data statistics.

Returns:

• tuple[dict, dict] –

  A tuple containing two dictionaries: - data_mean: mean values for input and target - data_std: standard deviation values for input and target

Create a dataloader for training.

Returns:

• DataLoader –

  Training dataloader.

Create a dataloader for validation.

Returns:

• DataLoader –

  Validation dataloader.

Instantiate N2VModule.

Parameters:

• algorithm_config (N2VAlgorithm or dict) –

  Configuration for the N2V algorithm, either as an N2VAlgorithm instance or a dictionary.

Configure optimizer and learning rate scheduler.

Returns:

• dict[str, Any] –

  A dictionary containing the optimizer and learning rate scheduler.

Forward pass.

Parameters:

• x (Tensor) –

  Input tensor.

Returns:

• Tensor –

  Model output tensor.

On fit start hook for N2V module.

Prediction step for N2V model.

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.

Training step for N2V model.

Parameters:

• batch (ImageRegionData or (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:

• Tensor –

  The loss value for the current training step.

Validation step for N2V model.

Parameters:

• batch (ImageRegionData or (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.

Override this to customize the metrics displayed in the progress bar.

Parameters:

• trainer (Trainer) –

  The trainer object.

• pl_module (LightningModule) –

  The LightningModule object, unused.

Returns:

• dict –

  A dictionary with the metrics to display in the progress bar.

Override this to customize the tqdm bar for testing.

Returns:

• tqdm –

  A tqdm bar.

Override this to customize the tqdm bar for training.

Returns:

• tqdm –

  A tqdm bar.

Override this to customize the tqdm bar for validation.

Returns:

• tqdm –

  A tqdm bar.

Initialize the callback with a stop condition.

Parameters:

• stop_condition (Callable[[], bool]) –

  Function that returns True when prediction should stop.

Check stop condition at the start of each prediction batch.

Parameters:

• trainer (Trainer) –

  PyTorch Lightning trainer instance.

• pl_module (LightningModule) –

  Lightning module being used for prediction.

• batch (Any) –

  Current batch of data.

• batch_idx (int) –

  Index of the current batch.

• dataloader_idx (int, default: 0 ) –

  Index of the dataloader, by default 0.

Raises:

• PredictionStoppedException –

  If stop_condition() returns True.

Lightning module for CAREamics.

This class encapsulates the a PyTorch model along with the training, validation, and testing logic. It is configured using an AlgorithmModel Pydantic class.

Parameters:

• algorithm_config (Union[AlgorithmModel, dict]) –

  Algorithm configuration.

Compute the PSNR for the current validation batch.

Parameters:

• model_output (tuple[Tensor, dict[str, Any]]) –

  Model output, a tuple with the predicted mean and (optionally) logvar, and the top-down data dictionary.

• target (Tensor) –

  Target tensor.

• psnr_func (Callable, default: scale_invariant_psnr ) –

  PSNR function to use, by default scale_invariant_psnr.

Returns:

• list[float] –

  PSNR for each channel in the current batch.

Configure optimizers and learning rate schedulers.

Returns:

• Any –

  Optimizer and learning rate scheduler.

Forward pass.

Parameters:

• x (Tensor) –

  Input tensor of shape (B, (1 + n_LC), [Z], Y, X), where n_LC is the number of lateral inputs.

Returns:

• tuple[Tensor, dict[str, Any]] –

  A tuple with the output tensor and additional data from the top-down pass.

Get the reconstructed tensor from the LVAE model outputs.

Parameters:

• model_outputs (tuple[Tensor, dict[str, Any]]) –

  Model outputs. It is a tuple with a tensor representing the predicted mean and (optionally) logvar, and the top-down data dictionary.

Returns:

• Tensor –

  Reconstructed tensor, i.e., the predicted mean.

Validation epoch end.

Prediction step.

Parameters:

• batch (Tensor) –

  Input batch.

• batch_idx (Any) –

  Batch index.

Returns:

• Any –

  Model output.

Reduce the running PSNR statistics and reset the running PSNR.

Returns:

• Optional[float] –

  Running PSNR averaged over the different output channels.

Set data mean and std for the noise model likelihood.

Parameters:

• data_mean (float) –

  Mean of the data.

• data_std (float) –

  Standard deviation of the data.

Training step.

Parameters:

• batch (tuple[Tensor, Tensor]) –

  Input batch. It is a tuple with the input tensor and the target tensor. The input tensor has shape (B, (1 + n_LC), [Z], Y, X), where n_LC is the number of lateral inputs. The target tensor has shape (B, C, [Z], Y, X), where C is the number of target channels (e.g., 1 in HDN, >1 in muSplit/denoiSplit).

• batch_idx (Any) –

  Batch index.

Returns:

• Any –

  Loss value.

Validation step.

Parameters:

• batch (tuple[Tensor, Tensor]) –

  Input batch. It is a tuple with the input tensor and the target tensor. The input tensor has shape (B, (1 + n_LC), [Z], Y, X), where n_LC is the number of lateral inputs. The target tensor has shape (B, C, [Z], Y, X), where C is the number of target channels (e.g., 1 in HDN, >1 in muSplit/denoiSplit).

• batch_idx (Any) –

  Batch index.