noise2void

Self-supervised denoiser implementation, based on Noise2Void.

@author: Nicola VIGANÒ, CEA-MEM, Grenoble, France

Classes:

• N2V –

  Self-supervised denoising from single images.

N2V

N2V(
    model: int | str | NetworkParams | Module | Mapping,
    data_scale_bias: DataScaleBias | None = None,
    reg_val: float | LossRegularizer | None = None,
    device: str = "cuda" if is_available() else "cpu",
    batch_size: int | None = None,
    augmentation: str | Sequence[str] | None = None,
    save_epochs_dir: str | None = None,
    verbose: bool = True,
)

Bases: Denoiser

Self-supervised denoising from single images.

Parameters:

• model (str | NetworkParams | Module | Mapping | None) –

  Type of neural network to use or a specific network (or state) to use

• data_scale_bias (DataScaleBias | None, default: None ) –

  Scale and bias of the input data, by default None

• reg_val (float | None, default: None ) –

  Regularization value, by default 1e-5

• device (str, default: 'cuda' if is_available() else 'cpu' ) –

  Device to use, by default "cuda" if cuda is available, otherwise "cpu"

• save_epochs_dir (str | None, default: None ) –

  Directory where to save network states at each epoch. If None disabled, by default None

• verbose (bool, default: True ) –

  Whether to produce verbose output, by default True

Methods:

• infer –

  Inference, given an initial stack of images.

• train –

  Self-supervised training.

Attributes:

• n_dims (int) –

  Returns the expected signal dimensions.

Source code in src/autoden/algorithms/denoiser.py

def __init__(
    self,
    model: int | str | NetworkParams | pt.nn.Module | Mapping,
    data_scale_bias: DataScaleBias | None = None,
    reg_val: float | LossRegularizer | None = None,
    device: str = "cuda" if pt.cuda.is_available() else "cpu",
    batch_size: int | None = None,
    augmentation: str | Sequence[str] | None = None,
    save_epochs_dir: str | None = None,
    verbose: bool = True,
) -> None:
    """Initialize the noise2noise method.

    Parameters
    ----------
    model : str | NetworkParams | pt.nn.Module | Mapping | None
        Type of neural network to use or a specific network (or state) to use
    data_scale_bias : DataScaleBias | None, optional
        Scale and bias of the input data, by default None
    reg_val : float | None, optional
        Regularization value, by default 1e-5
    device : str, optional
        Device to use, by default "cuda" if cuda is available, otherwise "cpu"
    save_epochs_dir : str | None, optional
        Directory where to save network states at each epoch.
        If None disabled, by default None
    verbose : bool, optional
        Whether to produce verbose output, by default True
    """
    if isinstance(model, int):
        if self.save_epochs_dir is None:
            raise ValueError("Directory for saving epochs not specified")

        model = load_model_state(self.save_epochs_dir, epoch_num=model)

    if isinstance(model, (str, NetworkParams, Mapping, pt.nn.Module)):
        self.model = create_network(model, device=device)
    else:
        raise ValueError(f"Invalid model {type(model)}")
    if verbose:
        get_num_parameters(self.model, verbose=True)

    if augmentation is None:
        augmentation = []
    elif isinstance(augmentation, str):
        augmentation = [augmentation.lower()]
    elif isinstance(augmentation, Sequence):
        augmentation = [str(a).lower() for a in augmentation]

    self.data_sb = data_scale_bias

    self.reg_val = reg_val
    self.device = device
    self.batch_size = batch_size
    self.augmentation = augmentation
    self.save_epochs_dir = save_epochs_dir
    self.verbose = verbose

n_dims property

n_dims: int

Returns the expected signal dimensions.

If the model is an instance of SerializableModel and has an init_params attribute containing the key "n_dims", this property returns the value associated with "n_dims". Otherwise, it defaults to 2.

Returns:

• int –

  The expected signal dimensions.

infer

infer(inp: NDArray) -> NDArray

Inference, given an initial stack of images.

Parameters:

• inp (NDArray) –

  The input stack of images

Returns:

• NDArray –

  The denoised stack of images

Source code in src/autoden/algorithms/denoiser.py

def infer(self, inp: NDArray) -> NDArray:
    """Inference, given an initial stack of images.

    Parameters
    ----------
    inp : NDArray
        The input stack of images

    Returns
    -------
    NDArray
        The denoised stack of images
    """
    # Rescale input
    if self.data_sb is not None:
        inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp

    inp_t = data_to_tensor(inp, device=self.device, n_dims=self.n_dims)

    self.model.eval()
    with pt.inference_mode():
        out_t: pt.Tensor = self.model(inp_t)
        output = out_t.squeeze(dim=(0, 1)).to("cpu").numpy()

    # Rescale output
    if self.data_sb is not None:
        output = (output + self.data_sb.bias_out) / self.data_sb.scale_out

    return output

train

train(
    inp: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    mask_shape: int | Sequence[int] | NDArray = 1,
    ratio_blind_spot: float = 0.015,
    algo: str = "adam",
    lower_limit: float | NDArray | None = None,
) -> dict[str, NDArray]

Self-supervised training.

Parameters:

• inp (NDArray) –

  The input images, which will also be targets

• epochs (int) –

  Number of training epochs

• tst_inds (Sequence[int] | NDArray) –

  The validation set indices

• mask_shape (int | Sequence[int] | NDArray, default: 1 ) –

  Shape of the blind spot mask, by default 1.

• algo (str, default: 'adam' ) –

  Optimizer algorithm to use, by default "adam"

• lower_limit (float | NDArray | None, default: None ) –

  The lower limit for the input data. If provided, the input data will be clipped to this limit. Default is None.

Source code in src/autoden/algorithms/noise2void.py

def train(
    self,
    inp: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    mask_shape: int | Sequence[int] | NDArray = 1,
    ratio_blind_spot: float = 0.015,
    algo: str = "adam",
    lower_limit: float | NDArray | None = None,
) -> dict[str, NDArray]:
    """Self-supervised training.

    Parameters
    ----------
    inp : NDArray
        The input images, which will also be targets
    epochs : int
        Number of training epochs
    tst_inds : Sequence[int] | NDArray
        The validation set indices
    mask_shape : int | Sequence[int] | NDArray
        Shape of the blind spot mask, by default 1.
    algo : str, optional
        Optimizer algorithm to use, by default "adam"
    lower_limit : float | NDArray | None, optional
        The lower limit for the input data. If provided, the input data will be clipped to this limit.
        Default is None.
    """
    num_imgs = inp.shape[0]
    tst_inds = np.array(tst_inds, dtype=int)
    if np.any(tst_inds < 0) or np.any(tst_inds >= num_imgs):
        raise ValueError(
            f"Each cross-validation index should be greater or equal than 0, and less than the number of images {num_imgs}"
        )
    trn_inds = np.delete(np.arange(num_imgs), obj=tst_inds)

    if self.data_sb is None:
        self.data_sb = compute_scaling_selfsupervised(inp)

    # Rescale the datasets
    inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp

    inp_trn = inp[trn_inds]
    inp_tst = inp[tst_inds]

    reg = self._get_regularization()
    losses = self._train_n2v_pixelmask_small(
        inp_trn,
        inp_tst,
        epochs=epochs,
        mask_shape=mask_shape,
        ratio_blind_spot=ratio_blind_spot,
        algo=algo,
        regularizer=reg,
        lower_limit=lower_limit,
    )

    if self.verbose:
        self._plot_loss_curves(losses, f"Self-supervised {self.__class__.__name__} {algo.upper()}")

    return losses