%matplotlib inline

import seaborn as sns
import matplotlib.pyplot as plt
from thunder import Colorize
image = Colorize.image
tile = Colorize.tile
sns.set_style('darkgrid')
sns.set_context('notebook')

"""
Explanation: Image registration
A common problem when working with large collections of related images is registrating them together, relative to a reference. Thunder implements a generic image registration API that supports a variety of registration algorithms, and exposes to the user both standard approaches and the ability to implement custom solutions. Here, we generate data for registration, apply a registration algorithm, and show the results.
If you are interested in contributing a new registration algorithm to Thunder, let us know in the chatroom!
Setup plotting
End of explanation
"""


data = tsc.loadExample('mouse-images')
data

"""
Explanation: Generating data
We will use the included toy example image data to test registration algorithms. These data do not actually have any motion, so to test the algorithms, we will induce fake motion. First we'll load and inspect the data.
End of explanation
"""


from numpy import random
from scipy.ndimage.filters import gaussian_filter
t = 500
dx = gaussian_filter(random.randn(t), 50) * 25
dy = gaussian_filter(random.randn(t), 50) * 25

plt.plot(dx);
plt.plot(dy);

"""
Explanation: There are 500 images (corresponding to 500 time points), and the data are two-dimensional, so we'll want to generate 500 random shifts in x and y. We'll use smoothing functions from scipy to make sure the drift varies slowly over time, which will be easier to look at.
End of explanation
"""


from scipy.ndimage import shift
shifted = data.apply(lambda (k, v): (k, shift(v, (dx[k], dy[k]), mode='nearest', order=0)))

"""
Explanation: Now let's use these drifts to shift the data. We'll use the apply method on our data, which applies an arbitrary function to each record; in this case, the function is to shift by an amount given by the corresponding entry in our list of shifts.
End of explanation
"""


im1 = data[0]
im2 = shifted[0]
tile([im1, im2, im1-im2], clim=[(0,300), (0,300), (-300,300)], grid=(1,3), size=14)

"""
Explanation: Look at the first entry of both the original images and the shifted images, and their difference
End of explanation
"""


tile([data.mean(), shifted.mean()], size=14)

"""
Explanation: It's also useful to look at the mean of the raw images and the shifted images, the mean of the shifted images should be much more blurry!
End of explanation
"""


from thunder import Registration
reg = Registration('crosscorr')

"""
Explanation: Registration
To run registration, first we create a registration method by specifying its name (current options include 'crosscorr' and 'planarcrosscorr')
End of explanation
"""


reg.prepare(shifted, startIdx=0, stopIdx=500);

"""
Explanation: This method computes a cross-correlation in parallel between every image and a reference. To compute that reference, we can use the prepare method, and either give it a reference, or have it compute one for us. For this method, the default prepare is to compute a mean, over some specified range. We call:
End of explanation
"""


image(reg.reference)

"""
Explanation: This adds a reference attribute to the reg object, which we can look at
End of explanation
"""


ref = shifted.filterOnKeys(lambda k: k > 0 and k < 500).mean()
reg.prepare(ref)
image(reg.reference)

"""
Explanation: We could have equivalently computed the reference ourselves (using the mean, or any other calculation) and passed it as an argument
End of explanation
"""


model = reg.fit(shifted)

"""
Explanation: Now we use the registration method reg and fit it to the shifted data, returning a fitted RegistrationModel
End of explanation
"""


model

"""
Explanation: Inspect the model
End of explanation
"""


model[0]

"""
Explanation: The model represents a list of transformations. You can inspect them:
End of explanation
"""


clrs = sns.color_palette('deep')
plt.plot(model.toArray()[:,0], color=clrs[0])
plt.plot(dx, color=clrs[0])
plt.plot(model.toArray()[:,1], color=clrs[1])
plt.plot(dy, color=clrs[1]);

"""
Explanation: You can also convert the full collection of transformations into an array, which is useful for plotting. Here we'll plot the estimated transformations relative to the ground truth, they should be fairly similar.
End of explanation
"""


reg.prepare(data, startIdx=0, stopIdx=500)
model = reg.fit(shifted)

"""
Explanation: Note that, while following a similar pattern as the ground truth, the estimates are not perfect. That's because we didn't use the true reference to estimate the displacements, but rather the mean of the displaced data. To see that we get the exact displacements back, let's compute a reference from the original, unshifted data.
End of explanation
"""


plt.plot(model.toArray()[:,0], color=clrs[0])
plt.plot(dx, color=clrs[0])
plt.plot(model.toArray()[:,1], color=clrs[1])
plt.plot(dy, color=clrs[1]);

"""
Explanation: Now the estimates should be exact (up to rounding error)! But note that this is sort of cheating, because in general we don't know the ground truth.
End of explanation
"""


corrected = model.transform(shifted)

"""
Explanation: We can now use our model to transform a set of images, which applies the estimated transformations. The API design makes it easy to apply the transformations to the dataset we used to estimate the transformations, or a different one. We'll use the model we just estimates, which used the true reference, because it will be easy to see that it did the right thing.
End of explanation
"""


im1 = data[0]
im2 = corrected[0]
tile([im1, im2, im1-im2], clim=[(0,300), (0,300), (-300,300)], grid=(1,3), size=14)

"""
Explanation: Let's again look at the first image from the orignal and corrected, and their difference. Whereas before they were different, now they should be the same, except for minor near the boundaries (where the image has been replaced with its nearest neighbors).
End of explanation
"""


tile([shifted.mean(), corrected.mean()], size=14)

"""
Explanation: As a final check on the registation, we can compare the mean of the shifted data, and the mean of the regsitered data. The latter should be much sharper.
End of explanation
"""


