Skip to content

deep_image_prior

Unsupervised denoiser implementation, based on the Deep Image Prior.

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

Classes:

  • DIP –

    Deep image prior.

DIP 

DIP(
    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

Deep image prior.

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.

  • train –

    Train the model in an unsupervised manner.

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 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
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 
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
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(
    tgt: NDArray,
    inp: NDArray | None = None,
    *,
    num_tst_ratio: float = 0.2,
    average_redundant: bool = False,
    channel_axis: int | None = None
) -> tuple[NDArray, NDArray, NDArray]

Prepare input data.

Parameters:

  • tgt (NDArray) –

    The target image array. The shape of the output noise array will match the spatial dimensions of this array.

  • inp (NDArray | None, default: None ) –

    The input image array. If provided, it will be used as the initial input for the DIP algorithm. If None, a random noise array will be generated. Default is None.

  • num_tst_ratio (float, default: 0.2 ) –

    The ratio of the test set size to the total dataset size. Default is 0.2.

  • average_redundant (bool, default: False ) –

    If True, average redundant realizations in the target array to match the expected number of dimensions. Default is False.

  • channel_axis (int | None, default: None ) –

    The axis of the target 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: - A random noise array with the same spatial dimensions as the target image. - The target image array. - A mask array indicating the training pixels.

Notes

This function generates a random noise array with the same spatial dimensions as the target image. The noise array is used as the initial input for the DIP algorithm. It also generates a mask array indicating the training pixels based on the provided ratio.

Source code in src/autoden/algorithms/deep_image_prior.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def prepare_data(
    self,
    tgt: NDArray,
    inp: NDArray | None = None,
    *,
    num_tst_ratio: float = 0.2,
    average_redundant: bool = False,
    channel_axis: int | None = None,
) -> tuple[NDArray, NDArray, NDArray]:
    """
    Prepare input data.

    Parameters
    ----------
    tgt : NDArray
        The target image array. The shape of the output noise array will match
        the spatial dimensions of this array.
    inp : NDArray | None, optional
        The input image array. If provided, it will be used as the initial input
        for the DIP algorithm. If None, a random noise array will be generated.
        Default is None.
    num_tst_ratio : float, optional
        The ratio of the test set size to the total dataset size.
        Default is 0.2.
    average_redundant : bool, optional
        If True, average redundant realizations in the target array to match the
        expected number of dimensions. Default is False.
    channel_axis : int | None, optional
        The axis of the target 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:
        - A random noise array with the same spatial dimensions as the target
          image.
        - The target image array.
        - A mask array indicating the training pixels.

    Notes
    -----
    This function generates a random noise array with the same spatial dimensions
    as the target image. The noise array is used as the initial input for the DIP
    algorithm. It also generates a mask array indicating the training pixels based
    on the provided ratio.
    """
    tgt, channel_axis_tgt = self._prepare_channel_axis(tgt, channel_axis)
    self._check_channel_axis_size(tgt, channel_axis_tgt, "n_channels_out")

    model_tgt_axes = self.n_dims + (self.n_channels_out > 1)
    if tgt.ndim < model_tgt_axes:
        raise ValueError(f"Target data should at least be of {model_tgt_axes} dimensions, but its shape is {tgt.shape}")

    if average_redundant and tgt.ndim > model_tgt_axes:
        tgt = tgt.mean(axis=tuple(np.arange(-tgt.ndim, -model_tgt_axes)))

    if inp is None:
        inp = np.random.normal(size=tgt.shape[-model_tgt_axes:], scale=0.25).astype(tgt.dtype)
    else:
        inp, channel_axis_inp = self._prepare_channel_axis(inp, channel_axis)
        self._check_channel_axis_size(inp, channel_axis_inp, "n_channels_inp")

        model_inp_axes = self.n_dims + (self.n_channels_in > 1)
        if inp.ndim < model_inp_axes:
            raise ValueError(f"Input data should at least be of {model_inp_axes} dimensions, but its shape is {inp.shape}")

    mask_trn = get_random_pixel_mask(tgt.shape, mask_pixel_ratio=num_tst_ratio)
    return inp, tgt, mask_trn

train 

train(
    inp: NDArray,
    tgt: NDArray,
    mask_trn: NDArray,
    *,
    epochs: int,
    learning_rate: float = 0.001,
    optimizer: str = "adam",
    lower_limit: float | NDArray | None = None
) -> dict[str, NDArray]

Train the model in an unsupervised manner.

Parameters:

  • inp (NDArray) –

    The input image.

  • tgt (NDArray) –

    The target image to be denoised.

  • mask_trn (NDArray) –

    The mask array indicating the training pixels.

  • epochs (int) –

    The number of training epochs.

  • learning_rate (float, default: 0.001 ) –

    The learning rate for the optimizer. Default is 1e-3.

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

    The optimization algorithm to use. Default is "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.

Returns:

  • dict[str, NDArray] –

    A dictionary containing the training losses.

Notes

This method trains the model using the deep image prior approach in an unsupervised manner. It uses a random initialization for the input image if not provided and applies a scaling and bias transformation to the input and target images. It then trains the model using the specified optimization algorithm and the provided mask array indicating the training pixels.

Source code in src/autoden/algorithms/deep_image_prior.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
def train(
    self,
    inp: NDArray,
    tgt: NDArray,
    mask_trn: NDArray,
    *,
    epochs: int,
    learning_rate: float = 1e-3,
    optimizer: str = "adam",
    lower_limit: float | NDArray | None = None,
) -> dict[str, NDArray]:
    """
    Train the model in an unsupervised manner.

    Parameters
    ----------
    inp : NDArray
        The input image.
    tgt : NDArray
        The target image to be denoised.
    mask_trn : NDArray
        The mask array indicating the training pixels.
    epochs : int
        The number of training epochs.
    learning_rate : float, optional
        The learning rate for the optimizer. Default is 1e-3.
    optimizer : str, optional
        The optimization algorithm to use. Default is "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.

    Returns
    -------
    dict[str, NDArray]
        A dictionary containing the training losses.

    Notes
    -----
    This method trains the model using the deep image prior approach in an unsupervised manner.
    It uses a random initialization for the input image if not provided and applies a scaling and bias
    transformation to the input and target images. It then trains the model using the specified optimization
    algorithm and the provided mask array indicating the training pixels.
    """
    if self.data_sb is None:
        self.data_sb = compute_scaling_supervised(inp, tgt)

    # Rescale the datasets
    tmp_inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp
    tmp_tgt = tgt * self.data_sb.scale_tgt - self.data_sb.bias_tgt

    reg = self._get_regularization()
    losses = self._train_pixelmask_small(
        inp=tmp_inp,
        tgt=tmp_tgt,
        mask_trn=mask_trn,
        epochs=epochs,
        learning_rate=learning_rate,
        optimizer=optimizer,
        regularizer=reg,
        lower_limit=lower_limit,
    )

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

    return losses