import gradio as gr
import shap
from fastai.vision.all import *
from efficientnet_pytorch import EfficientNet 



title = "COVID_19 Infection Detectation App!"
head = (
  "<div>"
    "This Space demonstrates model based on efficientnet base model. I has been trained to classify chest xray image." 
    " "
    "To test it, Use the Example Images Provided or Upload your own xray images the space provided."
    " "
    "The model is trained using [anasmohammedtahir/covidqu](https://www.kaggle.com/datasets/anasmohammedtahir/covidqu) dataset"  
  "</div>"
)
description = head

examples = [
    ['covid/covid_1038.png'], ['covid/covid_1034.png'], 
    ['covid/cd.png'], ['covid/covid_1021.png'], 
    ['covid/covid_1027.png'], ['covid/covid_1042.png'], 
    ['covid/covid_1031.png']
]

# Load the SHAP explainer for your model
explainer = shap.Explainer(model, data=None)

# Define a function to compute Shapley values for an image
def compute_shapley(image):
    # Preprocess the input image (e.g., resize, normalize) to match the model's input requirements
    # You should replace this with your specific preprocessing code
    image = preprocess_image(image)

    # Compute Shapley values for the input image
    shap_values = explainer(image)

    # Convert the Shapley values to a format that can be displayed (e.g., a heatmap)
    shap_values_image = shap.image_plot(shap_values[0], image)

    return shap_values_image

# Define the Gradio interface
image_input = gr.inputs.Image(shape=(224, 224))  # Adjust the shape as needed
shapley_output = gr.outputs.Image()

def interpreter_func(input_text):
    if input_text == "Compute Shapley":
        return compute_shapley(image_input)
    else:
        return None

interpreter_input = gr.inputs.Textbox(default="Type 'Compute Shapley' and click the button")
interpreter_button = "Generate Shapley"

#learn = load_learner('model/predictcovidfastaifinal18102023.pkl')
learn = load_learner('model/final_20102023_eb7_model.pkl')

categories = learn.dls.vocab

def predict_image(get_image):
   pred, idx, probs = learn.predict(get_image)
   return dict(zip(categories, map(float, probs)))

interpretation="default"
enable_queue=True

gr.Interface(fn=predict_image, inputs=gr.Image(shape=(224,224)),
             outputs = gr.Label(num_top_classes=3),title=title,description=description,examples=examples, interpretation=interpretation,enable_queue=enable_queue).launch(share=False)