# -*- coding: utf-8 -*-
"""
Created on Mon Apr 17 08:43:48 2023

@author: mritchey
"""
import keras
import streamlit as st
from PIL import Image
import pandas as pd
import numpy as np
from keras import layers  
import matplotlib.pyplot as plt 

def get_img_array(img_path, target_size):
    array = keras.utils.img_to_array(img)
    array = np.expand_dims(array, axis=0)
    return array
st.set_page_config(layout="wide")



model_type = st.sidebar.selectbox(
    'Select Model', ('VGG16', 'VGG19', 'ResNet50V2', 'MobileNetV2'))
models = {'VGG16': 'vgg16', 'VGG19': 'vgg19', 'ResNet50V2': 'resnet_v2',
          'MobileNetV2': 'mobilenet_v2'}
model_type2 = models[model_type]

top_n = st.sidebar.selectbox('Number of Results', (3, 5, 10))
results = st.sidebar.selectbox('Display Summary', ('No','Yes'))
display = st.sidebar.selectbox('Display Filtered Images', ('No','Yes'))

exec(f'from keras.applications.{model_type2} import {model_type}')
exec(
    f'from keras.applications.{model_type2} import preprocess_input, decode_predictions')
model = eval(f'{model_type}(weights="imagenet")')

img_path = st.file_uploader("Upload Picture")

try:
    img = Image.open(img_path)
except:
    img = Image.open('dog.jpg')


img = img.resize((224, 224))  # Resize to match VGG16 input size
x = np.array(img)
x = np.expand_dims(x, axis=0)
x = preprocess_input(x)

# Make predictions on the image
preds = model.predict(x)
# Convert the predictions to human-readable labels
decoded_preds = decode_predictions(preds, top=top_n)[0]

df = pd.DataFrame(decoded_preds)
df.columns = ['label', 'Object', 'Percent Certainty']
df.index = df.index+1
df = df[['Object', 'Percent Certainty']]
df['Percent Certainty'] = df['Percent Certainty'].apply(
    lambda x: '{:.2%}'.format(x))

# with st.container():

with st.container():
    col1, col2 = st.columns((1,3))
    with col1:
        st.image(img,width=400)
    with col2:
        st.dataframe(df)
        
    
with st.container():
    col1, col2 = st.columns((2, 4))
    if results=='Yes':
        with col1:
            stringlist = []
            model.summary(print_fn=lambda x: stringlist.append(x))
            short_model_summary = "\n".join(stringlist)
            print(short_model_summary)
            st.write(short_model_summary)
    
    
    if display =='Yes':
        img_tensor = get_img_array(img, target_size=(224, 224))
        layer_outputs = []
        layer_names = []
        for layer in model.layers:
            if isinstance(layer, (layers.Conv2D, layers.MaxPooling2D)):
                layer_outputs.append(layer.output)
                layer_names.append(layer.name)
        activation_model = keras.Model(inputs=model.input, outputs=layer_outputs)
        
        activations = activation_model.predict(img_tensor)
        
        first_layer_activation = activations[0]
        
        plt.matshow(first_layer_activation[0, :, :, 5], cmap="viridis")
        
        images_per_row = 16
        all_pngs=[]
        for layer_name, layer_activation in zip(layer_names, activations):
            n_features = layer_activation.shape[-1]
            size = layer_activation.shape[1]
            n_cols = n_features // images_per_row
            display_grid = np.zeros(((size + 1) * n_cols - 1,
                                     images_per_row * (size + 1) - 1))
            for col in range(n_cols):
                for row in range(images_per_row):
                    channel_index = col * images_per_row + row
                    channel_image = layer_activation[0, :, :, channel_index].copy()
                    if channel_image.sum() != 0:
                        channel_image -= channel_image.mean()
                        channel_image /= channel_image.std()
                        channel_image *= 64
                        channel_image += 128
                    channel_image = np.clip(channel_image, 0, 255).astype("uint8")
                    display_grid[
                        col * (size + 1): (col + 1) * size + col,
                        row * (size + 1) : (row + 1) * size + row] = channel_image
            scale = 1. / size
            plt.figure(figsize=(scale * display_grid.shape[1],
                                scale * display_grid.shape[0]))
            plt.title(layer_name)
            plt.grid(False)
            plt.axis("off")
            plt.imshow(display_grid, aspect="auto", cmap="viridis")
            
            filename=f'{layer_name}.png'
            plt.savefig(f'{layer_name}.png')
            all_pngs.append(filename)
        with col2:
            for i in all_pngs: st.image(i)