# coding: utf8
import gradio as gr
from transformers import pipeline

pipeline = pipeline(task="image-classification", model="dqnguyen/Diabetic_Foot_Ulcer_Image_Classification")

def predict(image):
    predictions = pipeline(image)
    #return {p["label"]: p["score"] for p in predictions}
    results = {}
    for p in predictions:
        if p["label"] == "MoHat":
            results["Granulation tissue"] = p["score"]
        elif p["label"] == "MoGiaMacNhiemKhuan":
            results["Pseudomembranous tissue with a bacterial infection"] = p["score"]
        elif p["label"] == "MoHoaiTu":
            results["Necrotic tissue"] = p["score"]
    return results
    
gr.Interface(
    predict,
    inputs=gr.inputs.Image(label="Upload a diabetic foot ulcer image", type="filepath"),
    outputs=gr.outputs.Label(num_top_classes=5),
    title="Diabetic Foot Ulcer Image Classification",
).launch()