import gradio as gr

from fastai.vision.all import *

# load model
learn = load_learner('./model.pkl')

categories = ('CLEAR', 'PNEUMONIA')

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

examples = ['./unseen_data/NORMAL/NORMAL2-IM-1427-0001.jpeg', 
            './unseen_data/PNEUMONIA/person1946_bacteria_4874.jpeg',
            './unseen_data/NORMAL/NORMAL2-IM-1430-0001.jpeg',
            './unseen_data/PNEUMONIA/person1946_bacteria_4875.jpeg']

title = 'Pneumonia Child Chest X-Ray Classifier'

description = """An X-Ray classifier trained on the Children Chest X-Ray Images."""

article="""
<p>License: CC BY 4.0</p><br
<p><a href='http://www.cell.com/cell/fulltext/S0092-8674(18)30154-5' target='_blank'>Citation: Identifying Medical Diagnoses and Treatable Diseases by Image-Based Deep Learning</a></p><br>
<p><a href='https://www.kaggle.com/datasets/paultimothymooney/chest-xray-pneumonia' target='_blank'>Kaggle Dataset</a></p><br>
<p>Trained using the fast.ai library</p>
"""

intf = gr.Interface(fn=classify_image, 
                    inputs=gr.Image(type='pil'), 
                    outputs=gr.Label(), 
                    examples=examples,
                    title = title,
                    description = description,
                    article = article)

intf.launch(inline=False, share=True)