import gradio as gr
import pickle
import numpy as np
from tensorflow.keras.models import load_model
from PIL import Image

# Update the path to the directory where your models are stored
model_directory = '/Users/keremoktay/Desktop/Gradio/'

# Load the decision tree and KNN models
dt_models = {
    "2": pickle.load(open('dt2.pkl', 'rb')),
    "3": pickle.load(open('dt3.pkl', 'rb')),
    "5": pickle.load(open('dt5.pkl', 'rb'))
}
knn_models = {
    "1": pickle.load(open('knn1.pkl', 'rb')),
    "3": pickle.load(open('knn3.pkl', 'rb')),
    "5": pickle.load(open('knn5.pkl', 'rb'))
}

# Load the neural network models
nn_models = {
    "1": load_model('cnn_1layer.h5'),
    "2": load_model('cnn_2layer.h5'),
    "3": load_model('cnn_3layer.h5')
}

def preprocess_image_for_ml(image_path):
    try:
        # For traditional ML models, just flatten the image without resizing
        image = Image.open(image_path)
        image_array = np.asarray(image).flatten().reshape(1, -1)
        return image_array
    except Exception as e:
        print(f"Error in preprocess_image_for_ml: {e}")
        raise e

def preprocess_image_for_cnn(image_path, target_size=(128, 128)):
    try:
        # For CNN models, convert to RGB, resize to 128x128, normalize, and add a batch dimension
        image = Image.open(image_path).convert('RGB')
        image = image.resize(target_size)
        image_array = np.asarray(image) / 255.0
        image_array = np.expand_dims(image_array, axis=0)
        return image_array
    except Exception as e:
        print(f"Error in preprocess_image_for_cnn: {e}")
        raise e

class_names = {
    0: "Angular Leaf Spot",
    1: "Bean Rust",
    2: "Healthy"
}

def classify_image_with_decision_tree(image, depth):
    try:
        image_array = preprocess_image_for_ml(image)
        model = dt_models.get(depth)
        prediction = model.predict(image_array)
        # Sayısal tahmini hastalık ismine dönüştür
        return class_names[int(prediction)]
    except Exception as e:
        print(f"Error in classify_image_with_decision_tree: {e}")
        raise e

def classify_image_with_knn(image, k):
    try:
        image_array = preprocess_image_for_ml(image)
        model = knn_models.get(k)
        prediction = model.predict(image_array)
        # Sayısal tahmini hastalık ismine dönüştür
        return class_names[int(prediction)]
    except Exception as e:
        print(f"Error in classify_image_with_knn: {e}")
        raise e

def classify_image_with_neural_network(image, layers):
    try:
        image_array = preprocess_image_for_cnn(image)
        model = nn_models.get(layers)
        prediction = model.predict(image_array)
        # En yüksek tahmin değerine sahip indeksi bul ve hastalık ismine dönüştür
        return class_names[np.argmax(prediction, axis=1)[0]]
    except Exception as e:
        print(f"Error in classify_image_with_neural_network: {e}")
        raise e

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("Image Classification using Different Models")
    
    with gr.Tab("Decision Tree"):
        with gr.Row():
            image_input_dt = gr.Image(type="filepath")
            dropdown_depth = gr.Dropdown(label="Select Depth", choices=["2", "3", "5"])
            output_dt = gr.Textbox(label="Decision Tree Output")
        classify_button_dt = gr.Button("Classify with Decision Tree")
        classify_button_dt.click(classify_image_with_decision_tree, inputs=[image_input_dt, dropdown_depth], outputs=output_dt)

    with gr.Tab("K-Nearest Neighbors (KNN)"):
        with gr.Row():
            image_input_knn = gr.Image(type="filepath")
            dropdown_k = gr.Dropdown(label="Select Number of Neighbors (k)", choices=["1", "3", "5"])
            output_knn = gr.Textbox(label="KNN Output")
        classify_button_knn = gr.Button("Classify with KNN")
        classify_button_knn.click(classify_image_with_knn, inputs=[image_input_knn, dropdown_k], outputs=output_knn)

    with gr.Tab("Neural Network"):
        with gr.Row():
            image_input_nn = gr.Image(type="filepath")
            dropdown_layers = gr.Dropdown(label="Select Number of Layers", choices=["1", "2", "3"])
            output_nn = gr.Textbox(label="Neural Network Output")
        classify_button_nn = gr.Button("Classify with Neural Network")
        classify_button_nn.click(classify_image_with_neural_network, inputs=[image_input_nn, dropdown_layers], outputs=output_nn)

demo.launch()