import gradio as gr
from huggingface_hub import hf_hub_download
from fastai.learner import load_learner

repo_id = "ethanmac/dr-macbloomber-retina-condition-classifier"

model = load_learner(
    hf_hub_download(repo_id, "model.pkl")
)

class_names = [
'Normal',
'Hollenhorst Emboli',
'Hypertensive Retinopathy',
'Coat\'s',
'Macroaneurism',
'Choroidal Neovascularization',
'Other',
'Branch Retinal Artery Occlusion',
'Cilio-Retinal Artery Occlusion',
'Branch Retinal Vein Occlusion',
'Central Retinal Vein Occlusion',
'Hemi-Central Retinal Vein Occlusion',
'Background Diabetic Retinopathy',
'Proliferative Diabetic Retinopathy',
'Arteriosclerotic Retinopathy'
]

categories = [c.replace('_', ' ').title() for c in class_names]

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

intf = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=5),
    title="Retinal Image Condition Classifier",
    examples=['healthy.jpg', 'crvo.jpg']
)

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