import h2o
import gradio as gr
from PIL import Image
import numpy as np
import pandas as pd

h2o.init()

model = h2o.load_model('DeepLearning_grid_1_AutoML_11_20230827_202719_model_3')

column_names = model._model_json['output']['names'][:-1]

label_mapping = {
    'p0': 'choroidal neovascularization',
    'p1': 'diabetic macular edema',
    'p2': 'drusen',
    'p3': 'normal'
}


def predict(image):
    desired_size = (28, 28)

    resized_image = image.resize(desired_size)
    gray_image = resized_image.convert("L")


    image_array = np.asarray(gray_image)
    image_sample_flatten = image_array.flatten()
    print(image_sample_flatten.shape)
    df_sample = pd.DataFrame([image_sample_flatten], columns=column_names)

    image_sampleh2o = h2o.H2OFrame(df_sample)

    prediction = model.predict(image_sampleh2o)

    pandas_df = prediction.as_data_frame()


    prediction_dict = pandas_df.to_dict(orient="list")

    del prediction_dict['predict']
    new_dict = {label_mapping[key]: value for key, value in prediction_dict.items()}


    return(new_dict)



examples=[
    ['choroidal neovascularization.png'],
    ['diabetic macular edema.png'],
    ['drusen.png'],
    ['normal.png'],
]
demo = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), examples=examples, outputs="label", title="Retinal OCT Medical Image Classification", description='Retinal OCT Image Classification app using OCTMNIST dataset', article='Dataset rights: Jiancheng Yang, Rui Shi, Donglai Wei, Zequan Liu, Lin Zhao, Bilian Ke, Hanspeter Pfister, Bingbing Ni. Yang, Jiancheng, et al. "MedMNIST v2-A large-scale lightweight benchmark for 2D and 3D biomedical image classification." Scientific Data, 2023.')
demo.launch()