import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches
import numpy as np
from sklearn.model_selection import train_test_split
import tensorflow as tf
import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from joblib import dump, load
from sklearn.ensemble import RandomForestRegressor
from skimage.feature import hog
from skimage import exposure

NN = keras.models.load_model('model4.keras')
CNN = keras.models.load_model('model3.keras')
DT = load('random_forest_model.joblib')
RR = load('ridge_regression_model.joblib')

def preprocess_images(images):
    hog_features = []
    for image in images:
        fd, hog_image = hog(image, orientations=8, pixels_per_cell=(8, 8), cells_per_block=(1, 1), visualize=True)
        hog_features.append(fd)
    return np.array(hog_features)

# Specify image module
input_module1 = gr.File(label = "Input Image File (Must be 128x128 Greyscale)")

# Specify method dropdown menu
input_module2 = gr.Dropdown(choices=['Convolutional Neural Network', 'Neural Network', 'Random Forest', 'Ridge Regression', 'All'], label = "Method")

# Specify output 1
output_module1 = gr.Plot(label = "Predicted Box")

def multi_inputs(input1, input2):
    import numpy as np
    ## processing inputs
    plt.clf()
    photo = plt.imread(input1)
    Dcells = 128*128
    if input2 == "Neural Network":
      flattened = photo.reshape(1,Dcells)
      predicted_boxnn = NN.predict((flattened))
      box_xnn = predicted_boxnn[0,0]
      box_ynn = predicted_boxnn[0,1]
      box_widthnn = predicted_boxnn[0,2]
      box_heightnn = predicted_boxnn[0,3]
      plt.imshow(photo, cmap='gray')
      plt.gca().add_patch(matplotlib.patches.Rectangle((box_xnn, box_ynn), box_widthnn, box_heightnn, ec='r', fc='none'))
      plt.gca().annotate("NN", xy = (0,0), xytext = (box_xnn, box_ynn+box_heightnn+4), color='r',fontsize = 8)
    elif input2 == "Convolutional Neural Network":
      shaped = photo.reshape(1,128,128)
      predicted_boxcnn = CNN.predict(shaped)
      box_xcnn = predicted_boxcnn[0,0]
      box_ycnn = predicted_boxcnn[0,1]
      box_widthcnn = predicted_boxcnn[0,2]
      box_heightcnn = predicted_boxcnn[0,3]
      plt.imshow(photo, cmap='gray')
      plt.gca().add_patch(matplotlib.patches.Rectangle((box_xcnn, box_ycnn), box_widthcnn, box_heightcnn, ec='b', fc='none'))
      plt.gca().annotate("CNN", xy = (0,0), xytext = (box_xcnn, box_ycnn+box_heightcnn+4), color='b',fontsize = 8)
    elif input2 == "Random Forest":
      processed = preprocess_images(photo.reshape(1,128,128))
      predicted_boxRF = DT.predict(processed)
      box_xRF = predicted_boxRF[0,0]
      box_yRF = predicted_boxRF[0,1]
      box_widthRF = predicted_boxRF[0,2]
      box_heightRF = predicted_boxRF[0,3]
      plt.imshow(photo, cmap='gray')
      plt.gca().add_patch(matplotlib.patches.Rectangle((box_xRF, box_yRF), box_widthRF, box_heightRF, ec='y', fc='none'))
      plt.gca().annotate("RF", xy = (0,0), xytext = (box_xRF, box_yRF+box_heightRF+4), color='y',fontsize = 8)
    elif input2 == "Ridge Regression":
      #RR
      processed = preprocess_images(photo.reshape(1,128,128))
      predicted_boxRR = RR.predict(processed)
      box_xRR = predicted_boxRR[0,0]
      box_yRR = predicted_boxRR[0,1]
      box_widthRR = predicted_boxRR[0,2]
      box_heightRR = predicted_boxRR[0,3]
      plt.imshow(photo, cmap='gray')
      plt.gca().add_patch(matplotlib.patches.Rectangle((box_xRR, box_yRR), box_widthRR, box_heightRR, ec='c', fc='none'))
      plt.gca().annotate("RR", xy = (0,0), xytext = (box_xRR, box_yRR+box_heightRR+4), color='c',fontsize = 8)
    else:
      #NN
      flattened = photo.reshape(1,Dcells)
      predicted_boxnn = NN.predict((flattened))
      box_xnn = predicted_boxnn[0,0]
      box_ynn = predicted_boxnn[0,1]
      box_widthnn = predicted_boxnn[0,2]
      box_heightnn = predicted_boxnn[0,3]
      plt.imshow(photo, cmap='gray')
      plt.gca().add_patch(matplotlib.patches.Rectangle((box_xnn, box_ynn), box_widthnn, box_heightnn, ec='r', fc='none'))
      plt.gca().annotate("NN", xy = (0,0), xytext = (0,112), color='r',fontsize = 8)
      #CNN
      shaped = photo.reshape(1,128,128)
      predicted_boxcnn = CNN.predict(shaped)
      box_xcnn = predicted_boxcnn[0,0]
      box_ycnn = predicted_boxcnn[0,1]
      box_widthcnn = predicted_boxcnn[0,2]
      box_heightcnn = predicted_boxcnn[0,3]
      plt.imshow(photo, cmap='gray')
      plt.gca().add_patch(matplotlib.patches.Rectangle((box_xcnn, box_ycnn), box_widthcnn, box_heightcnn, ec='b', fc='none'))
      plt.gca().annotate("CNN", xy = (0,0), xytext = (0,117), color='b',fontsize = 8)
      #RF
      processed = preprocess_images(photo.reshape(1,128,128))
      predicted_boxRF = DT.predict(processed)
      box_xRF = predicted_boxRF[0,0]
      box_yRF = predicted_boxRF[0,1]
      box_widthRF = predicted_boxRF[0,2]
      box_heightRF = predicted_boxRF[0,3]
      plt.imshow(photo, cmap='gray')
      plt.gca().add_patch(matplotlib.patches.Rectangle((box_xRF, box_yRF), box_widthRF, box_heightRF, ec='y', fc='none'))
      plt.gca().annotate("RF", xy = (0,0), xytext = (0,122), color='y',fontsize = 8)
      #RR
      processed = preprocess_images(photo.reshape(1,128,128))
      predicted_boxRR = RR.predict(processed)
      box_xRR = predicted_boxRR[0,0]
      box_yRR = predicted_boxRR[0,1]
      box_widthRR = predicted_boxRR[0,2]
      box_heightRR = predicted_boxRR[0,3]
      plt.imshow(photo, cmap='gray')
      plt.gca().add_patch(matplotlib.patches.Rectangle((box_xRR, box_yRR), box_widthRR, box_heightRR, ec='c', fc='none'))
      plt.gca().annotate("RR", xy = (0,0), xytext = (0,127), color='c',fontsize = 8)
    return plt

gr.Interface(fn=multi_inputs,
             inputs=[input_module1, input_module2],
             outputs=[output_module1]
            ).launch(debug = True)