Simplify Noise2Void

[jdeschamps]

Description

N2VManipulate as a transform makes everything complicated, including a lot of constraints on the Pydantic models and the potential for users to miss-configure their training, in particular in the Lightning API.

In addition, this transform returns a different number of objects has it returns also the original patch. The reason for that is the current implementation of the loss.

Overall, N2V requires many additional checks throughout the code.

Which part of the code?

Examples of N2V checks:

careamics/src/careamics/config/data_model.py

Line 238 in cd2eae2

if SupportedTransform.N2V_MANIPULATE in transform_list:

N2V loss:

careamics/src/careamics/losses/fcn/losses.py

Line 53 in cd2eae2

errors = (original_patches - manipulated_patches) ** 2

Potential solutions

Simplify transform

In N2VManipulate:

return patch * mask

Simplify loss

In n2v_loss, simply pass target:

mask = (target != 0).float()
errors = (target - manipulated_patches) ** 2
loss = torch.sum(errors * mask) / torch.sum(mask)

N2VModule as a subclass of UNetModule

N2VModule as a subclass of UNetModule would allow injecting the N2VManipulate transform. We would need to see how to instantiate it properly based on the configuration.

This would also simplify using N2V in the Lightning API, which is one way users have told us they intend to use N2V.