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

def load_h5_model(model_path):
    loaded_model = load_model(model_path)
    return(loaded_model)

def prepro_img(img):
    # resize the image to 130x130
    img = img.resize((130,130))
    # converto it to array with shape (1,130,130,3)
    img_array = np.array(img)
    img_array = np.array([img_array])
    # return result
    return(img_array)

def make_prediction(img):
    # preprocess image
    img = prepro_img(img)
    # make prediction
    prediction = model.predict(img)
    prediction = int(prediction[0][0])
    # return prediction label
    if prediction == 1:
        return('Uninfected cell')
    else:
        return('Parasitized cell')    

model = load_h5_model('cell_classifier_model.h5')

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            cell_img = gr.Image(label="Cell Image",
                                type='pil')
            examples = gr.Examples(['para_1.png', 'para_2.png', 'para_3.png'],
                                   inputs=cell_img,
                                   label='Parasitized Cells')
            examples = gr.Examples(['uninf_1.png', 'uninf_2.png', 'uninf_3.png'],
                                   inputs=cell_img,
                                   label='Uninfected Cells')
        with gr.Column():
            cell_class = gr.Label(value='...')
            predict_btn = gr.Button("Predict")

    predict_btn.click(fn=make_prediction, inputs=cell_img, outputs=cell_class)

demo.launch(debug=True)