2D +3D unified affine transforms API

[alisterburt]

We should have a nice unified API for defining/applying affine transforms in 2D and 3D.

We can use the existing code in https://github.com/teamtomo/libtilt/blob/main/src/libtilt/transformations.py combined with centered coordinate grid generation from https://github.com/teamtomo/libtilt/blob/main/src/libtilt/grids/coordinate_grid.py

Once the affine API is in place we can add specific useful affines like rotate_2d, rotate_3d, stretch_2d etc

cc @McHaillet

[McHaillet]

Maybe @the-lay has some input on this! He wrote a 3D affine transform package using GPU texture memory (https://github.com/the-lay/voltools)

[the-lay]

Thanks for tagging @McHaillet!

While voltools has some cool unique features (cubic b-spline interpolation, texture memory for super quick operations), it's quite a bit outdated, clunky and has a dependency on CuPy, so I don't think it's a good idea to integrate voltools in any way. Moreover, nowadays, I really think PyTorch is a great foundation!

However, maybe I can suggest looking at voltool's high-level API for inspiration, or for how not to do, no hard feelings :)

Matrix generations: https://github.com/the-lay/voltools/blob/master/voltools/utils/matrices.py#L111
Operations: https://github.com/the-lay/voltools/blob/master/voltools/transforms.py#L25
There's also a StaticVolume class as a representation for a volume that needs to be rotated many times (e.g. for template matching): https://github.com/the-lay/voltools/blob/master/voltools/volume.py#L13

[McHaillet]

While voltools has some cool unique features (cubic b-spline interpolation, texture memory for super quick operations), it's quite a bit outdated, clunky and has a dependency on CuPy, so I don't think it's a good idea to integrate voltools in any way. Moreover, nowadays, I really think PyTorch is a great foundation!

For sure, also not a good idea to mix the two backends. I meant mainly for inspiration, because it is a nice package :)

[McHaillet]

Just saw that pytorch also has functionality for generating a coordinate grid from an affine matrix and input shape: https://pytorch.org/docs/stable/generated/torch.nn.functional.affine_grid.html

@alisterburt do you think that would be good or better to go with libtilt's coordinate_grid?

[McHaillet]

For 2D transformations we can add 3x3 matrix generation for images to transformations.py.

[alisterburt]

Sorry for not chiming in here, my partner has been visiting so I've not been spending much time online. Will check out the voltools API when back, looks nice!

Re: PyTorch's affine, I think using it works if you're okay adopting their conventions about rotation centers etc -> the torch API doesn't match our usual center as far as I know and their coordinate ordering is weird in grid_sample, which I think also transfers over to affine...

My overall feeling is better to keep this API under our control than adopt the 'different to what we as a community might expect and potentially janky in a few places' torch one, do you agree?

[McHaillet]

Sorry for not chiming in here, my partner has been visiting so I've not been spending much time online. Will check out the voltools API when back, looks nice!

No hurry at all! Enjoy the time off :)

My overall feeling is better to keep this API under our control than adopt the 'different to what we as a community might expect and potentially janky in a few places' torch one, do you agree?

If the rotation centering is weird, I absolutely agree. Experienced multiple problems with rotation centers during the phd already haha