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
        | Augmentation
        | Sequence[str | Augmentation]
        | 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.
prepare_data –

Prepare input data for training.
train –

Self-supervised training.

Attributes:

n_channels_in (int) –

Returns the number of input channels of the model.
n_channels_out (int) –

Returns the number of output channels of the model.
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 | Augmentation | Sequence[str | Augmentation] | 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_channels_in `property` ¶

n_channels_in: int

Returns the number of input channels of the model.

If the model is an instance of SerializableModel and has an init_params attribute containing the key "n_channels_in", this property returns its value. Otherwise, it defaults to 1.

Returns:

int –

The number of input channels.

n_channels_out `property` ¶

n_channels_out: int

Returns the number of output channels of the model.

If the model is an instance of SerializableModel and has an init_params attribute containing the key "n_channels_out", this property returns its value. Otherwise, it defaults to 1.

Returns:

int –

The number of output channels.

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 its value. Otherwise, it defaults to 2.

Returns:

int –

The expected signal dimensions.

infer ¶

infer(
    inp: NDArray, channel_axis_dst: int | None = None
) -> NDArray

Inference, given an initial stack of images.

Parameters:

inp (NDArray) –

The input stack of images
channel_axis_dst (int | None, default: None ) –

The desired channel axis for the output. If None, the output will have the same channel axis as the input.

Returns:

NDArray –

The denoised stack of images

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

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

    Parameters
    ----------
    inp : NDArray
        The input stack of images
    channel_axis_dst : int | None, optional
        The desired channel axis for the output. If None, the output will have the same channel axis as the input.

    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

    channel_ax_inp = -self.n_dims - 1 if self.n_channels_in > 1 else None
    inp_t = data_to_tensor(inp, device=self.device, n_dims=self.n_dims, channel_axis=channel_ax_inp)

    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

    if channel_axis_dst is not None:
        output = self._move_output_channel_axis(output, channel_axis_dst)

    return output

prepare_data ¶

prepare_data(
    inp: NDArray,
    num_tst_ratio: float = 0.2,
    channel_axis: int | None = None,
) -> tuple[NDArray, list[int]]

Prepare input data for training.

Parameters:

inp (NDArray) –

The input data to be used for training. This should be a NumPy array of shape (N, H, W), where N is the number of samples, and H and W are the height and width of each sample, respectively.
num_tst_ratio (float, default: 0.2 ) –

The ratio of the input data to be used for testing. The remaining data will be used for training. Default is 0.2.
channel_axis (int | None, default: None ) –

The axis of the input array that corresponds to the spectral dimension. If None, the spectral dimension is assumed to not be present. Default is None.

Returns:

tuple[NDArray, NDArray, NDArray] –

A tuple containing: - The input data array. - The mask array indicating the training pixels.

Notes

This function generates input-target pairs based on the specified strategy. It also generates a mask array indicating the training pixels based on the provided ratio.

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

def prepare_data(
    self,
    inp: NDArray,
    num_tst_ratio: float = 0.2,
    channel_axis: int | None = None,
) -> tuple[NDArray, list[int]]:
    """
    Prepare input data for training.

    Parameters
    ----------
    inp : NDArray
        The input data to be used for training. This should be a NumPy array of shape (N, H, W), where N is the
        number of samples, and H and W are the height and width of each sample, respectively.
    num_tst_ratio : float, optional
        The ratio of the input data to be used for testing. The remaining data will be used for training.
        Default is 0.2.
    channel_axis : int | None, optional
        The axis of the input array that corresponds to the spectral dimension.
        If None, the spectral dimension is assumed to not be present.
        Default is None.

    Returns
    -------
    tuple[NDArray, NDArray, NDArray]
        A tuple containing:
        - The input data array.
        - The mask array indicating the training pixels.

    Notes
    -----
    This function generates input-target pairs based on the specified strategy. It also generates a mask array
    indicating the training pixels based on the provided ratio.
    """
    inp, channel_axis = self._prepare_channel_axis(inp, channel_axis)
    self._check_channel_axis_size(inp, channel_axis, "n_channels_in")
    self._check_channel_axis_size(inp, channel_axis, "n_channels_out")

    model_n_axes = self.n_dims + (channel_axis is not None)
    if inp.ndim < model_n_axes:
        raise ValueError(f"Target data should at least be of {model_n_axes + 1} dimensions, but its shape is {inp.shape}")

    batch_length = inp.shape[0]
    mask_tst = get_random_image_indices(batch_length, num_tst_ratio=num_tst_ratio)

    return inp, mask_tst

train ¶

train(
    inp: NDArray,
    tst_inds: Sequence[int] | NDArray,
    *,
    epochs: int,
    mask_shape: int | Sequence[int] | NDArray = 1,
    ratio_blind_spot: float = 0.015,
    learning_rate: float = 0.001,
    optimizer: 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
tst_inds (Sequence[int] | NDArray) –

The validation set indices (indices if Sequence[int])
epochs (int) –

Number of training epochs
mask_shape (int | Sequence[int] | NDArray, default: 1 ) –

Shape of the blind spot mask, by default 1.
ratio_blind_spot (float, default: 0.015 ) –

Ratio of the blind spot size to the total image size, by default 0.015.
learning_rate (float, default: 0.001 ) –

Learning rate for the optimizer, by default 1e-3.
optimizer (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,
    tst_inds: Sequence[int] | NDArray,
    *,
    epochs: int,
    mask_shape: int | Sequence[int] | NDArray = 1,
    ratio_blind_spot: float = 0.015,
    learning_rate: float = 1e-3,
    optimizer: 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
    tst_inds : Sequence[int] | NDArray
        The validation set indices (indices if Sequence[int])
    epochs : int
        Number of training epochs
    mask_shape : int | Sequence[int] | NDArray
        Shape of the blind spot mask, by default 1.
    ratio_blind_spot : float
        Ratio of the blind spot size to the total image size, by default 0.015.
    learning_rate : float
        Learning rate for the optimizer, by default 1e-3.
    optimizer : 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.
    """
    batch_length = inp.shape[0]
    tst_inds = np.array(tst_inds, dtype=int)
    if np.any(tst_inds < 0) or np.any(tst_inds >= batch_length):
        raise ValueError(
            f"Each cross-validation index should be greater or equal than 0, and less than the number of images {batch_length}"
        )
    trn_inds = np.delete(np.arange(batch_length), 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,
        learning_rate=learning_rate,
        optimizer=optimizer,
        regularizer=reg,
        lower_limit=lower_limit,
    )

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

    return losses

noise2void ¶

N2V ¶

n_channels_in property ¶

n_channels_out property ¶

n_dims property ¶

infer ¶

prepare_data ¶

train ¶

n_channels_in `property` ¶

n_channels_out `property` ¶

n_dims `property` ¶