#!/usr/bin/env python
# coding: utf-8

# # An extension of 4
# Here we show how to build a basic neural network for identifying dog or muffin

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt # plotting
from skimage.io import imread # read in images
from skimage.segmentation import mark_boundaries # mark labels
from sklearn.metrics import roc_curve, auc # roc curve tools
from skimage.color import label2rgb
import numpy as np # linear algebra / matrices
# make the notebook interactive
from ipywidgets import interact, interactive, fixed 
import ipywidgets as widgets #add new widgets
from IPython.display import display
import os


# In[2]:


base_path = '04-files'
seg_path = os.path.join(base_path,'DogVsMuffin_seg_bw.jpg')
rgb_path = os.path.join(base_path,'DogVsMuffin.jpg')
face_path = os.path.join(base_path,'DogVsMuffin_face.jpg')
seg_img = imread(seg_path)[80:520,:450]
rgb_img = imread(rgb_path)[80:520,:450,:]
face_img = imread(face_path)
print('RGB Size',rgb_img.shape,'Seg Size',seg_img.shape,'Face Size',face_img.shape)


# # Calculate the baseline ROC curve

# In[3]:


ground_truth_labels = seg_img.flatten()>0
score_value = 1-np.mean(rgb_img.astype(np.float32),2).flatten()/255.0
fpr, tpr, _ = roc_curve(ground_truth_labels,score_value)
roc_auc = auc(fpr,tpr)


# In[4]:


get_ipython().run_line_magic('matplotlib', 'inline')
fig, (ax1, ax2, ax3) = plt.subplots(1,3,figsize = (20,5))
ax1.imshow(rgb_img) # show the color image
ax1.set_title("Color Image")
ax2.imshow(seg_img, cmap='gray') # show the segments
ax2.set_title("Ground Truth")
ax3.imshow(mark_boundaries(rgb_img,seg_img))
ax3.set_title("Labeled Image")


# In[5]:


from keras.models import Sequential
from keras.layers import Conv2D, BatchNormalization, Input
simple_model = Sequential(name = 'SingleConvLayer')
simple_model.add(BatchNormalization(input_shape = (None, None, 3)))
simple_model.add(Conv2D(1, kernel_size = (25, 25), activation = 'sigmoid', padding = 'same'))
simple_model.compile(optimizer = 'sgd', loss = 'binary_crossentropy', metrics = ['binary_accuracy', 'mse'])
simple_model.summary()


# In[6]:


simple_model.fit(np.expand_dims(rgb_img,0), 
                 np.expand_dims(np.expand_dims(seg_img/255.0, -1),0), 
                 epochs=1)


# In[7]:


get_ipython().run_line_magic('matplotlib', 'inline')
fig, (ax1, ax2, ax3) = plt.subplots(1,3,figsize = (20,5))
ax1.imshow(rgb_img) # show the color image
ax1.set_title("Color Image")
ax2.imshow(seg_img, vmin = 0, vmax = 1, cmap = 'jet') # show the segments
ax2.set_title("Ground Truth")
m_output = simple_model.predict(np.expand_dims(rgb_img,0))
ax3.imshow(m_output[0, :, :, 0], vmin = 0, vmax = 1, cmap = 'jet')
ax3.set_title("CNN Prediction");


# In[8]:


loss_history = simple_model.fit(np.expand_dims(rgb_img,0), 
                 np.expand_dims(np.expand_dims(seg_img/255.0, -1),0), 
                 epochs=10, 
                                verbose = False)
loss_history.history['binary_accuracy'][-1]


# In[9]:


get_ipython().run_line_magic('matplotlib', 'inline')
fig, (ax1, ax2, ax3) = plt.subplots(1,3,figsize = (20,5))
ax1.imshow(rgb_img) # show the color image
ax1.set_title("Color Image")
ax2.imshow(seg_img, vmin = 0, vmax = 1, cmap = 'jet') # show the segments
ax2.set_title("Ground Truth")
m_output = simple_model.predict(np.expand_dims(rgb_img,0))
ax3.imshow(m_output[0, :, :, 0], vmin = 0, vmax = 1, cmap = 'jet')
ax3.set_title("CNN Prediction");


# In[10]:


fpr_cnn, tpr_cnn, _ = roc_curve(ground_truth_labels,m_output.ravel())
roc_auc_cnn = auc(fpr_cnn,tpr_cnn)


# In[11]:


get_ipython().run_line_magic('matplotlib', 'inline')
fig, ax = plt.subplots(1,1)
ax.plot(fpr, tpr, label='RGB Value (area = %0.2f)' % roc_auc)
ax.plot(fpr_cnn, tpr_cnn, label='CNN Output (area = %0.2f)' % roc_auc_cnn)
ax.plot([0, 1], [0, 1], 'k--')
ax.set_xlim([0.0, 1.0])
ax.set_ylim([0.0, 1.05])
ax.set_xlabel('False Positive Rate')
ax.set_ylabel('True Positive Rate')
ax.set_title('Receiver operating characteristic example')
ax.legend(loc="lower right")


# In[12]:


from keras.preprocessing.image import ImageDataGenerator
idg = ImageDataGenerator(rotation_range = 2, 
                   zoom_range = 0.1, 
                   width_shift_range = .1, 
                   height_shift_range=0.1, 
                   vertical_flip=True, 
                   horizontal_flip=True)
img_gen = idg.flow(np.expand_dims(rgb_img,0), seed = 1234) 
seg_gen = idg.flow(np.expand_dims(np.expand_dims(seg_img/255.0, -1),0), seed = 1234)


# In[13]:


simple_model.fit_generator(zip(img_gen, seg_gen), steps_per_epoch=5, epochs=5)


# In[14]:


get_ipython().run_line_magic('matplotlib', 'inline')
fig, (ax1, ax2, ax3) = plt.subplots(1,3,figsize = (20,5))
ax1.imshow(rgb_img) # show the color image
ax1.set_title("Color Image")
ax2.imshow(seg_img, vmin = 0, vmax = 1, cmap = 'jet') # show the segments
ax2.set_title("Ground Truth")
m_output = simple_model.predict(np.expand_dims(rgb_img,0))
ax3.imshow(m_output[0, :, :, 0], vmin = 0, vmax = 1, cmap = 'jet')
ax3.set_title("CNN Prediction");


# ### Tasks
# 1. How can you improve the neural network model (what can you add, or change)?
# 2. Where might morphological operations fit in?
# 3. What is wrong with our approach for validating the model here? (what data are we using to measure the accuracy)
# 
# 

# In[ ]: