import joblib
import gradio as gr

dt_model = joblib.load('dt_model.joblib')
rfc_model = joblib.load('rfc_model.joblib')
abc_model = joblib.load('abc_model.joblib')
knn_model = joblib.load('knn_model.joblib')

from sklearn.metrics.cluster import normalized_mutual_info_score
import cv2
import numpy as np

def preprocess_image(image):
  grayscale_img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  normalized_img = grayscale_img / 255.0
  resized_img = cv2.resize(normalized_img, (28,28))
  flattened_img = resized_img.flatten()
  processed_img = np.array([flattened_img])

  return processed_img

def pred(input1, input2):
  class_names = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]
  processed_image = preprocess_image(input1)

  if input2 == "Decision Tree":
    preds = dt_model.predict_proba(processed_image)[0]
  elif input2 == "Random Forest":
    preds = rfc_model.predict_proba(processed_image)[0]
  elif input2 == "AdaBoost":
    preds = abc_model.predict_proba(processed_image)[0]
  elif input2 == "KNN":
    preds = knn_model.predict_proba(processed_image)[0]

  #print(class_names)
  #print(preds)

  return {label: float(pred) for label, pred in zip(class_names, preds)}


input_module1 = gr.inputs.Image(label = "Input Sample Image")
input_module2 = gr.inputs.Dropdown(choices=['Decision Tree', 'Random Forest', 'AdaBoost', 'KNN'], label = "ML Models Dropdown")

output_module1 = gr.outputs.Label(label = "Predicted Class")

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