Feature: Transform multiple arrays identically to `patch` and `target`

melisande-c commented 3 weeks ago

Description

What: Add option in XYFilp and XYRandomRotate90 to transform additional arrays identically to the patch and target. Include the ability to compose these transforms. Also fixed some type hint issues, but some remain.
Why: The LVAE training requires the identical transformation of multiple arrays, including images and noise arrays. Previously it was implemented with the Albumentations library. Original motivation came from this comment chain https://github.com/CAREamics/careamics/pull/223#discussion_r1733656634.
How: A few options were considered, but ultimately I went with what I thought would require the least amount of the code base to change. Additional arrays can passed to the transform __call__ function as kwargs, and they will be returned as a dict after the augmented patch and target. Composition of transforms with additional arrays has been kept as a separate method, Compose.transform_with_additional_arrays, mostly to not have to change too much of the code base, this can be revised and integrated with Compose.__call__ later if desired.

Changes Made

Added:
- Tests for the new functionality.
- _chain_transforms_additional_arrays and transform_with_additional_arrays methods to Compose.
Modified:
- add **additional_arrays: NDArray argument to the __call__ method of some transforms.
- A few existing tests that needed to modify the unpacking of transform outputs, now that they output a tuple of 3 rather than 2.

Breaking changes

Code calling transforms directly, the output is now a tuple of 3 rather than 2.

Additional Notes and Examples

Note

One of the biggest pain points that hasn't been fully resolved is that the N2VManipulate transform doesn't have the same function signature in the __call__ method as the other transforms, it returns 3 arrays rather than 2. We know it gets applied last in the algorithm but we could make it more explicit somehow, this would resolve some annoying type hinting.

Converting LVAE transform code

https://github.com/CAREamics/careamics/blob/b008187292f269455272c92c9009d45a330e7ca0/src/careamics/lvae_training/dataset/vae_dataset.py#L867-L881

To convert the code snippet above, I think (not tested) the following can be done, assuming self.rotation_transform is created using the CAREamics Compose instead of that from Albumentations:

img_kwargs = {}
for i, img in enumerate(img_tuples):
    for k in range(len(img)):
        img_kwargs[f"img{i}_{k}"] = img[k]

noise_kwargs = {}
for i, nimg in enumerate(noise_tuples):
    for k in range(len(nimg)):
        noise_kwargs[f"noise{i}_{k}"] = nimg[k]

img, _, rot_dic = self._rotation_transform.transform_with_additional_arrays(
    img_tuples[0][0], **img_kwargs, **noise_kwargs
)

Please ensure your PR meets the following requirements:

[x] Code builds and passes tests locally, including doctests
[x] New tests have been added (for bug fixes/features)
[x] Pre-commit passes
[ ] PR to the documentation exists (for bug fixes / features)

codecov[bot] commented 3 weeks ago

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 87.09%. Comparing base (427fa26) to head (d357239). Report is 6 commits behind head on main.

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #238 +/- ## ========================================== - Coverage 87.20% 87.09% -0.11% ========================================== Files 130 127 -3 Lines 4876 4868 -8 ========================================== - Hits 4252 4240 -12 - Misses 624 628 +4 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

jdeschamps commented 3 weeks ago

Yes, the N2VManipulate is annoying. We could pack some of the return arrays into one... Dirty and not clear...

Another way would be to inspect again the loss. If I remember it there is a reason to return the three arguments, but maybe we can rewrite it to only need mask * original_patch. (Unfortunately it probably needs the original patch for something).

jdeschamps commented 3 weeks ago

For the sake of the LVAE submission, let's merge it as it. But we should open issues about the following:

One can just be discussed briefly in person: should target be folded in the additional arrays? Would require an additional PR.
One issue to discuss what to do with the N2VManipulate signature, and whether n2v loss can be modified to not require patch and mask, but just patch * mask.

CAREamics / careamics