POT
POT copied to clipboard
dim>1 for ot.bregman.barycenter
🚀 Feature
Extend ot.bregman.barycenter to higher dimensions (>1D data)
Motivation
Computing barycenters for color images requires at least 3 dimensions (rgb or lab).
Pitch
Recent papers have suggested this for data augmentation
Hello @gabrieldernbach,
Note that the current solver ot.bregman.barycenter works in any dimensions since you provide the solver with a distance matrix that can be computed in any dimensions (it won't scale to millions of pixels though).
We also have solvers in 2d for regular grids here: https://pythonot.github.io/auto_examples/barycenters/plot_convolutional_barycenter.html#sphx-glr-auto-examples-barycenters-plot-convolutional-barycenter-py
And the solver with free support works also very well in any dimension: https://pythonot.github.io/auto_examples/barycenters/plot_free_support_barycenter.html#sphx-glr-auto-examples-barycenters-plot-free-support-barycenter-py
I agree that some higher dimesiond solvers for grid (3d color histograms) would be interesting and we will look into implementing it in the future.
Dear rflamary thank you a lot for your quick reply. I am happy to hear that this is already possible. I had a try, but could not figure the right input shapes. May I suggest to add a note to the docs?
In the documentation it says that
A (ndarray, shape (dim, n_hists)) – n_hists training distributions a_i of size dim
M (ndarray, shape (dim, dim)) – loss matrix for OT
Lets say I have 128 pixel in 3 dimension (lab). Then this runs through, as expected
A = np.random.rand(128, 3)
M = np.random.rand(128, 128)
ot.bregman.barycenter_sinkhorn(A, M, reg=0.01)
If I now have 5 images for which I wish to find the barycenters
A = np.random.rand(5, 128, 3)
M = np.random.rand(5, 128, 128) # also tried (128, 128)
ot.bregman.barycenter_sinkhorn(A, M, reg=0.01)
results in a dimension misalignment error. I tried some permutations but couldn't find the appropriate version.
Following the tutorial example, where the arguments are lists, I also tried
A = [np.random.rand(128, 3) for _ in range(3)]
M = np.random.rand(128, 128)
ot.bregman.barycenter_sinkhorn(A, M, reg=0.01)
which fails by the fact that A is supposed to have a shape property.
Thanks for your help.
This is something I am also struggling with. I would prefer to use ot.bregman.barycenter over the free_support version because then I can specify the cost matrix. Thanks