# change to the shogun-data directory
import os
SHOGUN_DATA_DIR=os.getenv('SHOGUN_DATA_DIR', '../../../data')
os.chdir(os.path.join(SHOGUN_DATA_DIR, 'ica'))

from PIL import Image
import numpy as np

# Load Images as grayscale images and convert to numpy arrays
s1 = np.asarray(Image.open("lena.jpg").convert('L'))
s2 = np.asarray(Image.open("monalisa.jpg").convert('L'))

# Save Image Dimensions
# we'll need these later for reshaping the images
rows = s1.shape[0]
cols = s1.shape[1]

"""
Explanation: Blind Source Separation on Images with Shogun
by Kevin Hughes
This notebook illustrates <a href="http://en.wikipedia.org/wiki/Blind_signal_separation">Blind Source Seperation</a>(BSS) on images using <a href="http://en.wikipedia.org/wiki/Independent_component_analysis">Independent Component Analysis</a> (ICA) in Shogun. This is very similar to the <a href="http://www.shogun-toolbox.org/static/notebook/current/bss_audio.html">BSS audio notebook</a> except that here we have used images instead of audio signals.
The first step is to load 2 images from the Shogun data repository:
End of explanation
"""


%matplotlib inline
import matplotlib.pyplot as plt

# Show Images
f,(ax1,ax2) = plt.subplots(1,2)
ax1.imshow(s1, cmap=plt.gray()) # set the color map to gray, only needs to be done once!
ax2.imshow(s2)

"""
Explanation: Displaying the images using pylab:
End of explanation
"""


# Convert Images to row vectors
# and stack into a Data Matrix
S = np.c_[s1.flatten(), s2.flatten()].T

"""
Explanation: In our previous ICA examples the input data or source signals were already 1D but these images are obviously 2D. One common way to handle this case is to simply "flatten" the 2D image matrix into a 1D row vector. The same idea can also be applied to 3D data, for example a 3 channel RGB image can be converted a row vector by reshaping each 2D channel into a row vector and then placing them after each other length wise.
Lets prep the data:
End of explanation
"""


# Mixing Matrix
A = np.array([[1, 0.5], [0.5, 1]])

# Mix Signals
X = np.dot(A,S)

# Show Images
f,(ax1,ax2) = plt.subplots(1,2)
ax1.imshow(X[0,:].reshape(rows,cols))
ax2.imshow(X[1,:].reshape(rows,cols))

"""
Explanation: It is pretty easy using a nice library like numpy.
Next we need to mix our source signals together. We do this exactly the same way we handled the audio data - take a look!
End of explanation
"""


import shogun as sg

mixed_signals = sg.create_features(X)

# Separating
jade = sg.create_transformer('Jade')
jade.fit(mixed_signals)
signals = jade.transform(mixed_signals)
S_ = signals.get('feature_matrix')

# Show Images
f,(ax1,ax2) = plt.subplots(1,2)
ax1.imshow(S_[0,:].reshape(rows,cols) *-1)
ax2.imshow(S_[1,:].reshape(rows,cols))

"""
Explanation: Notice how we had to reshape from a 1D row vector back into a 2D matrix of the correct shape. There is also another nuance that I would like to mention here: pylab is actually doing quite a lot for us here that you might not be aware of. It does a pretty good job determining the value range of the image to be shown and then it applies the color map. Many other libraries (for example OpenCV's highgui) won't be this helpful and you'll need to remember to scale the image appropriately on your own before trying to display it. 
Now onto the exciting step, unmixing the images using ICA! Again this step is the same as when using Audio data. Again we need to reshape the images before viewing them and an additional nuance was to add the *-1 to the first separated signal. I did this after viewing the result the first time as the image was clearly inversed, this can happen because ICA can't necessarily capture the correct phase.
End of explanation
"""


