Feature: HDN and muSplit

[jdeschamps]

What

HDN and muSplit integration into CAREamics will pose a certain number of challenges:

Specific datasets

muSplit requires specific datasets.

Choice of Dataset will now depend on the algorithm

e.g. muSplit requires a specific dataset spitting out two targets for each input, but also multiple scaled down images for the LC.

Different ways to deal with that:

• Pass the whole configuration to the LightningDataModule, but that means asking for much more parameters in the TrainingDataWrapper
• Create a SplitLigthningDataModule that specifically only instantiates splitting datasets.
• Add a parameter to the DataConfig for the type of datasets and validate it from the Configuration. This is a bit hacky but might be the simpler solution.

Model parameters used by the Dataset

Related to the previous point, the number of LC levels is both a parameter of the model and of the dataset, as it also influences how many lower res images are generated by the dataset.

Model output

muSplit outputs multiple channel, in a single Tensor, so this should not be a problem with CAREamics (as opposed to outputting multiple tensors).

Sampling, generation and averaging

We will need to implement prediction methods for sampling, averaging and generation of images.

Calibration

Calibration is a fitting/learning process that uses many sampled images and the pixel-wise std.

Specific losses

That should not be an issue!

Noise Models

Noise models will need to be implemented in CAREamics. Currently, old code (that was never used) currently resides in https://github.com/CAREamics/careamics/tree/jd/refac/save_noisemodels.

Model Zoo export

What challenges await us? @mese79

---

Any other issue we foresee?

@CatEek @federico-carrara @ashesh-0

[mese79]

This is the wrapper I wrote for LadderVAE model in uSplit:

class LVAE(nn.Module):
    def __init__(self):
        super().__init__()
        # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # configs
        config = get_config()
        config = ml_collections.ConfigDict(config)

        mean = np.array([[[[404.2586]], [[404.2586]]]])
        std = np.array([[[[235.67084]], [[235.67084]]]])
        # mean = np.array([[[[0.]], [[0.]]]])
        # std = np.array([[[[1.0]], [[1.0]]]])

        self.lvae_model = LadderVAE(mean, std, config)
        self.lvae_model.non_stochastic_version = True

    def forward(self, x):
        lvae_out, _ = self.lvae_model(x)
        out = self.get_img_from_forward_output(lvae_out)

        return out

    def get_img_from_forward_output(self, forward_out):
        recons_img = self.lvae_model.likelihood.get_mean_lv(forward_out)[0]
        recons_img = recons_img.cpu() * self.lvae_model.data_std + self.lvae_model.data_mean

        return recons_img

def get_config():
    config =  {
        ...
    }

    return config

But this only works for a batch of 5x64x64 patches and there is no data pipeline in there.

[ashesh-0]

I assume it is some deadline to submit the models which is making @mese79 do ^^, otherwise @federico-carrara has already HDN,uSplit,denoiSplit all implemented as nn.Module. A bit of debugging still remains, ofcourse ;)

But more specifically on this,

1. model.reset_for_different_output_size() should be called if you want to call it for different spatial size.
2. 5 in 5x64x64 is the LC dimension. Since the network was trained with that, it should not be changed.

In case reset_for_different_output_size is not there in the branch you are working for, here you go:

    def reset_for_different_output_size(self, output_size):
        """
        This should be called if we want to predict for a different input/output size.
        """
        for i in range(self.n_layers):
            sz = output_size // 2**(1 + i)
            self.bottom_up_layers[i].output_expected_shape = (sz, sz)
            self.top_down_layers[i].latent_shape = (output_size, output_size)