import numpy as np
import os
import keras
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from keras.models import Sequential
from PIL import Image
from keras.layers import Conv2D, Flatten, Dense, Dropout, BatchNormalization, MaxPooling2D
from sklearn.preprocessing import OneHotEncoder
import pickle
import gradio as gr

def load_model():
    save_path = 'model.pkl'
    with open(save_path, 'rb') as file:
        model = pickle.load(file)
    return model

def predict_dementia(images, model):
    predictions = []
    for image in images:
        img = Image.fromarray(image.astype('uint8'))
        img = img.resize((128, 128))
        img = np.array(img)
        img = img.reshape(1, 128, 128, 3)
        
        prediction = model.predict(img)
        prediction_class = np.argmax(prediction)
        predictions.append(names(prediction_class))
    return predictions

def names(number):
    if number == 0:
        return 'Non Demented'
    elif number == 1:
        return 'Mild Dementia'
    elif number == 2:
        return 'Moderate Dementia'
    elif number == 3:
        return 'Very Mild Dementia'
    else:
        return 'Error in Prediction'

def main(images):
    model = load_model()
    predictions = predict_dementia(images, model)
    return predictions

iface = gr.Interface(fn=main,
                     inputs="image",
                     outputs="text",
                     title="Dementia Classification",
                     description="Classify dementia based on brain images",
                     examples=[["Non(1).jpg"],["Moderate.jpg"],["Mild.jpg"]])

iface.launch(debug =True)