careamics 0.1.0rc6__py3-none-any.whl → 0.1.0rc7__py3-none-any.whl

careamics/lvae_training/metrics.py

@@ -0,0 +1,214 @@
+"""
+This script contains the functions/classes to compute loss and metrics used to train and evaluate the performance of the model.
+"""
+
+import numpy as np
+import torch
+from skimage.metrics import structural_similarity
+from torchmetrics.image import MultiScaleStructuralSimilarityIndexMeasure
+
+from careamics.models.lvae.utils import allow_numpy
+
+
+class RunningPSNR:
+    """
+    This class allows to compute the running PSNR during validation step in training.
+    In this way it is possible to compute the PSNR on the entire validation set one batch at the time.
+    """
+
+    def __init__(self):
+        # number of elements seen so far during the epoch
+        self.N = None
+        # running sum of the MSE over the self.N elements seen so far
+        self.mse_sum = None
+        # running max and min values of the self.N target images seen so far
+        self.max = self.min = None
+        self.reset()
+
+    def reset(self):
+        """
+        Used to reset the running PSNR (usually called at the end of each epoch).
+        """
+        self.mse_sum = 0
+        self.N = 0
+        self.max = self.min = None
+
+    def update(self, rec: torch.Tensor, tar: torch.Tensor) -> None:
+        """
+        Given a batch of reconstructed and target images, it updates the MSE and.
+
+        Parameters
+        ----------
+        rec: torch.Tensor
+            Batch of reconstructed images (B, H, W).
+        tar: torch.Tensor
+            Batch of target images (B, H, W).
+        """
+        ins_max = torch.max(tar).item()
+        ins_min = torch.min(tar).item()
+        if self.max is None:
+            assert self.min is None
+            self.max = ins_max
+            self.min = ins_min
+        else:
+            self.max = max(self.max, ins_max)
+            self.min = min(self.min, ins_min)
+
+        mse = (rec - tar) ** 2
+        elementwise_mse = torch.mean(mse.view(len(mse), -1), dim=1)
+        self.mse_sum += torch.nansum(elementwise_mse)
+        self.N += len(elementwise_mse) - torch.sum(torch.isnan(elementwise_mse))
+
+    def get(self):
+        """
+        The get the actual PSNR value given the running statistics.
+        """
+        if self.N == 0 or self.N is None:
+            return None
+        rmse = torch.sqrt(self.mse_sum / self.N)
+        return 20 * torch.log10((self.max - self.min) / rmse)
+
+
+def zero_mean(x):
+    return x - torch.mean(x, dim=1, keepdim=True)
+
+
+def fix_range(gt, x):
+    a = torch.sum(gt * x, dim=1, keepdim=True) / (torch.sum(x * x, dim=1, keepdim=True))
+    return x * a
+
+
+def fix(gt, x):
+    gt_ = zero_mean(gt)
+    return fix_range(gt_, zero_mean(x))
+
+
+def _PSNR_internal(gt, pred, range_=None):
+    if range_ is None:
+        range_ = torch.max(gt, dim=1).values - torch.min(gt, dim=1).values
+
+    mse = torch.mean((gt - pred) ** 2, dim=1)
+    return 20 * torch.log10(range_ / torch.sqrt(mse))
+
+
+@allow_numpy
+def PSNR(gt, pred, range_=None):
+    """
+    Compute PSNR.
+
+    Parameters
+    ----------
+    gt: array
+        Ground truth image.
+    pred: array
+        Predicted image.
+    """
+    assert len(gt.shape) == 3, "Images must be in shape: (batch,H,W)"
+
+    gt = gt.view(len(gt), -1)
+    pred = pred.view(len(gt), -1)
+    return _PSNR_internal(gt, pred, range_=range_)
+
+
+@allow_numpy
+def RangeInvariantPsnr(gt: torch.Tensor, pred: torch.Tensor):
+    """
+    NOTE: Works only for grayscale images.
+    Adapted from https://github.com/juglab/ScaleInvPSNR/blob/master/psnr.py
+    It rescales the prediction to ensure that the prediction has the same range as the ground truth.
+    """
+    assert len(gt.shape) == 3, "Images must be in shape: (batch,H,W)"
+    gt = gt.view(len(gt), -1)
+    pred = pred.view(len(gt), -1)
+    ra = (torch.max(gt, dim=1).values - torch.min(gt, dim=1).values) / torch.std(
+        gt, dim=1
+    )
+    gt_ = zero_mean(gt) / torch.std(gt, dim=1, keepdim=True)
+    return _PSNR_internal(zero_mean(gt_), fix(gt_, pred), ra)
+
+
+def _avg_psnr(target, prediction, psnr_fn):
+    output = np.mean(
+        [
+            psnr_fn(target[i : i + 1], prediction[i : i + 1]).item()
+            for i in range(len(prediction))
+        ]
+    )
+    return round(output, 2)
+
+
+def avg_range_inv_psnr(target, prediction):
+    return _avg_psnr(target, prediction, RangeInvariantPsnr)
+
+
+def avg_psnr(target, prediction):
+    return _avg_psnr(target, prediction, PSNR)
+
+
+def compute_masked_psnr(mask, tar1, tar2, pred1, pred2):
+    mask = mask.astype(bool)
+    mask = mask[..., 0]
+    tmp_tar1 = tar1[mask].reshape((len(tar1), -1, 1))
+    tmp_pred1 = pred1[mask].reshape((len(tar1), -1, 1))
+    tmp_tar2 = tar2[mask].reshape((len(tar2), -1, 1))
+    tmp_pred2 = pred2[mask].reshape((len(tar2), -1, 1))
+    psnr1 = avg_range_inv_psnr(tmp_tar1, tmp_pred1)
+    psnr2 = avg_range_inv_psnr(tmp_tar2, tmp_pred2)
+    return psnr1, psnr2
+
+
+def avg_ssim(target, prediction):
+    ssim = [
+        structural_similarity(
+            target[i], prediction[i], data_range=(target[i].max() - target[i].min())
+        )
+        for i in range(len(target))
+    ]
+    return np.mean(ssim), np.std(ssim)
+
+
+@allow_numpy
+def range_invariant_multiscale_ssim(gt_, pred_):
+    """
+    Computes range invariant multiscale ssim for one channel.
+    This has the benefit that it is invariant to scalar multiplications in the prediction.
+    """
+    shape = gt_.shape
+    gt_ = torch.Tensor(gt_.reshape((shape[0], -1)))
+    pred_ = torch.Tensor(pred_.reshape((shape[0], -1)))
+    gt_ = zero_mean(gt_)
+    pred_ = zero_mean(pred_)
+    pred_ = fix(gt_, pred_)
+    pred_ = pred_.reshape(shape)
+    gt_ = gt_.reshape(shape)
+
+    ms_ssim = MultiScaleStructuralSimilarityIndexMeasure(
+        data_range=gt_.max() - gt_.min()
+    )
+    return ms_ssim(torch.Tensor(pred_[:, None]), torch.Tensor(gt_[:, None])).item()
+
+
+def compute_multiscale_ssim(gt_, pred_, range_invariant=True):
+    """
+    Computes multiscale ssim for each channel.
+    Args:
+    gt_: ground truth image with shape (N, H, W, C)
+    pred_: predicted image with shape (N, H, W, C)
+    range_invariant: whether to use range invariant multiscale ssim
+    """
+    ms_ssim_values = {i: None for i in range(gt_.shape[-1])}
+    for ch_idx in range(gt_.shape[-1]):
+        tar_tmp = gt_[..., ch_idx]
+        pred_tmp = pred_[..., ch_idx]
+        if range_invariant:
+            ms_ssim_values[ch_idx] = range_invariant_multiscale_ssim(tar_tmp, pred_tmp)
+        else:
+            ms_ssim = MultiScaleStructuralSimilarityIndexMeasure(
+                data_range=tar_tmp.max() - tar_tmp.min()
+            )
+            ms_ssim_values[ch_idx] = ms_ssim(
+                torch.Tensor(pred_tmp[:, None]), torch.Tensor(tar_tmp[:, None])
+            ).item()
+
+    output = [ms_ssim_values[i] for i in range(gt_.shape[-1])]
+    return output

careamics/lvae_training/train_lvae.py

@@ -0,0 +1,339 @@
+"""
+This script is meant to load data, intialize the model, and provide the logic for training it.
+"""
+
+import glob
+import os
+import socket
+import sys
+from typing import Dict
+
+import pytorch_lightning as pl
+import torch
+from absl import app, flags
+from ml_collections.config_flags import config_flags
+from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
+from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+from pytorch_lightning.loggers import WandbLogger
+from torch.utils.data import DataLoader
+
+sys.path.append(os.path.join(os.path.dirname(__file__), "../.."))
+print(sys.path)
+
+from careamics.lvae_training.data_modules import LCMultiChDloader, MultiChDloader
+from careamics.lvae_training.data_utils import DataSplitType
+from careamics.lvae_training.lightning_module import LadderVAELight
+from careamics.lvae_training.train_utils import *
+
+FLAGS = flags.FLAGS
+
+config_flags.DEFINE_config_file(
+    "config", None, "Training configuration.", lock_config=False
+)
+flags.DEFINE_string("workdir", None, "Work directory.")
+flags.DEFINE_enum("mode", None, ["train", "eval"], "Running mode: train or eval")
+flags.DEFINE_string(
+    "logdir", "/group/jug/federico/wandb_backup/", "The folder name for storing logging"
+)
+flags.DEFINE_string(
+    "datadir", "/group/jug/federico/careamics_training/data/BioSR", "Data directory."
+)
+flags.DEFINE_boolean("use_max_version", False, "Overwrite the max version of the model")
+flags.DEFINE_string(
+    "load_ckptfpath",
+    "",
+    "The path to a previous ckpt from which the weights should be loaded",
+)
+flags.mark_flags_as_required(["workdir", "config", "mode"])
+
+
+def create_dataset(
+    config,
+    datadir,
+    eval_datasplit_type=DataSplitType.Val,
+    raw_data_dict=None,
+    skip_train_dataset=False,
+    kwargs_dict=None,
+):
+
+    if kwargs_dict is None:
+        kwargs_dict = {}
+
+    datapath = datadir
+
+    # Hard-coded parameters (used to be in the config file)
+    normalized_input = True
+    use_one_mu_std = True
+    train_aug_rotate = False
+    enable_random_cropping = True
+    lowres_supervision = False
+
+    # 1) Data loader for Lateral Contextualization
+    if (
+        "multiscale_lowres_count" in config.data
+        and config.data.multiscale_lowres_count is not None
+    ):
+        # Get padding attributes
+        if "padding_kwargs" not in kwargs_dict:
+            padding_kwargs = {}
+            if "padding_mode" in config.data and config.data.padding_mode is not None:
+                padding_kwargs["mode"] = config.data.padding_mode
+            else:
+                padding_kwargs["mode"] = "reflect"
+            if "padding_value" in config.data and config.data.padding_value is not None:
+                padding_kwargs["constant_values"] = config.data.padding_value
+            else:
+                padding_kwargs["constant_values"] = None
+        else:
+            padding_kwargs = kwargs_dict.pop("padding_kwargs")
+
+        train_data = (
+            None
+            if skip_train_dataset
+            else LCMultiChDloader(
+                config.data,
+                datapath,
+                datasplit_type=DataSplitType.Train,
+                val_fraction=0.1,
+                test_fraction=0.1,
+                normalized_input=normalized_input,
+                use_one_mu_std=use_one_mu_std,
+                enable_rotation_aug=train_aug_rotate,
+                enable_random_cropping=enable_random_cropping,
+                num_scales=config.data.multiscale_lowres_count,
+                lowres_supervision=lowres_supervision,
+                padding_kwargs=padding_kwargs,
+                **kwargs_dict,
+                allow_generation=True,
+            )
+        )
+        max_val = train_data.get_max_val()
+
+        val_data = LCMultiChDloader(
+            config.data,
+            datapath,
+            datasplit_type=eval_datasplit_type,
+            val_fraction=0.1,
+            test_fraction=0.1,
+            normalized_input=normalized_input,
+            use_one_mu_std=use_one_mu_std,
+            enable_rotation_aug=False,  # No rotation aug on validation
+            enable_random_cropping=False,
+            # No random cropping on validation. Validation is evaluated on determistic grids
+            num_scales=config.data.multiscale_lowres_count,
+            lowres_supervision=lowres_supervision,
+            padding_kwargs=padding_kwargs,
+            allow_generation=False,
+            **kwargs_dict,
+            max_val=max_val,
+        )
+    # 2) Vanilla data loader
+    else:
+        train_data_kwargs = {"allow_generation": True, **kwargs_dict}
+        val_data_kwargs = {"allow_generation": False, **kwargs_dict}
+
+        train_data_kwargs["enable_random_cropping"] = enable_random_cropping
+        val_data_kwargs["enable_random_cropping"] = False
+
+        train_data = (
+            None
+            if skip_train_dataset
+            else MultiChDloader(
+                data_config=config.data,
+                fpath=datapath,
+                datasplit_type=DataSplitType.Train,
+                val_fraction=0.1,
+                test_fraction=0.1,
+                normalized_input=normalized_input,
+                use_one_mu_std=use_one_mu_std,
+                enable_rotation_aug=train_aug_rotate,
+                **train_data_kwargs,
+            )
+        )
+
+        max_val = train_data.get_max_val()
+        val_data = MultiChDloader(
+            data_config=config.data,
+            fpath=datapath,
+            datasplit_type=eval_datasplit_type,
+            val_fraction=0.1,
+            test_fraction=0.1,
+            normalized_input=normalized_input,
+            use_one_mu_std=use_one_mu_std,
+            enable_rotation_aug=False,  # No rotation aug on validation
+            max_val=max_val,
+            **val_data_kwargs,
+        )
+
+    # For normalizing, we should be using the training data's mean and std.
+    mean_val, std_val = train_data.compute_mean_std()
+    train_data.set_mean_std(mean_val, std_val)
+    val_data.set_mean_std(mean_val, std_val)
+
+    return train_data, val_data
+
+
+def create_model_and_train(
+    config: ml_collections.ConfigDict,
+    data_mean: Dict[str, torch.Tensor],
+    data_std: Dict[str, torch.Tensor],
+    logger: WandbLogger,
+    checkpoint_callback: ModelCheckpoint,
+    train_loader: DataLoader,
+    val_loader: DataLoader,
+):
+    # tensorboard previous files.
+    for filename in glob.glob(config.workdir + "/events*"):
+        os.remove(filename)
+
+    # checkpoints
+    for filename in glob.glob(config.workdir + "/*.ckpt"):
+        os.remove(filename)
+
+    if "num_targets" in config.model:
+        target_ch = config.model.num_targets
+    else:
+        target_ch = config.data.get("num_channels", 2)
+
+    # Instantiate the model (lightning wrapper)
+    model = LadderVAELight(
+        data_mean=data_mean, data_std=data_std, config=config, target_ch=target_ch
+    )
+
+    # Load pre-trained weights if any
+    if config.training.pre_trained_ckpt_fpath:
+        print("Starting with pre-trained model", config.training.pre_trained_ckpt_fpath)
+        checkpoint = torch.load(config.training.pre_trained_ckpt_fpath)
+        _ = model.load_state_dict(checkpoint["state_dict"], strict=False)
+
+    estop_monitor = config.model.get("monitor", "val_loss")
+    estop_mode = MetricMonitor(estop_monitor).mode()
+
+    callbacks = [
+        EarlyStopping(
+            monitor=estop_monitor,
+            min_delta=1e-6,
+            patience=config.training.earlystop_patience,
+            verbose=True,
+            mode=estop_mode,
+        ),
+        checkpoint_callback,
+        LearningRateMonitor(logging_interval="epoch"),
+    ]
+
+    logger.experiment.config.update(config.to_dict())
+    # wandb.init(config=config)
+    trainer = pl.Trainer(
+        accelerator="gpu",
+        max_epochs=config.training.max_epochs,
+        gradient_clip_val=config.training.grad_clip_norm_value,
+        gradient_clip_algorithm=config.training.gradient_clip_algorithm,
+        logger=logger,
+        callbacks=callbacks,
+        # limit_train_batches = config.training.limit_train_batches,
+        precision=config.training.precision,
+    )
+    trainer.fit(model, train_loader, val_loader)
+
+
+def train_network(
+    train_loader: DataLoader,
+    val_loader: DataLoader,
+    data_mean: Dict[str, torch.Tensor],
+    data_std: Dict[str, torch.Tensor],
+    config: ml_collections.ConfigDict,
+    model_name: str,
+    logdir: str,
+):
+    ckpt_monitor = config.model.get("monitor", "val_loss")
+    ckpt_mode = MetricMonitor(ckpt_monitor).mode()
+    checkpoint_callback = ModelCheckpoint(
+        monitor=ckpt_monitor,
+        dirpath=config.workdir,
+        filename=model_name + "_best",
+        save_last=True,
+        save_top_k=1,
+        mode=ckpt_mode,
+    )
+    checkpoint_callback.CHECKPOINT_NAME_LAST = model_name + "_last"
+    logger = WandbLogger(
+        name=os.path.join(config.hostname, config.exptname),
+        save_dir=logdir,
+        project="Disentanglement",
+    )
+
+    create_model_and_train(
+        config=config,
+        data_mean=data_mean,
+        data_std=data_std,
+        logger=logger,
+        checkpoint_callback=checkpoint_callback,
+        train_loader=train_loader,
+        val_loader=val_loader,
+    )
+
+
+def main(argv):
+    config = FLAGS.config
+
+    assert os.path.exists(FLAGS.workdir)
+    cur_workdir, relative_path = get_workdir(
+        config, FLAGS.workdir, FLAGS.use_max_version
+    )
+    print(f"Saving training to {cur_workdir}")
+
+    config.workdir = cur_workdir
+    config.exptname = relative_path
+    config.hostname = socket.gethostname()
+    config.datadir = FLAGS.datadir
+    config.training.pre_trained_ckpt_fpath = FLAGS.load_ckptfpath
+
+    if FLAGS.mode == "train":
+        set_logger(workdir=cur_workdir)
+        raw_data_dict = None
+
+        # From now on, config cannot be changed.
+        config = ml_collections.FrozenConfigDict(config)
+        log_config(config, cur_workdir)
+
+        train_data, val_data = create_dataset(
+            config, FLAGS.datadir, raw_data_dict=raw_data_dict
+        )
+
+        mean_dict, std_dict = get_mean_std_dict_for_model(config, train_data)
+
+        batch_size = config.training.batch_size
+        shuffle = True
+        train_dloader = DataLoader(
+            train_data,
+            pin_memory=False,
+            num_workers=config.training.num_workers,
+            shuffle=shuffle,
+            batch_size=batch_size,
+        )
+        val_dloader = DataLoader(
+            val_data,
+            pin_memory=False,
+            num_workers=config.training.num_workers,
+            shuffle=False,
+            batch_size=batch_size,
+        )
+
+        train_network(
+            train_loader=train_dloader,
+            val_loader=val_dloader,
+            data_mean=mean_dict,
+            data_std=std_dict,
+            config=config,
+            model_name="BaselineVAECL",
+            logdir=FLAGS.logdir,
+        )
+
+    elif FLAGS.mode == "eval":
+        pass
+    else:
+        raise ValueError(f"Mode {FLAGS.mode} not recognized.")
+
+
+if __name__ == "__main__":
+    app.run(main)

careamics/lvae_training/train_utils.py

@@ -0,0 +1,121 @@
+"""
+This script contains the utility functions for training the LVAE model.
+These functions are mainly used in `train.py` script.
+"""
+
+import logging
+import os
+import pickle
+import time
+from copy import deepcopy
+from datetime import datetime
+from pathlib import Path
+
+import ml_collections
+
+
+def log_config(config: ml_collections.ConfigDict, cur_workdir: str) -> None:
+    # Saving config file.
+    with open(os.path.join(cur_workdir, "config.pkl"), "wb") as f:
+        pickle.dump(config, f)
+    print(f"Saved config to {cur_workdir}/config.pkl")
+
+
+def set_logger(workdir: str) -> None:
+    os.makedirs(workdir, exist_ok=True)
+    fstream = open(os.path.join(workdir, "stdout.txt"), "w")
+    handler = logging.StreamHandler(fstream)
+    formatter = logging.Formatter(
+        "%(levelname)s - %(filename)s - %(asctime)s - %(message)s"
+    )
+    handler.setFormatter(formatter)
+    logger = logging.getLogger()
+    logger.addHandler(handler)
+    logger.setLevel("INFO")
+
+
+def get_new_model_version(model_dir: str) -> str:
+    """
+    A model will have multiple runs. Each run will have a different version.
+    """
+    versions = []
+    for version_dir in os.listdir(model_dir):
+        try:
+            versions.append(int(version_dir))
+        except:
+            print(
+                f"Invalid subdirectory:{model_dir}/{version_dir}. Only integer versions are allowed"
+            )
+            exit()
+    if len(versions) == 0:
+        return "0"
+    return f"{max(versions) + 1}"
+
+
+def get_model_name(config: ml_collections.ConfigDict) -> str:
+    return "LVAE_denoiSplit"
+
+
+def get_workdir(
+    config: ml_collections.ConfigDict,
+    root_dir: str,
+    use_max_version: bool,
+    nested_call: int = 0,
+):
+    rel_path = datetime.now().strftime("%y%m")
+    cur_workdir = os.path.join(root_dir, rel_path)
+    Path(cur_workdir).mkdir(exist_ok=True)
+
+    rel_path = os.path.join(rel_path, get_model_name(config))
+    cur_workdir = os.path.join(root_dir, rel_path)
+    Path(cur_workdir).mkdir(exist_ok=True)
+
+    if use_max_version:
+        # Used for debugging.
+        version = int(get_new_model_version(cur_workdir))
+        if version > 0:
+            version = f"{version - 1}"
+
+        rel_path = os.path.join(rel_path, str(version))
+    else:
+        rel_path = os.path.join(rel_path, get_new_model_version(cur_workdir))
+
+    cur_workdir = os.path.join(root_dir, rel_path)
+    try:
+        Path(cur_workdir).mkdir(exist_ok=False)
+    except FileExistsError:
+        print(
+            f"Workdir {cur_workdir} already exists. Probably because someother program also created the exact same directory. Trying to get a new version."
+        )
+        time.sleep(2.5)
+        if nested_call > 10:
+            raise ValueError(
+                f"Cannot create a new directory. {cur_workdir} already exists."
+            )
+
+        return get_workdir(config, root_dir, use_max_version, nested_call + 1)
+
+    return cur_workdir, rel_path
+
+
+def get_mean_std_dict_for_model(config, train_dset):
+    """
+    Computes the mean and std for the model. This will be subsequently passed to the model.
+    """
+    mean_dict, std_dict = train_dset.get_mean_std()
+
+    return deepcopy(mean_dict), deepcopy(std_dict)
+
+
+class MetricMonitor:
+    def __init__(self, metric):
+        assert metric in ["val_loss", "val_psnr"]
+        self.metric = metric
+
+    def mode(self):
+        if self.metric == "val_loss":
+            return "min"
+        elif self.metric == "val_psnr":
+            return "max"
+        else:
+            raise ValueError(f"Invalid metric:{self.metric}")