app.py

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

# Define the functions to get the x and y values from the input dictionary - in this case, the x value is the image and the y value is the diagnosis
# needed to load the model since we defined them during training
def get_x(r): return ""
        
def get_y(r): return r['diagnosis']

learn = load_learner('model.pkl')
labels = learn.dls.vocab

# Define the mapping from label numbers to descriptions
label_descriptions = {
    0: "No DR",
    1: "Mild",
    2: "Moderate",
    3: "Severe",
    4: "Proliferative DR"
}

def predict(img):
    img = PILImage.create(img)
    pred,pred_idx,probs = learn.predict(img)
    # Use the label_descriptions dictionary to return descriptions instead of numbers
    return {label_descriptions[labels[i]]: float(probs[i]) for i in range(len(labels))}

title = "Diabetic Retinopathy Detection"
description = """Detects severity of diabetic retinopathy from a given retina image taken using fundus photography -

    0 - No DR

    1 - Mild

    2 - Moderate

    3 - Severe

    4 - Proliferative DR
"""
article = """
<p style='text-align: center'>
    <a href='https://www.kaggle.com/code/josemauriciodelgado/proliferative-retinopathy' target='_blank'>Kaggle Training Notebook</a> | 
    <a href='https://huggingface.co/jdelgado2002/diabetic_retinopathy_detection' target='_blank'>Model Card</a>
</p>
"""
# Get a list of all image paths in the test folder
test_folder = "test"  # replace with the actual path to your test folder
image_paths = [os.path.join(test_folder, img) for img in os.listdir(test_folder) if img.endswith(('.png', '.jpg', '.jpeg'))]

gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=5),
    examples=image_paths,  # set the examples parameter to the list of image paths
    article=article,
    title=title,
    description=description,
).launch()