from fastai.vision.all import *
import gradio as gr


def classify_image(img):
    img_pil = PILImage.create(img)
    img_resized = img_pil.resize((224, 224))

    pred, idx, probs = learn.predict(img_resized)
    return f"Wound class: {categories[idx]} \nHealing Time: {healing_times[idx]}"


learn = load_learner('model.pkl')

image = gr.Image()
label = gr.Label()

categories = ['Abrasion', 'Bruise', 'Burn', 'Cut', 'Diabetic Wound', 'Laseration', 'Normal', 'Pressure Wound',
              'Surgical Wound', 'Venous Wound']

healing_times = ['1-2 Weeks', '1-2 Weeks', '2-12 Weeks', '', '2-4 Months', '2-8 Weeks', '0', '',
                 '2 Weeks', '6-12 Months']


iface = gr.Interface(fn=classify_image, inputs=image, outputs=label)
iface.launch(inline=False, share=True)