from joblib import dump,load
import pandas as pd
import warnings
import gradio as gr
import cv2
warnings.filterwarnings("ignore")
small_X_train_flatten = pd.read_csv('Homework01_trainX_image_flatten.csv')
small_y_train = pd.read_csv('Homework01_trainy_image_flatten.csv')
best_knn = load("best_knn.joblib")
best_log = load("best_log.joblib")
best_knn.fit(small_X_train_flatten,small_y_train)
best_log.fit(small_X_train_flatten,small_y_train)

def preprocess_image(image):
    resized_image = cv2.resize(image, (28, 28))
    grayscale_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2GRAY)
    flattened_image = grayscale_image.flatten()
    normalized_image = flattened_image / 255.0
    return normalized_image

class_names = [
    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]

def classify_image(input_image, classifier):
    preprocessed_image = preprocess_image(input_image)
    reshaped_image = preprocessed_image.reshape(1, -1)
    if classifier == "Logistic Regression":
        output1 = best_log.predict(reshaped_image)[0]
        output2 = dict(zip(class_names, best_log.predict_proba(reshaped_image)[0]))
    elif classifier == "K-Nearest Neighbors":
        output1 = best_knn.predict(reshaped_image)[0]
        output2 = dict(zip(class_names, best_knn.predict_proba(reshaped_image)[0]))
    return class_names[output1], output2

gr.Interface(
    fn=classify_image,
    title="Fashion MNIST Classifier",
    inputs=[
        gr.Image(type="numpy", label="Input Image"),
        gr.Dropdown(
            label="Select Classifiers",
            choices=["Logistic Regression", "K-Nearest Neighbors"]
        )
    ],
    outputs=[
        gr.Textbox(label="Predicted Class"),
        gr.Label(label="Predicted Label Distribution")
    ]
).launch(share=True)