search

teamtomo / libtilt

Image processing for cryo-electron microscopy in PyTorch
Other
14 stars 5 forks source link

2D +3D unified affine transforms API #64

Open alisterburt opened 7 months ago

alisterburt commented 7 months ago

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 commented 7 months ago

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 commented 7 months ago

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 commented 7 months ago

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 commented 6 months ago

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 commented 6 months ago

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

alisterburt commented 6 months ago

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 commented 6 months ago

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